#!/usr/bin/env python3
"""
VPS Monitor Agent — à déployer sur chaque VPS.
Expose une API REST utilisée par le backend central pour interroger les conteneurs Docker.
"""

import asyncio
import json
import os
import re
import shutil
import socket as socket_module
import subprocess
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

import docker
import psutil
from docker.errors import DockerException, NotFound
from fastapi import Depends, FastAPI, HTTPException, Security, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import APIKeyHeader
from pydantic import BaseModel

# ─── Config ───────────────────────────────────────────────────────────────────

AGENT_VERSION = "1.4.0"

REPO_BASE     = os.getenv("AGENT_REPO_BASE", "https://git.jeanbonapp.com/jeanbon/ScriptVPS/raw/branch/main")
INSTALL_DIR   = os.getenv("AGENT_INSTALL_DIR", "/opt/vps-monitor-agent")

API_KEY     = os.getenv("AGENT_API_KEY", "changeme-please")
AGENT_PORT  = int(os.getenv("AGENT_PORT", "8001"))

# ─── App ──────────────────────────────────────────────────────────────────────

app = FastAPI(title="VPS Monitor Agent", version="1.0.0", docs_url=None, redoc_url=None)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST", "PUT"],
    allow_headers=["*"],
)

api_key_header = APIKeyHeader(name="X-API-Key", auto_error=True)


def require_api_key(key: str = Security(api_key_header)) -> None:
    if key != API_KEY:
        raise HTTPException(status_code=403, detail="Clé API invalide")


def get_docker_client():
    try:
        return docker.from_env()
    except DockerException as e:
        raise HTTPException(status_code=503, detail=f"Docker inaccessible : {e}")


# ─── Routes ───────────────────────────────────────────────────────────────────

@app.get("/health")
def health():
    """Vérification de disponibilité — sans authentification."""
    return {"status": "ok", "timestamp": datetime.now(timezone.utc).isoformat()}


@app.get("/version")
def get_version():
    """Retourne la version de l'agent — sans authentification."""
    return {"version": AGENT_VERSION}


@app.post("/self-update")
def self_update(_: None = Depends(require_api_key)):
    """Télécharge la dernière version de l'agent depuis le dépôt et redémarre le service."""

    def _do_update():
        time.sleep(0.5)  # laisse la réponse HTTP partir
        try:
            for filename in ("agent.py", "requirements.txt"):
                src = f"{REPO_BASE}/vps-monitor/agent/{filename}"
                dst = f"{INSTALL_DIR}/{filename}"
                subprocess.run(
                    ["curl", "-fsSL", src, "-o", dst],
                    timeout=60,
                    check=True,
                )
            subprocess.run(
                [f"{INSTALL_DIR}/venv/bin/pip", "install", "-r",
                 f"{INSTALL_DIR}/requirements.txt", "-q"],
                timeout=120,
                check=True,
            )
            subprocess.run(
                ["systemctl", "restart", "vps-monitor-agent"],
                timeout=30,
                check=True,
            )
        except Exception:
            pass

    threading.Thread(target=_do_update, daemon=True).start()
    return {"status": "update_started"}


@app.get("/containers")
def list_containers(_: None = Depends(require_api_key)):
    """Retourne tous les conteneurs (actifs et arrêtés)."""
    client = get_docker_client()
    result = []
    for c in client.containers.list(all=True):
        image_tag = c.image.tags[0] if c.image.tags else c.image.short_id
        health_state = c.attrs.get("State", {}).get("Health", {})
        health = health_state.get("Status", "none") if health_state else "none"
        result.append({
            "id": c.short_id,
            "name": c.name.lstrip("/"),
            "status": c.status,
            "health": health,
            "image": image_tag,
            "created": c.attrs.get("Created", ""),
            "compose_project": c.labels.get("com.docker.compose.project", ""),
            "compose_service": c.labels.get("com.docker.compose.service", ""),
            "compose_working_dir": c.labels.get("com.docker.compose.project.working_dir", ""),
            "ports": {
                host: [{"HostIp": b["HostIp"], "HostPort": b["HostPort"]} for b in bindings]
                for host, bindings in (c.ports or {}).items()
                if bindings
            },
        })
    return sorted(result, key=lambda x: x["name"])


