#!/usr/bin/env python3
"""Authenticated LAN bridge from CodexPi to a PC-local Ollama server."""
from __future__ import annotations

import argparse
import hmac
import json
import os
import platform
import secrets
import shutil
import socket
import subprocess
import sys
import tempfile
import threading
import urllib.error
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path


MAX_REQUEST_BYTES = 64 * 1024 * 1024
ALLOWED_GET = {"/api/tags", "/v1/models", "/api/experimental/model-recommendations"}
ALLOWED_POST = {"/api/show", "/api/chat", "/api/generate", "/api/pull", "/v1/chat/completions", "/v1/responses"}


def checked_url(value: str, *, local_only: bool = False) -> str:
    parsed = urllib.parse.urlsplit(value)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password:
        raise ValueError("Adresse HTTP invalide")
    if local_only and parsed.hostname not in {"127.0.0.1", "localhost", "::1"}:
        raise ValueError("Ollama doit rester attaché à l’adresse locale du PC")
    return value.rstrip("/")


def json_request(server: str, path: str, body: dict) -> dict:
    data = json.dumps(body, ensure_ascii=False).encode("utf-8")
    request = urllib.request.Request(
        checked_url(server) + path,
        data=data,
        method="POST",
        headers={"Accept": "application/json", "Content-Type": "application/json", "User-Agent": "CodexPi-AI-Bridge/1"},
    )
    try:
        with urllib.request.urlopen(request, timeout=30) as response:
            raw = response.read()
    except urllib.error.HTTPError as exc:
        try:
            detail = json.loads(exc.read()).get("detail")
        except Exception:
            detail = None
        raise RuntimeError(detail or f"Erreur CodexPi HTTP {exc.code}") from exc
    value = json.loads(raw or b"{}")
    if not isinstance(value, dict):
        raise RuntimeError("Réponse CodexPi invalide")
    return value


def ollama_models(upstream: str) -> list[dict]:
    try:
        with urllib.request.urlopen(checked_url(upstream, local_only=True) + "/api/tags", timeout=5) as response:
            value = json.loads(response.read(2_000_000) or b"{}")
    except (OSError, urllib.error.URLError, ValueError) as exc:
        raise RuntimeError("Ollama ne répond pas sur ce PC. Ouvrez Ollama puis réessayez.") from exc
    models = value.get("models") if isinstance(value, dict) else None
    if not isinstance(models, list) or not models:
        raise RuntimeError("Ollama fonctionne mais aucun modèle n’est installé")
    return models


