"""
Axiom Zero — Remote Frontend TCP Bridge
=======================================
Runs inside the daemon (axiom_launcher.py) as a background thread.
Listens on TCP port 7443 with TLS, accepts connections from the
GTK frontend running on a DIFFERENT machine, and relays all IPC
events to it using the same length-prefixed JSON frame format
used by the UNIX socket.

If no TLS certificates are found, generates self-signed ones
automatically on first run and writes them to /var/lib/axiom-zero/.

Setup (automatic on first daemon start):
  Server cert → /var/lib/axiom-zero/server.pem + server.key
  CA cert      → /var/lib/axiom-zero/ca.pem
  Client cert  → generated and printed to console so operator can
                 copy ~/.axiom-zero/client.pem on their desktop

The frontend (gtk_ipc_listener.py) connects here automatically
if ~/.axiom-zero/connection.json specifies a remote server_host.
"""

import asyncio
import ssl
import json
import struct
import os
import logging
import ipaddress
import socket
import datetime
from pathlib import Path

user_dir = os.path.expanduser("~/.axiom-zero")
default_dir = "/var/lib/axiom-zero" if os.access("/var/lib", os.W_OK) else user_dir
RUNTIME_DIR = Path(os.environ.get("AXIOM_RUNTIME_DIR", default_dir))
SERVER_CERT  = RUNTIME_DIR / "server.pem"
SERVER_KEY   = RUNTIME_DIR / "server.key"
CA_CERT      = RUNTIME_DIR / "ca.pem"
CLIENT_CERT  = RUNTIME_DIR / "client.pem"
CLIENT_KEY   = RUNTIME_DIR / "client.key"
LISTEN_PORT  = 7443
log = logging.getLogger("axiom.tcp_bridge")


# ── Certificate generation ──────────────────────────────────────────

def generate_certs_if_missing():
    """Generate self-signed CA + server + client certs if not present."""
    RUNTIME_DIR.mkdir(parents=True, exist_ok=True)
    if SERVER_CERT.exists() and SERVER_KEY.exists():
        return  # Already exists

    try:
        from cryptography import x509
        from cryptography.x509.oid import NameOID
        from cryptography.hazmat.primitives import hashes, serialization
        from cryptography.hazmat.primitives.asymmetric import ec
        import datetime

        def _make_key():
            return ec.generate_private_key(ec.SECP384R1())

        def _write_pem(path, obj):
            if hasattr(obj, "private_bytes"):
                data = obj.private_bytes(serialization.Encoding.PEM,
                                         serialization.PrivateFormat.TraditionalOpenSSL,
                                         serialization.NoEncryption())
            else:
                data = obj.public_bytes(serialization.Encoding.PEM)
            path.write_bytes(data)
            path.chmod(0o600)

        now = datetime.datetime.utcnow()
        exp = now + datetime.timedelta(days=3650)

        # ── CA ────────────────────────────────────────────────────────
        ca_key  = _make_key()
        ca_name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "Axiom Zero CA")])
        ca_cert = (x509.CertificateBuilder()
                   .subject_name(ca_name).issuer_name(ca_name)
                   .public_key(ca_key.public_key())
                   .serial_number(x509.random_serial_number())
                   .not_valid_before(now).not_valid_after(exp)
                   .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
                   .sign(ca_key, hashes.SHA384()))

        # ── Server cert (SANs include hostname + local IPs) ───────────
        hostname   = socket.gethostname()
        local_ips  = _get_local_ips()
        san_dns    = [x509.DNSName(hostname), x509.DNSName("localhost")]
        san_ips    = [x509.IPAddress(ipaddress.IPv4Address(ip)) for ip in local_ips]
        san_ips   += [x509.IPAddress(ipaddress.IPv4Address("127.0.0.1"))]
        srv_key    = _make_key()
        srv_cert   = (x509.CertificateBuilder()
                      .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, hostname)]))
                      .issuer_name(ca_name)
                      .public_key(srv_key.public_key())
                      .serial_number(x509.random_serial_number())
                      .not_valid_before(now).not_valid_after(exp)
                      .add_extension(x509.SubjectAlternativeName(san_dns + san_ips), critical=False)
                      .sign(ca_key, hashes.SHA384()))

        # ── Client cert ───────────────────────────────────────────────
        cli_key  = _make_key()
        cli_cert = (x509.CertificateBuilder()
                    .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "axiom-operator")]))
                    .issuer_name(ca_name)
                    .public_key(cli_key.public_key())
                    .serial_number(x509.random_serial_number())
                    .not_valid_before(now).not_valid_after(exp)
                    .add_extension(x509.ExtendedKeyUsage([x509.ExtendedKeyUsageOID.CLIENT_AUTH]), critical=False)
                    .sign(ca_key, hashes.SHA384()))

        _write_pem(CA_CERT,     ca_cert)
        _write_pem(SERVER_CERT, srv_cert)
        _write_pem(SERVER_KEY,  srv_key)
        _write_pem(CLIENT_CERT, cli_cert)
        _write_pem(CLIENT_KEY,  cli_key)

        log.info("=" * 60)
        log.info("AXIOM ZERO — TLS CERTIFICATES GENERATED")
        log.info(f"  Server cert : {SERVER_CERT}")
        log.info(f"  Client cert : {CLIENT_CERT}")
        log.info(f"  Client key  : {CLIENT_KEY}")
        log.info(f"  CA cert     : {CA_CERT}")
        log.info("")
        log.info("TO CONNECT FROM YOUR DESKTOP, COPY THESE 3 FILES:")
        log.info(f"  scp root@{hostname}:{CLIENT_CERT} ~/.axiom-zero/client.pem")
        log.info(f"  scp root@{hostname}:{CLIENT_KEY}  ~/.axiom-zero/client.key")
        log.info(f"  scp root@{hostname}:{CA_CERT}     ~/.axiom-zero/ca.pem")
        log.info("")
        log.info("THEN RUN ON YOUR DESKTOP:")
        log.info(f"  axiom-zero-wizard")
        log.info(f"  (enter server IP: {local_ips[0] if local_ips else hostname})")
        log.info("=" * 60)

    except ImportError:
        log.warning("cryptography library not installed — TLS cert generation skipped")
        log.warning("Install with: pip3 install cryptography")