@app.get("/containers/{container_id}/logs")
def get_logs(container_id: str, lines: int = 100, _: None = Depends(require_api_key)):
    """Retourne les N dernières lignes de logs d'un conteneur."""
    client = get_docker_client()
    try:
        c = client.containers.get(container_id)
        raw = c.logs(tail=lines, timestamps=True)
        return {"logs": raw.decode("utf-8", errors="replace")}
    except NotFound:
        raise HTTPException(status_code=404, detail="Conteneur introuvable")


@app.post("/containers/{container_id}/action")
def container_action(container_id: str, action: str, _: None = Depends(require_api_key)):
    """Effectue une action sur un conteneur : start, stop, restart."""
    if action not in ("start", "stop", "restart"):
        raise HTTPException(status_code=400, detail=f"Action invalide : {action}")
    client = get_docker_client()
    try:
        c = client.containers.get(container_id)
        getattr(c, action)()
        return {"status": "ok", "action": action, "container": container_id}
    except NotFound:
        raise HTTPException(status_code=404, detail="Conteneur introuvable")


# ─── Terminal interactif (WebSocket) ─────────────────────────────────────────

# Choisit bash s'il est présent dans l'image, sinon sh — beaucoup d'images
# (alpine, distroless-ish) n'embarquent pas bash.
_SHELL_PICKER = "if command -v bash >/dev/null 2>&1; then exec bash; else exec sh; fi"

# Nombre maximal de sessions terminal simultanées, pour éviter qu'un client
# emballé n'ouvre des exec à l'infini.
_MAX_EXEC_SESSIONS = int(os.getenv("AGENT_MAX_EXEC_SESSIONS", "10"))
_exec_sessions = 0
_exec_sessions_lock = threading.Lock()


def _sock_recv(sock, size=4096):
    """Lecture bloquante — `recv` sur une vraie socket, `read` sur un SocketIO."""
    try:
        if hasattr(sock, "recv"):
            return sock.recv(size)
        return sock.read(size)
    except OSError:
        return b""


def _sock_send(sock, data: bytes) -> None:
    if hasattr(sock, "sendall"):
        sock.sendall(data)
    else:
        sock.write(data)
        sock.flush()


def _sock_close(sock) -> None:
    """Ferme la socket : débloque aussi le thread resté dans recv()."""
    try:
        if hasattr(sock, "shutdown"):
            sock.shutdown(socket_module.SHUT_RDWR)
    except OSError:
        pass
    try:
        sock.close()
    except OSError:
        pass


