#!/usr/bin/env python3
"""
core/grounding/infigui_server.py — Service grounding persistant InfiGUI

Charge InfiGUI-G1-3B en 4-bit une fois (~2.4 GB VRAM), puis sert les requêtes
de grounding via un Unix socket. Évite le coût de chargement (~10s) à chaque
appel que paie le subprocess one-shot.

Protocole (length-prefixed JSON) :
    Requête  : [4 octets uint32 BE = longueur] + payload JSON UTF-8
    Réponse  : [4 octets uint32 BE = longueur] + payload JSON UTF-8

Opérations supportées (champ "op", défaut "ground") :
    - "ping"     → {"ok": true, "vram_gb": float, "uptime_s": float}
    - "ground"   → {"x": int|None, "y": int|None, "confidence": float, ...}
    - "shutdown" → {"ok": true} puis arrêt propre du serveur

Le payload "ground" reprend exactement le format de l'ancien worker one-shot :
    {"target": str, "description": str, "image_path": str, "anchor_image_path": str}

Les images restent passées via fichiers (/tmp/...) — pas de bytes sur le socket.

Lancement (manuel ou via systemd user unit rpa-grounding.service) :
    cd ~/ai/rpa_vision_v3
    .venv/bin/python -m core.grounding.infigui_server

Variables d'environnement :
    RPA_GROUNDING_SOCKET   chemin du socket (défaut: $XDG_RUNTIME_DIR/rpa-grounding.sock
                           sinon /tmp/rpa-grounding.sock)
    RPA_GROUNDING_BACKLOG  taille listen backlog (défaut: 4)
"""

from __future__ import annotations

import json
import os
import signal
import socket
import struct
import sys
import threading
import time
import traceback
from typing import Any, Dict, Optional

# Réutilise la logique de chargement et d'inférence du worker one-shot.
# load_model() et infer() sont conçus pour être appelés en process indépendant ;
# on les appelle ici dans un process unique de longue durée.
from core.grounding.infigui_worker import infer, load_model


# ── Configuration ────────────────────────────────────────────────────────


def _default_socket_path() -> str:
    # /run/rpa/ est la convention "production" (RuntimeDirectory=rpa partagé
    # entre les services systemd). Cohérent avec ui_tars_grounder._default_socket_path.
    if os.path.isdir("/run/rpa"):
        return "/run/rpa/grounding.sock"
    runtime_dir = os.environ.get("XDG_RUNTIME_DIR")
    if runtime_dir and os.path.isdir(runtime_dir):
        return os.path.join(runtime_dir, "rpa-grounding.sock")
    return "/tmp/rpa-grounding.sock"


SOCKET_PATH = os.environ.get("RPA_GROUNDING_SOCKET") or _default_socket_path()
LISTEN_BACKLOG = int(os.environ.get("RPA_GROUNDING_BACKLOG", "4"))

# Limite raisonnable pour un payload JSON (la requête contient juste des chemins
# et du texte court — 4 MB suffit largement pour parer un client buggé).
MAX_PAYLOAD_BYTES = 4 * 1024 * 1024


# ── Protocole length-prefixed ────────────────────────────────────────────


def _recv_exact(conn: socket.socket, n: int) -> Optional[bytes]:
    """Lit exactement n octets ou retourne None si la connexion ferme avant."""
    chunks = []
    remaining = n
    while remaining > 0:
        chunk = conn.recv(remaining)
        if not chunk:
            return None
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def recv_message(conn: socket.socket) -> Optional[Dict[str, Any]]:
    header = _recv_exact(conn, 4)
    if header is None:
        return None
    (length,) = struct.unpack(">I", header)
    if length == 0 or length > MAX_PAYLOAD_BYTES:
        raise ValueError(f"Longueur payload invalide: {length}")
    payload = _recv_exact(conn, length)
    if payload is None:
        return None
    return json.loads(payload.decode("utf-8"))


def send_message(conn: socket.socket, obj: Dict[str, Any]) -> None:
    payload = json.dumps(obj, ensure_ascii=False).encode("utf-8")
    conn.sendall(struct.pack(">I", len(payload)) + payload)


# ── Serveur ──────────────────────────────────────────────────────────────


