#!/usr/bin/env python3
"""
Lite Sources Helper - Privileged operations helper for Lite Sources
Authors: Jerry Bezencon, Ralphy
Website: https://www.linuxliteos.com
"""

import os
import sys
import socket
import subprocess
import signal

SOURCES_LIST = "/etc/apt/sources.list.d/linuxlite.list"
SOCKET_PATH = "/tmp/lite-sources-helper.sock"


def set_mirror(url, release):
    """Write the mirror to sources list."""
    try:
        with open(SOURCES_LIST, 'w') as f:
            f.write(f"deb {url} {release} main\n")
        return True, "Mirror set successfully"
    except Exception as e:
        return False, str(e)


def update_cache():
    """Run apt-get update."""
    try:
        result = subprocess.run(
            ['apt-get', 'update'],
            capture_output=True,
            text=True,
            timeout=180
        )
        if result.returncode == 0:
            return True, "Cache updated"
        else:
            return False, result.stderr
    except subprocess.TimeoutExpired:
        return False, "Update timed out"
    except Exception as e:
        return False, str(e)


def handle_command(cmd):
    """Handle a command and return response."""
    parts = cmd.strip().split('|')
    action = parts[0] if parts else ''

    if action == 'set-mirror':
        if len(parts) >= 3:
            success, msg = set_mirror(parts[1], parts[2])
            return "OK" if success else f"ERROR: {msg}"
        return "ERROR: Missing arguments"

    elif action == 'update-cache':
        success, msg = update_cache()
        return "OK" if success else f"ERROR: {msg}"

    elif action == 'set-mirror-and-update':
        if len(parts) >= 3:
            success, msg = set_mirror(parts[1], parts[2])
            if success:
                success, msg = update_cache()
            return "OK" if success else f"ERROR: {msg}"
        return "ERROR: Missing arguments"

    elif action == 'ping':
        return "PONG"

    elif action == 'quit':
        return "OK"

    return "ERROR: Unknown command"


def cleanup(signum=None, frame=None):
    """Clean up socket file."""
    try:
        if os.path.exists(SOCKET_PATH):
            os.remove(SOCKET_PATH)
    except Exception:
        pass
    sys.exit(0)


def run_daemon():
    """Run as a daemon listening on Unix socket."""
    # Clean up any existing socket
    cleanup_socket()

    # Set up signal handlers
    signal.signal(signal.SIGTERM, cleanup)
    signal.signal(signal.SIGINT, cleanup)

    # Create Unix socket
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        server.bind(SOCKET_PATH)
        os.chmod(SOCKET_PATH, 0o666)  # Allow non-root users to connect
        server.listen(1)

        while True:
            conn, addr = server.accept()
            try:
                data = conn.recv(4096).decode().strip()
                if data:
                    response = handle_command(data)
                    conn.sendall(response.encode())

                    if data.strip() == 'quit':
                        conn.close()
                        break
            except Exception as e:
                try:
                    conn.sendall(f"ERROR: {e}".encode())
                except Exception:
                    pass
            finally:
                conn.close()

    finally:
        server.close()
        cleanup()


def cleanup_socket():
    """Remove existing socket file."""
    try:
        if os.path.exists(SOCKET_PATH):
            os.remove(SOCKET_PATH)
    except Exception:
        pass


def run_single_command(args):
    """Run a single command (non-daemon mode)."""
    if args[0] == 'set-mirror':
        if len(args) >= 3:
            success, msg = set_mirror(args[1], args[2])
            print(msg)
            sys.exit(0 if success else 1)
        else:
            print("Usage: set-mirror <url> <release>")
            sys.exit(1)

    elif args[0] == 'update-cache':
        success, msg = update_cache()
        print(msg)
        sys.exit(0 if success else 1)

    elif args[0] == 'set-mirror-and-update':
        if len(args) >= 3:
            success, msg = set_mirror(args[1], args[2])
            if success:
                print(msg)
                success, msg = update_cache()
            print(msg)
            sys.exit(0 if success else 1)
        else:
            print("Usage: set-mirror-and-update <url> <release>")
            sys.exit(1)

    else:
        print(f"Unknown command: {args[0]}")
        print("Usage: lite-sources-helper {daemon|set-mirror|set-mirror-and-update|update-cache}")
        sys.exit(1)


def main():
    if len(sys.argv) < 2:
        print("Usage: lite-sources-helper {daemon|set-mirror|set-mirror-and-update|update-cache}")
        sys.exit(1)

    if sys.argv[1] == 'daemon':
        run_daemon()
    else:
        run_single_command(sys.argv[1:])


if __name__ == "__main__":
    main()