@app.websocket("/containers/{container_id}/exec")
async def exec_terminal(websocket: WebSocket, container_id: str, cols: int = 80, rows: int = 24):
    """Shell interactif dans un conteneur (équivalent `docker exec -it`).

    Protocole :
      • client → agent : JSON texte {"t":"i","d":"…"} (saisie) ou
        {"t":"r","cols":N,"rows":N} (redimensionnement) ;
      • agent → client : trames binaires (sortie brute du TTY) et JSON texte
        pour les messages d'état ({"t":"error"|"exit"}).

    L'authentification se fait par l'en-tête X-API-Key : le navigateur ne peut
    pas en poser sur un WebSocket, mais c'est le backend central qui se connecte
    ici, et lui le peut.
    """
    global _exec_sessions

    if websocket.headers.get("x-api-key") != API_KEY:
        await websocket.close(code=1008, reason="Clé API invalide")
        return

    with _exec_sessions_lock:
        if _exec_sessions >= _MAX_EXEC_SESSIONS:
            await websocket.close(code=1013, reason="Trop de sessions terminal ouvertes")
            return
        _exec_sessions += 1

    stream = None
    sock   = None
    try:
        await websocket.accept()

        try:
            client    = docker.from_env()
            container = client.containers.get(container_id)
        except NotFound:
            await websocket.send_text(json.dumps({"t": "error", "m": "Conteneur introuvable."}))
            return
        except DockerException as e:
            await websocket.send_text(json.dumps({"t": "error", "m": f"Docker inaccessible : {e}"}))
            return

        if container.status != "running":
            await websocket.send_text(json.dumps({
                "t": "error",
                "m": f"Le conteneur est « {container.status} » — démarrez-le pour ouvrir un terminal.",
            }))
            return

        api     = client.api
        exec_id = api.exec_create(
            container.id,
            cmd=["/bin/sh", "-c", _SHELL_PICKER],
            stdin=True, stdout=True, stderr=True, tty=True,
        )["Id"]

        stream = api.exec_start(exec_id, tty=True, socket=True, demux=False)
        sock   = getattr(stream, "_sock", stream)

        try:
            api.exec_resize(exec_id, height=rows, width=cols)
        except Exception:
            pass  # non bloquant : le shell tourne, seule la taille est approximative

        loop = asyncio.get_running_loop()

        async def pump_output():
            """Sortie du TTY → client, en binaire (xterm gère l'UTF-8 partiel)."""
            while True:
                chunk = await loop.run_in_executor(None, _sock_recv, sock)
                if not chunk:
                    break
                await websocket.send_bytes(chunk)

        async def pump_input():
            """Saisie et redimensionnements du client → TTY."""
            while True:
                message = await websocket.receive()
                if message["type"] == "websocket.disconnect":
                    break

                text = message.get("text")
                if text is not None:
                    try:
                        payload = json.loads(text)
                    except json.JSONDecodeError:
                        continue
                    kind = payload.get("t")
                    if kind == "i":
                        await loop.run_in_executor(
                            None, _sock_send, sock, payload.get("d", "").encode("utf-8")
                        )
                    elif kind == "r":
                        try:
                            api.exec_resize(
                                exec_id,
                                height=int(payload.get("rows", rows)),
                                width=int(payload.get("cols", cols)),
                            )
                        except Exception:
                            pass
                    continue

                data = message.get("bytes")
                if data:
                    await loop.run_in_executor(None, _sock_send, sock, data)

        tasks = [asyncio.create_task(pump_output()), asyncio.create_task(pump_input())]
        try:
            _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        finally:
            # Fermer la socket débloque le thread encore dans recv().
            _sock_close(sock)
            for task in tasks:
                task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)

        try:
            code = api.exec_inspect(exec_id).get("ExitCode")
            await websocket.send_text(json.dumps({"t": "exit", "code": code}))
        except Exception:
            pass

    except WebSocketDisconnect:
        pass
    finally:
        if sock is not None:
            _sock_close(sock)
        if stream is not None and stream is not sock:
            try:
                stream.close()
            except Exception:
                pass
        with _exec_sessions_lock:
            _exec_sessions -= 1
        try:
            await websocket.close()
        except Exception:
            pass


@app.get("/system")
def system_info(_: None = Depends(require_api_key)):
    """Retourne les informations système : CPU, RAM et bande passante."""
    cpu_percent = psutil.cpu_percent(interval=0.5)
    mem = psutil.virtual_memory()

    net1 = psutil.net_io_counters()
    time.sleep(0.5)
    net2 = psutil.net_io_counters()
    net_sent_per_sec = (net2.bytes_sent - net1.bytes_sent) * 2
    net_recv_per_sec = (net2.bytes_recv - net1.bytes_recv) * 2

    return {
        "cpu_percent":      cpu_percent,
        "ram_used":         mem.used,
        "ram_total":        mem.total,
        "ram_percent":      mem.percent,
        "net_sent_per_sec": net_sent_per_sec,
        "net_recv_per_sec": net_recv_per_sec,
        "net_bytes_sent":   net2.bytes_sent,
        "net_bytes_recv":   net2.bytes_recv,
    }