def _get_local_ips() -> list:
    ips = []
    try:
        import netifaces
        for iface in netifaces.interfaces():
            addrs = netifaces.ifaddresses(iface).get(netifaces.AF_INET, [])
            for a in addrs:
                ip = a.get("addr", "")
                if ip and not ip.startswith("127."):
                    ips.append(ip)
    except ImportError:
        # Fallback
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            s.connect(("8.8.8.8", 80))
            ips.append(s.getsockname()[0])
        except Exception:
            pass
        finally:
            s.close()
    return ips or ["127.0.0.1"]


# ── TCP server ──────────────────────────────────────────────────────

class RemoteFrontendBridge:
    """
    Asyncio TCP server. Accepts remote GTK frontend connections.
    Event broadcasts are pushed to all connected clients.
    """

    def __init__(self, event_queue: asyncio.Queue = None):
        self._clients    = set()
        self._lock       = asyncio.Lock()
        self.event_queue = event_queue or asyncio.Queue()
        self._server     = None

    async def start(self):
        generate_certs_if_missing()

        if not SERVER_CERT.exists():
            log.warning("No TLS cert found — remote frontend bridge disabled")
            return

        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ctx.load_cert_chain(certfile=str(SERVER_CERT), keyfile=str(SERVER_KEY))

        # Require client cert (mTLS) if CA is present
        if CA_CERT.exists():
            ctx.load_verify_locations(str(CA_CERT))
            ctx.verify_mode = ssl.CERT_REQUIRED

        self._server = await asyncio.start_server(
            self._handle_client, "0.0.0.0", LISTEN_PORT, ssl=ctx
        )
        log.info(f"Remote frontend bridge listening on TCP :{LISTEN_PORT} (mTLS)")

        # Start broadcaster coroutine
        asyncio.create_task(self._broadcast_loop())

        async with self._server:
            await self._server.serve_forever()

    async def _handle_client(self, reader: asyncio.StreamReader,
                              writer: asyncio.StreamWriter):
        peer = writer.get_extra_info("peername")
        log.info(f"Remote frontend connected from {peer}")
        async with self._lock:
            self._clients.add(writer)
        try:
            # Keep connection alive until client disconnects
            while True:
                data = await reader.read(1024)
                if not data:
                    break
        except Exception:
            pass
        finally:
            async with self._lock:
                self._clients.discard(writer)
            try:
                writer.close()
                await writer.wait_closed()
            except Exception:
                pass
            log.info(f"Remote frontend disconnected: {peer}")

    async def _broadcast_loop(self):
        """Picks events off the queue and sends to all connected remote clients."""
        while True:
            event = await self.event_queue.get()
            frame = self._encode_frame(event)
            async with self._lock:
                dead = set()
                for writer in self._clients:
                    try:
                        writer.write(frame)
                        await writer.drain()
                    except Exception:
                        dead.add(writer)
                self._clients -= dead

    def push_event(self, event: dict):
        """Thread-safe event push from the daemon's sync code."""
        try:
            self.event_queue.put_nowait(event)
        except asyncio.QueueFull:
            pass

    @staticmethod
    def _encode_frame(event: dict) -> bytes:
        payload = json.dumps(event).encode("utf-8")
        header  = struct.pack(">I", len(payload))
        return header + payload