def hardware_info() -> dict:
    """Return only coarse sizing data needed to recommend a local model."""
    memory_bytes = 0
    try:
        if platform.system() == "Windows":
            import ctypes

            class MemoryStatus(ctypes.Structure):
                _fields_ = [
                    ("length", ctypes.c_ulong),
                    ("memory_load", ctypes.c_ulong),
                    ("total_physical", ctypes.c_ulonglong),
                    ("available_physical", ctypes.c_ulonglong),
                    ("total_page_file", ctypes.c_ulonglong),
                    ("available_page_file", ctypes.c_ulonglong),
                    ("total_virtual", ctypes.c_ulonglong),
                    ("available_virtual", ctypes.c_ulonglong),
                    ("available_extended_virtual", ctypes.c_ulonglong),
                ]

            status = MemoryStatus()
            status.length = ctypes.sizeof(status)
            if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
                memory_bytes = int(status.total_physical)
        elif hasattr(os, "sysconf"):
            memory_bytes = int(os.sysconf("SC_PAGE_SIZE")) * int(os.sysconf("SC_PHYS_PAGES"))
    except (AttributeError, OSError, TypeError, ValueError):
        memory_bytes = 0

    gpu_name = ""
    gpu_vram_bytes = 0
    executable = shutil.which("nvidia-smi")
    if executable:
        try:
            result = subprocess.run(
                [executable, "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
                check=True,
                capture_output=True,
                text=True,
                timeout=5,
            )
            candidates: list[tuple[int, str]] = []
            for line in result.stdout.splitlines():
                name, separator, memory = line.rpartition(",")
                if separator:
                    candidates.append((int(memory.strip()) * 1024 * 1024, name.strip()[:120]))
            if candidates:
                gpu_vram_bytes, gpu_name = max(candidates)
        except (OSError, subprocess.SubprocessError, ValueError):
            pass

    try:
        disk = shutil.disk_usage(Path.home())
        disk_free_bytes = int(disk.free)
    except OSError:
        disk_free_bytes = 0
    return {
        "platform": platform.system(),
        "architecture": platform.machine(),
        "memory_bytes": memory_bytes,
        "disk_free_bytes": disk_free_bytes,
        "gpu_name": gpu_name,
        "gpu_vram_bytes": gpu_vram_bytes,
    }


def save_config(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    descriptor, temporary = tempfile.mkstemp(prefix="ai-bridge-", suffix=".json", dir=path.parent)
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
            json.dump(value, handle, ensure_ascii=False, indent=2)
            handle.write("\n")
        try:
            os.chmod(temporary, 0o600)
        except OSError:
            pass
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def discover_address(server: str) -> str:
    parsed = urllib.parse.urlsplit(checked_url(server))
    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    addresses = socket.getaddrinfo(parsed.hostname, port, type=socket.SOCK_DGRAM)
    last_error: OSError | None = None
    for family, socket_type, protocol, _, target in addresses:
        probe = socket.socket(family, socket_type, protocol)
        try:
            probe.connect(target)
            return str(probe.getsockname()[0])
        except OSError as exc:
            last_error = exc
        finally:
            probe.close()
    raise RuntimeError(f"Impossible de déterminer l’adresse locale du PC : {last_error}")


def handler_for(upstream: str, token: str):
    class BridgeHandler(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.0"
        server_version = "CodexPiAIBridge/1"

        def log_message(self, fmt: str, *args) -> None:
            print(f"{self.client_address[0]} · {fmt % args}")

        def _authorized(self) -> bool:
            scheme, _, supplied = self.headers.get("Authorization", "").partition(" ")
            return scheme.lower() == "bearer" and hmac.compare_digest(supplied, token)

        def _json_error(self, status: int, message: str) -> None:
            body = json.dumps({"error": message}, ensure_ascii=False).encode("utf-8")
            self.send_response(status)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def _path(self) -> str:
            return urllib.parse.urlsplit(self.path).path

        def _proxy(self, method: str) -> None:
            if not self._authorized():
                self._json_error(401, "Jeton du bridge IA invalide")
                return
            path = self._path()
            allowed = ALLOWED_GET if method == "GET" else ALLOWED_POST
            if path not in allowed:
                self._json_error(404, "Route Ollama non autorisée")
                return
            length = int(self.headers.get("Content-Length") or 0)
            if length < 0 or length > MAX_REQUEST_BYTES:
                self._json_error(413, "Requête trop volumineuse")
                return
            data = self.rfile.read(length) if length else None
            headers = {"Accept": self.headers.get("Accept", "application/json"), "User-Agent": "CodexPi-AI-Bridge/1"}
            if data is not None:
                headers["Content-Type"] = self.headers.get("Content-Type", "application/json")
            request = urllib.request.Request(
                upstream + self.path,
                data=data,
                method=method,
                headers=headers,
            )
            try:
                with urllib.request.urlopen(request, timeout=900) as response:
                    self.send_response(response.status)
                    self.send_header("Content-Type", response.headers.get("Content-Type", "application/json"))
                    self.send_header("Cache-Control", "no-store")
                    self.send_header("Connection", "close")
                    self.end_headers()
                    while True:
                        chunk = response.read(64 * 1024)
                        if not chunk:
                            break
                        self.wfile.write(chunk)
                        self.wfile.flush()
            except urllib.error.HTTPError as exc:
                body = exc.read(2_000_000)
                self.send_response(exc.code)
                self.send_header("Content-Type", exc.headers.get("Content-Type", "application/json"))
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
            except (BrokenPipeError, ConnectionResetError):
                return
            except (OSError, urllib.error.URLError) as exc:
                self._json_error(502, f"Ollama est indisponible : {exc}")

        def do_GET(self) -> None:
            if self._path() == "/codexpi/info":
                if not self._authorized():
                    self._json_error(401, "Jeton du bridge IA invalide")
                    return
                body = json.dumps(hardware_info(), ensure_ascii=False).encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "application/json; charset=utf-8")
                self.send_header("Content-Length", str(len(body)))
                self.send_header("Cache-Control", "no-store")
                self.end_headers()
                self.wfile.write(body)
                return
            if self._path() == "/health":
                if not self._authorized():
                    self._json_error(401, "Jeton du bridge IA invalide")
                    return
                try:
                    models = ollama_models(upstream)
                except RuntimeError as exc:
                    self._json_error(503, str(exc))
                    return
                body = json.dumps({"ok": True, "runtime": "ollama", "models": len(models)}).encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.send_header("Cache-Control", "no-store")
                self.end_headers()
                self.wfile.write(body)
                return
            self._proxy("GET")

        def do_POST(self) -> None:
            self._proxy("POST")

    return BridgeHandler


def create_server(listen: str, port: int, upstream: str, token: str) -> ThreadingHTTPServer:
    return ThreadingHTTPServer((listen, port), handler_for(checked_url(upstream, local_only=True), token))


def pair(
    pairing_path: Path,
    config_path: Path,
    listen: str,
    port: int,
    upstream: str,
    name: str,
    advertise_address: str | None,
) -> tuple[dict, ThreadingHTTPServer]:
    pairing = json.loads(pairing_path.read_text(encoding="utf-8"))
    server_url = checked_url(str(pairing["server"]))
    ollama_models(upstream)
    token = secrets.token_urlsafe(48)
    server = create_server(listen, port, upstream, token)
    thread = threading.Thread(target=server.serve_forever, name="codexpi-ai-bridge", daemon=True)
    thread.start()
    address = advertise_address or discover_address(server_url)
    try:
        result = json_request(server_url, "/api/local-ai/agent/pair", {
            "pairing_id": pairing["pairing_id"],
            "code": pairing["code"],
            "name": name,
            "platform": platform.platform()[:160],
            "address": address,
            "port": port,
            "token": token,
        })
    except Exception:
        server.shutdown()
        server.server_close()
        raise
    config = {
        "listen": listen,
        "port": port,
        "upstream": checked_url(upstream, local_only=True),
        "token": token,
        "bridge_id": result.get("id"),
        "codexpi": server_url,
    }
    save_config(config_path, config)
    print(f"PC associé à CodexPi : {name}")
    print(f"Adresse privée annoncée : http://{address}:{port}")
    print(f"Modèle par défaut : {result.get('default_model') or 'à choisir dans CodexPi'}")
    print(f"Configuration privée : {config_path}")
    return config, server


def main() -> int:
    default_config = Path.home() / ".codexpi" / "ai-bridge.json"
    parser = argparse.ArgumentParser(description="Relie CodexPi à Ollama sur ce PC sans exposer l’API Ollama brute.")
    parser.add_argument("--pairing", type=Path, help="fichier .codexpi-ai téléchargé depuis Comptes IA")
    parser.add_argument("--config", type=Path, default=default_config, help="configuration privée du bridge")
    parser.add_argument("--listen", default="0.0.0.0", help="adresse d’écoute du bridge")
    parser.add_argument("--port", type=int, default=11435, help="port privé du bridge")
    parser.add_argument("--upstream", default="http://127.0.0.1:11434", help="adresse locale d’Ollama")
    parser.add_argument("--name", default=socket.gethostname(), help="nom affiché dans CodexPi")
    parser.add_argument("--advertise-address", help="adresse IP privée du PC visible depuis le Pi")
    args = parser.parse_args()
    if not 1024 <= args.port <= 65535:
        parser.error("--port doit être compris entre 1024 et 65535")
    active_server: ThreadingHTTPServer | None = None
    background_server = False
    try:
        if args.pairing:
            _, active_server = pair(
                args.pairing,
                args.config,
                args.listen,
                args.port,
                args.upstream,
                args.name,
                args.advertise_address,
            )
            background_server = True
        else:
            config = json.loads(args.config.read_text(encoding="utf-8"))
            ollama_models(str(config["upstream"]))
            active_server = create_server(
                str(config.get("listen") or args.listen),
                int(config.get("port") or args.port),
                str(config["upstream"]),
                str(config["token"]),
            )
            print(f"Bridge IA CodexPi prêt sur le port {config.get('port') or args.port}")
        print("Laissez cette fenêtre ouverte. Ctrl+C arrête le bridge.")
        if background_server:
            stop = threading.Event()
            while not stop.wait(3600):
                pass
        else:
            active_server.serve_forever()
        return 0
    except KeyboardInterrupt:
        print("\nBridge IA arrêté.")
        return 0
    except Exception as exc:
        print(f"Erreur : {exc}", file=sys.stderr)
        return 1
    finally:
        if active_server is not None:
            if background_server:
                active_server.shutdown()
            active_server.server_close()


if __name__ == "__main__":
    raise SystemExit(main())