@app.post("/compose/update")
def compose_update(project: str, _: None = Depends(require_api_key)):
    """Pull les nouvelles images et recrée les conteneurs d'un projet compose."""
    client = get_docker_client()

    working_dir = None
    for c in client.containers.list(all=True):
        if c.labels.get("com.docker.compose.project") == project:
            working_dir = c.labels.get("com.docker.compose.project.working_dir")
            if working_dir:
                break

    if not working_dir:
        raise HTTPException(
            status_code=404,
            detail=f"Projet compose '{project}' introuvable ou sans répertoire de travail",
        )

    output = ""

    pull = subprocess.run(
        ["docker", "compose", "pull"],
        cwd=working_dir,
        capture_output=True,
        text=True,
        timeout=300,
    )
    output += pull.stdout + pull.stderr

    up = subprocess.run(
        ["docker", "compose", "up", "-d", "--remove-orphans"],
        cwd=working_dir,
        capture_output=True,
        text=True,
        timeout=120,
    )
    output += up.stdout + up.stderr

    return {"output": output, "project": project, "working_dir": working_dir}


# ─── Applications compose ────────────────────────────────────────────────────
#
# Convention : une application = un dossier sous /home contenant son fichier
# compose, soit /home/<application>/compose.yaml (ou l'un des autres noms
# reconnus par Docker).

APPS_DIR = Path(os.getenv("AGENT_APPS_DIR", "/home"))

_COMPOSE_NAMES = ("compose.yaml", "compose.yml", "docker-compose.yml", "docker-compose.yaml")
_APP_NAME_RE   = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
_MAX_COMPOSE_BYTES = 512 * 1024


class ComposeWriteRequest(BaseModel):
    content: str
    deploy:  bool = False


class AppCreateRequest(BaseModel):
    name:     str
    content:  str
    filename: str = "compose.yaml"
    deploy:   bool = False


def _resolve_app_dir(name: str) -> Path:
    """Chemin du dossier d'une application, en refusant tout ce qui sort d'APPS_DIR."""
    if not _APP_NAME_RE.match(name or ""):
        raise HTTPException(status_code=400, detail="Nom d'application invalide")
    base = APPS_DIR.resolve()
    path = (base / name).resolve()
    # Le nom ne contient ni « / » ni « .. » (regex) ; ce contrôle attrape en plus
    # les liens symboliques qui pointeraient hors de /home.
    if path.parent != base:
        raise HTTPException(status_code=400, detail="Chemin d'application invalide")
    return path


def _find_compose_file(app_dir: Path) -> Path | None:
    for candidate in _COMPOSE_NAMES:
        path = app_dir / candidate
        if path.is_file():
            return path
    return None


def _validate_compose(app_dir: Path, content: str, display_name: str) -> tuple[bool, str]:
    """Valide le YAML avec `docker compose config` avant d'écrire quoi que ce soit.

    La validation se fait sur un fichier temporaire dans le dossier de
    l'application, pour que les chemins relatifs et le `.env` soient résolus
    exactement comme au déploiement.
    """
    tmp = app_dir / f".compose-check-{os.getpid()}.yaml"
    try:
        tmp.write_text(content, encoding="utf-8")
        proc = subprocess.run(
            ["docker", "compose", "-f", str(tmp), "config", "-q"],
            cwd=app_dir, capture_output=True, text=True, timeout=30,
        )
        message = (proc.stderr or proc.stdout).strip().replace(tmp.name, display_name)
        return proc.returncode == 0, message
    except FileNotFoundError:
        return True, ""   # docker compose absent : on n'empêche pas l'enregistrement
    except subprocess.TimeoutExpired:
        return False, "La validation `docker compose config` a expiré."
    except OSError as e:
        raise HTTPException(status_code=500, detail=f"Écriture impossible : {e}")
    finally:
        tmp.unlink(missing_ok=True)