class InfiGUIServer:
    """Daemon qui sert les requêtes de grounding sur un Unix socket."""

    def __init__(self, socket_path: str = SOCKET_PATH):
        self.socket_path = socket_path
        self._server_sock: Optional[socket.socket] = None
        self._stop = threading.Event()
        # CUDA n'est pas thread-safe sur le même modèle ; sérialise les inférences.
        self._infer_lock = threading.Lock()
        self._model = None
        self._processor = None
        self._start_time = time.time()
        self._request_count = 0

    # ── Lifecycle ────────────────────────────────────────────────────────

    def start(self) -> None:
        # 1. Charger le modèle AVANT d'ouvrir le socket : si le chargement échoue,
        #    on n'expose pas un endpoint à moitié fonctionnel aux clients.
        print(f"[infigui-server] Chargement InfiGUI-G1-3B...")
        self._model, self._processor = load_model()

        # 2. Ouvrir le Unix socket (suppression d'un éventuel ancien socket fantôme)
        try:
            os.unlink(self.socket_path)
        except FileNotFoundError:
            pass

        self._server_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self._server_sock.bind(self.socket_path)
        # Permissions : owner-only (rw-------) pour éviter les autres utilisateurs
        os.chmod(self.socket_path, 0o600)
        self._server_sock.listen(LISTEN_BACKLOG)
        # Petit timeout sur accept pour que la boucle réagisse aux signaux
        self._server_sock.settimeout(1.0)

        print(f"[infigui-server] Écoute sur {self.socket_path}")

        # 3. Signaux d'arrêt propre
        signal.signal(signal.SIGTERM, self._on_signal)
        signal.signal(signal.SIGINT, self._on_signal)

        # 4. Boucle accept
        try:
            while not self._stop.is_set():
                try:
                    conn, _ = self._server_sock.accept()
                except socket.timeout:
                    continue
                except OSError:
                    if self._stop.is_set():
                        break
                    raise
                # Une requête à la fois (CUDA non thread-safe). On gère néanmoins
                # la connexion dans un thread pour pouvoir lire/écrire sans bloquer
                # l'accept boucle quand le client traîne — l'inférence elle-même
                # est sérialisée par self._infer_lock.
                threading.Thread(
                    target=self._handle_client,
                    args=(conn,),
                    daemon=True,
                ).start()
        finally:
            self._cleanup()

    def _on_signal(self, signum, _frame) -> None:
        print(f"[infigui-server] Signal {signum} reçu, arrêt...")
        self._stop.set()
        # Casse un éventuel accept() bloqué
        try:
            if self._server_sock is not None:
                self._server_sock.close()
        except OSError:
            pass

    def _cleanup(self) -> None:
        try:
            if self._server_sock is not None:
                self._server_sock.close()
        except OSError:
            pass
        try:
            os.unlink(self.socket_path)
        except FileNotFoundError:
            pass
        print(f"[infigui-server] Arrêté ({self._request_count} requêtes traitées)")

    # ── Gestion d'une connexion ──────────────────────────────────────────

    def _handle_client(self, conn: socket.socket) -> None:
        try:
            # Une connexion = N requêtes (keep-alive). Le client peut envoyer
            # plusieurs grounds successifs sans repayer le coût TCP/socket.
            while not self._stop.is_set():
                try:
                    req = recv_message(conn)
                except (ValueError, json.JSONDecodeError) as e:
                    self._safe_send(conn, {"error": f"requête invalide: {e}"})
                    return
                if req is None:
                    return  # client a fermé proprement

                op = req.get("op", "ground")
                if op == "ping":
                    self._safe_send(conn, self._do_ping())
                elif op == "shutdown":
                    self._safe_send(conn, {"ok": True})
                    self._stop.set()
                    try:
                        if self._server_sock is not None:
                            self._server_sock.close()
                    except OSError:
                        pass
                    return
                elif op == "ground":
                    resp = self._do_ground(req)
                    self._safe_send(conn, resp)
                else:
                    self._safe_send(conn, {"error": f"op inconnue: {op}"})
        except Exception as e:
            traceback.print_exc()
            self._safe_send(conn, {"error": str(e)})
        finally:
            try:
                conn.close()
            except OSError:
                pass

    def _safe_send(self, conn: socket.socket, obj: Dict[str, Any]) -> None:
        try:
            send_message(conn, obj)
        except OSError:
            # Client parti ; pas de quoi paniquer
            pass

    # ── Opérations ───────────────────────────────────────────────────────

    def _do_ping(self) -> Dict[str, Any]:
        try:
            import torch

            vram_gb = round(torch.cuda.memory_allocated() / 1e9, 2) if torch.cuda.is_available() else 0.0
        except Exception:
            vram_gb = 0.0
        return {
            "ok": True,
            "vram_gb": vram_gb,
            "uptime_s": round(time.time() - self._start_time, 1),
            "requests": self._request_count,
        }

    def _do_ground(self, req: Dict[str, Any]) -> Dict[str, Any]:
        with self._infer_lock:
            self._request_count += 1
            try:
                return infer(self._model, self._processor, req)
            except Exception as e:
                traceback.print_exc()
                return {
                    "x": None,
                    "y": None,
                    "error": str(e),
                }


def main() -> int:
    server = InfiGUIServer()
    try:
        server.start()
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"[infigui-server] Erreur fatale: {e}")
        traceback.print_exc()
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())