def _compose_up(app_dir: Path) -> str:
    output = ""
    for cmd, timeout in ((["docker", "compose", "pull"], 600),
                         (["docker", "compose", "up", "-d", "--remove-orphans"], 300)):
        proc = subprocess.run(cmd, cwd=app_dir, capture_output=True, text=True, timeout=timeout)
        output += proc.stdout + proc.stderr
    return output


def _app_payload(app_dir: Path, compose: Path, running: dict[str, int]) -> dict:
    stat = compose.stat()
    return {
        "name":     app_dir.name,
        "path":     str(app_dir),
        "file":     compose.name,
        "size":     stat.st_size,
        "modified": datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc).isoformat(),
        "running_containers": running.get(str(app_dir), 0),
    }


def _running_by_working_dir() -> dict[str, int]:
    """Conteneurs démarrés, indexés par répertoire de projet compose."""
    counts: dict[str, int] = {}
    try:
        client = docker.from_env()
        for container in client.containers.list():
            working_dir = container.labels.get("com.docker.compose.project.working_dir")
            if working_dir:
                counts[working_dir] = counts.get(working_dir, 0) + 1
    except DockerException:
        pass
    return counts


@app.get("/apps")
def list_apps(_: None = Depends(require_api_key)):
    """Liste les applications : dossiers de /home contenant un fichier compose."""
    if not APPS_DIR.is_dir():
        return []

    running = _running_by_working_dir()
    apps = []
    try:
        entries = sorted(APPS_DIR.iterdir(), key=lambda p: p.name.lower())
    except OSError as e:
        raise HTTPException(status_code=500, detail=f"Lecture de {APPS_DIR} impossible : {e}")

    for entry in entries:
        if not entry.is_dir() or entry.name.startswith("."):
            continue
        compose = _find_compose_file(entry)
        if compose:
            apps.append(_app_payload(entry, compose, running))
    return apps


@app.get("/apps/{name}/compose")
def read_compose(name: str, _: None = Depends(require_api_key)):
    """Retourne le contenu du fichier compose d'une application."""
    app_dir = _resolve_app_dir(name)
    if not app_dir.is_dir():
        raise HTTPException(status_code=404, detail="Application introuvable")
    compose = _find_compose_file(app_dir)
    if not compose:
        raise HTTPException(status_code=404, detail="Aucun fichier compose dans ce dossier")
    try:
        content = compose.read_text(encoding="utf-8")
    except (OSError, UnicodeDecodeError) as e:
        raise HTTPException(status_code=500, detail=f"Lecture impossible : {e}")

    payload = _app_payload(app_dir, compose, _running_by_working_dir())
    payload["content"] = content
    return payload


@app.put("/apps/{name}/compose")
def write_compose(name: str, body: ComposeWriteRequest, _: None = Depends(require_api_key)):
    """Écrit le fichier compose, après validation et sauvegarde de l'ancienne version."""
    app_dir = _resolve_app_dir(name)
    if not app_dir.is_dir():
        raise HTTPException(status_code=404, detail="Application introuvable")
    if len(body.content.encode("utf-8")) > _MAX_COMPOSE_BYTES:
        raise HTTPException(status_code=413, detail="Fichier compose trop volumineux")

    target = _find_compose_file(app_dir) or (app_dir / "compose.yaml")

    valid, message = _validate_compose(app_dir, body.content, target.name)
    if not valid:
        raise HTTPException(status_code=400, detail=message or "Fichier compose invalide")

    backup = None
    try:
        if target.exists():
            backup = target.with_name(target.name + ".bak")
            shutil.copy2(target, backup)
        target.write_text(body.content, encoding="utf-8")
    except OSError as e:
        raise HTTPException(status_code=500, detail=f"Écriture impossible : {e}")

    result = {
        "status": "ok",
        "file":   target.name,
        "path":   str(target),
        "backup": backup.name if backup else None,
    }
    if body.deploy:
        try:
            result["output"] = _compose_up(app_dir)
        except subprocess.TimeoutExpired:
            raise HTTPException(status_code=504, detail="Le déploiement a expiré")
    return result


@app.post("/apps", status_code=201)
def create_app(body: AppCreateRequest, _: None = Depends(require_api_key)):
    """Crée une application : dossier /home/<nom> et son fichier compose."""
    if body.filename not in _COMPOSE_NAMES:
        raise HTTPException(
            status_code=400,
            detail=f"Nom de fichier invalide — attendu : {', '.join(_COMPOSE_NAMES)}",
        )
    if len(body.content.encode("utf-8")) > _MAX_COMPOSE_BYTES:
        raise HTTPException(status_code=413, detail="Fichier compose trop volumineux")

    app_dir = _resolve_app_dir(body.name)
    if _find_compose_file(app_dir):
        raise HTTPException(status_code=409, detail="Cette application existe déjà")

    created_dir = not app_dir.exists()
    try:
        app_dir.mkdir(parents=True, exist_ok=True)
    except OSError as e:
        raise HTTPException(status_code=500, detail=f"Création du dossier impossible : {e}")

    valid, message = _validate_compose(app_dir, body.content, body.filename)
    if not valid:
        if created_dir:
            try:
                app_dir.rmdir()   # ne laisse pas un dossier vide derrière soi
            except OSError:
                pass
        raise HTTPException(status_code=400, detail=message or "Fichier compose invalide")

    target = app_dir / body.filename
    try:
        target.write_text(body.content, encoding="utf-8")
    except OSError as e:
        raise HTTPException(status_code=500, detail=f"Écriture impossible : {e}")

    result = {
        "status":      "created",
        "name":        app_dir.name,
        "path":        str(target),
        "file":        target.name,
        "created_dir": created_dir,
    }
    if body.deploy:
        try:
            result["output"] = _compose_up(app_dir)
        except subprocess.TimeoutExpired:
            raise HTTPException(status_code=504, detail="Le déploiement a expiré")
    return result


@app.post("/apps/{name}/up")
def deploy_app(name: str, _: None = Depends(require_api_key)):
    """Lance `docker compose pull` puis `up -d` dans le dossier de l'application."""
    app_dir = _resolve_app_dir(name)
    compose = _find_compose_file(app_dir) if app_dir.is_dir() else None
    if not compose:
        raise HTTPException(status_code=404, detail="Application introuvable")
    try:
        return {"output": _compose_up(app_dir), "name": app_dir.name, "path": str(app_dir)}
    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=504, detail="Le déploiement a expiré")


@app.get("/services")
def list_services(_: None = Depends(require_api_key)):
    """Retourne la liste des services systemd (hors Docker) avec leur état."""
    try:
        result = subprocess.run(
            ["systemctl", "list-units", "--type=service", "--no-legend", "--no-pager", "--all"],
            capture_output=True, text=True, timeout=10,
        )
    except FileNotFoundError:
        raise HTTPException(status_code=501, detail="systemctl introuvable — système non-systemd")
    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=504, detail="systemctl a expiré")

    _DOCKER_SERVICES = {"docker.service", "containerd.service", "docker.socket"}

    services = []
    for line in result.stdout.strip().splitlines():
        # Supprime les puces (● ○) et les espaces de début
        line = line.lstrip("●○").strip()
        if not line:
            continue
        parts = line.split(None, 4)
        if len(parts) < 4:
            continue
        name = parts[0]
        if not name.endswith(".service"):
            continue
        if name.lower() in _DOCKER_SERVICES or name.lower().startswith("docker"):
            continue
        services.append({
            "name":        name,
            "load":        parts[1],
            "active":      parts[2],
            "sub":         parts[3],
            "description": parts[4].strip() if len(parts) > 4 else "",
        })

    return sorted(services, key=lambda s: s["name"])


# ─── Entrée ───────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=AGENT_PORT)