"""
Chat IA Multi-Modèles — Serveur de synchronisation & proxy de génération (V9.4)
================================================================================

Ce serveur est prévu pour tourner dans un HuggingFace Space (16 Go RAM / 2 CPU,
port 7860) et fait 3 choses :

1. Authentification (inscription / connexion) avec des tokens JWT.
2. Synchronisation de l'état complet du chat (chats, clés API, réglages) par
   utilisateur, stocké dans /data (volume persistant du Space).
3. Proxy de génération : le client n'appelle plus l'API du provider (OpenAI,
   Groq, Gemini, custom) directement. Il demande au serveur de le faire.
   Le serveur lance une tâche asyncio EN ARRIÈRE-PLAN, indépendante de la
   requête HTTP du client : si le client se déconnecte (ferme l'onglet, perd
   le réseau, change d'appareil), la génération continue côté serveur et le
   texte est accumulé en mémoire (+ persisté à la fin dans /data). N'importe
   quel appareil connecté au même compte peut alors se rebrancher sur le flux
   (SSE) en cours ou récupérer le résultat déjà terminé.

Stockage disque (dans /data) :
  /data/users.json            -> {username: {password_hash, created_at}}
  /data/states/<username>.json -> état complet du chat de cet utilisateur
  /data/secret.key            -> clé secrète JWT (générée une fois, persistée)

Tout est volontairement simple (fichiers JSON + verrous asyncio) : c'est
largement suffisant pour un usage personnel / petit groupe d'utilisateurs sur
un Space à 2 CPU. Pas de base de données externe requise.
"""

import asyncio
import json
import os
import secrets
import time
import uuid
from pathlib import Path
from typing import Optional
from urllib.parse import urljoin

import bcrypt
import httpx
import jwt
from fastapi import FastAPI, Header, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

# --------------------------------------------------------------------------
# Configuration & stockage
# --------------------------------------------------------------------------

DATA_DIR = Path(os.environ.get("DATA_DIR", "/data"))
USERS_DIR = DATA_DIR
STATES_DIR = DATA_DIR / "states"
SECRET_FILE = DATA_DIR / "secret.key"
USERS_FILE = DATA_DIR / "users.json"

DATA_DIR.mkdir(parents=True, exist_ok=True)
STATES_DIR.mkdir(parents=True, exist_ok=True)

if not SECRET_FILE.exists():
    SECRET_FILE.write_text(secrets.token_hex(32))
JWT_SECRET = SECRET_FILE.read_text().strip()
JWT_ALGO = "HS256"
JWT_TTL_SECONDS = 60 * 60 * 24 * 30  # 30 jours

EMPTY_STATE = {
    "chats": [],
    "activeChatId": None,
    "virtualFiles": {},
    "currentImage": None,
    "currentTextFile": None,
    "keys": {"openai": "", "groq": "", "gemini": "", "deepai": "", "custom": []},
    "settings": {"webSearch": False, "fileSystem": False, "visionModel": "", "typeSpeed": 18},
}

# Grâce à ce verrou par utilisateur, deux écritures concurrentes sur le même
# fichier d'état ne se corrompent pas mutuellement.
_user_locks: dict[str, asyncio.Lock] = {}
_users_lock = asyncio.Lock()


def get_user_lock(username: str) -> asyncio.Lock:
    if username not in _user_locks:
        _user_locks[username] = asyncio.Lock()
    return _user_locks[username]


def _atomic_write_json(path: Path, data: dict):
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(data, ensure_ascii=False))
    os.replace(tmp, path)


async def load_users() -> dict:
    async with _users_lock:
        if not USERS_FILE.exists():
            return {}
        try:
            return json.loads(USERS_FILE.read_text())
        except Exception:
            return {}


async def save_users(users: dict):
    async with _users_lock:
        _atomic_write_json(USERS_FILE, users)


def state_path(username: str) -> Path:
    safe = "".join(c for c in username if c.isalnum() or c in ("-", "_")) or "user"
    return STATES_DIR / f"{safe}.json"


async def load_state(username: str) -> dict:
    path = state_path(username)
    if not path.exists():
        return json.loads(json.dumps(EMPTY_STATE))
    try:
        data = json.loads(path.read_text())
    except Exception:
        return json.loads(json.dumps(EMPTY_STATE))
    for k, v in EMPTY_STATE.items():
        if k not in data:
            data[k] = v
    return data


async def save_state(username: str, data: dict):
    async with get_user_lock(username):
        _atomic_write_json(state_path(username), data)


# --------------------------------------------------------------------------
# Auth
# --------------------------------------------------------------------------

def make_token(username: str) -> str:
    payload = {"sub": username, "exp": int(time.time()) + JWT_TTL_SECONDS}
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGO)


def verify_token(authorization: Optional[str]) -> str:
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Token manquant")
    token = authorization[len("Bearer "):]
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGO])
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Session expirée, reconnectez-vous")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Token invalide")
    return payload["sub"]


class AuthBody(BaseModel):
    username: str
    password: str


class SyncBody(BaseModel):
    data: dict


# --------------------------------------------------------------------------
# Jobs de génération (proxy streaming résilient)
# --------------------------------------------------------------------------

# jobs[job_id] = {
#   username, chat_id, message_id, provider, model,
#   buffer, status, error, tokens, promptTokens, completionTokens,
#   created_at, updated_at, cancel_event
# }
JOBS: dict[str, dict] = {}
JOB_GRACE_SECONDS = 60 * 30  # on garde un job terminé 30 min pour permettre une reprise tardive


def new_job(username: str, chat_id, message_id: str, provider: str, model: str) -> str:
    job_id = uuid.uuid4().hex
    JOBS[job_id] = {
        "id": job_id,
        "username": username,
        "chat_id": chat_id,
        "message_id": message_id,
        "provider": provider,
        "model": model,
        "buffer": "",
        "status": "running",
        "error": None,
        "tokens": None,
        "promptTokens": None,
        "completionTokens": None,
        "created_at": time.time(),
        "updated_at": time.time(),
        "cancel_event": asyncio.Event(),
        "status_text": None,  # statut agentique transitoire (ex: recherche web en cours)
    }
    return job_id


async def persist_job_result(job: dict):
    """Écrit le contenu final du message dans l'état persistant de l'utilisateur,
    et crée le chat/message s'il n'existe pas encore côté serveur (évite une
    course avec le debounce de /sync côté client)."""
    username = job["username"]
    async with get_user_lock(username):
        state = await load_state(username)
        chat = next((c for c in state["chats"] if str(c.get("id")) == str(job["chat_id"])), None)
        if chat is None:
            chat = {
                "id": job["chat_id"], "title": "Nouvelle discussion",
                "messages": [], "model": job["model"], "provider": job["provider"], "extraPricing": None,
            }
            state["chats"].insert(0, chat)
        msg = next((m for m in chat["messages"] if m.get("id") == job["message_id"]), None)
        if msg is None:
            msg = {"id": job["message_id"], "role": "assistant", "content": ""}
            chat["messages"].append(msg)
        msg["content"] = job["buffer"]
        msg["isThinking"] = False
        if job["tokens"]:
            msg["tokens"] = job["tokens"]
        if job["promptTokens"]:
            msg["promptTokens"] = job["promptTokens"]
        if job["completionTokens"]:
            msg["completionTokens"] = job["completionTokens"]
        if job["status"] == "error" and job["error"]:
            msg["content"] = (msg["content"] or "") + f"\n**Erreur API:** {job['error']}"
        _atomic_write_json(state_path(username), state)


async def upsert_pending_message(username: str, chat_id, chat_title: str, recent_messages: list):
    """Enregistre immédiatement le message utilisateur + le placeholder assistant
    (appelé au lancement du job) pour que le chat existe côté serveur sans
    attendre le prochain /sync du client."""
    async with get_user_lock(username):
        state = await load_state(username)
        chat = next((c for c in state["chats"] if str(c.get("id")) == str(chat_id)), None)
        if chat is None:
            chat = {
                "id": chat_id, "title": chat_title or "Nouvelle discussion",
                "messages": [], "model": "", "provider": "", "extraPricing": None,
            }
            state["chats"].insert(0, chat)
        existing_ids = {m.get("id") for m in chat["messages"]}
        for m in recent_messages or []:
            if m.get("id") and m["id"] not in existing_ids:
                chat["messages"].append(m)
                existing_ids.add(m["id"])
        _atomic_write_json(state_path(username), state)


def _sse(event: str, data: dict) -> bytes:
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n".encode("utf-8")


async def run_openai_like(job: dict, url: str, api_key: str, model: str, system: str, messages: list):
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    full_messages = [{"role": "system", "content": system}] + [
        {"role": m.get("role", "user"), "content": m.get("content", "")} for m in messages
    ]
    body = {"model": model, "messages": full_messages, "stream": True, "stream_options": {"include_usage": True}}

    async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=30.0)) as client:
        async with client.stream("POST", url, headers=headers, json=body) as resp:
            if resp.status_code >= 400:
                err_text = (await resp.aread()).decode("utf-8", "ignore")
                raise RuntimeError(f"HTTP {resp.status_code}: {err_text[:500]}")
            async for line in resp.aiter_lines():
                if job["cancel_event"].is_set():
                    return
                line = line.strip()
                if not line.startswith("data:"):
                    continue
                data_str = line[5:].strip()
                if data_str == "[DONE]":
                    continue
                try:
                    data = json.loads(data_str)
                except Exception:
                    continue
                choices = data.get("choices") or []
                if choices:
                    delta = (choices[0].get("delta") or {}).get("content")
                    if delta:
                        job["buffer"] += delta
                        job["updated_at"] = time.time()
                usage = data.get("usage")
                if usage:
                    job["tokens"] = usage.get("total_tokens") or job["tokens"]
                    job["promptTokens"] = usage.get("prompt_tokens") or job["promptTokens"]
                    job["completionTokens"] = usage.get("completion_tokens") or job["completionTokens"]


async def run_gemini(job: dict, api_key: str, model: str, system: str, messages: list):
    url = (
        f"https://generativelanguage.googleapis.com/v1beta/models/{model}:streamGenerateContent"
        f"?key={api_key}&alt=sse"
    )
    contents = [
        {"role": "user" if m.get("role") == "user" else "model", "parts": [{"text": m.get("content", "")}]}
        for m in messages
    ]
    body = {"contents": contents, "systemInstruction": {"parts": [{"text": system}]}}
    headers = {"Content-Type": "application/json"}

    async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=30.0)) as client:
        async with client.stream("POST", url, headers=headers, json=body) as resp:
            if resp.status_code >= 400:
                err_text = (await resp.aread()).decode("utf-8", "ignore")
                raise RuntimeError(f"HTTP {resp.status_code}: {err_text[:500]}")
            async for line in resp.aiter_lines():
                if job["cancel_event"].is_set():
                    return
                line = line.strip()
                if not line.startswith("data:"):
                    continue
                data_str = line[5:].strip()
                if not data_str:
                    continue
                try:
                    data = json.loads(data_str)
                except Exception:
                    continue
                candidates = data.get("candidates") or []
                if candidates:
                    parts = (candidates[0].get("content") or {}).get("parts") or []
                    if parts and parts[0].get("text"):
                        job["buffer"] += parts[0]["text"]
                        job["updated_at"] = time.time()
                usage = data.get("usageMetadata")
                if usage:
                    job["tokens"] = usage.get("candidatesTokenCount") or usage.get("totalTokenCount") or job["tokens"]
                    job["promptTokens"] = usage.get("promptTokenCount") or job["promptTokens"]
                    job["completionTokens"] = usage.get("candidatesTokenCount") or job["completionTokens"]


async def run_deepai(job: dict, api_key: str, system: str, messages: list):
    prompt = f"system: {system}\n"
    for m in messages:
        prompt += f"{m.get('role','user')}: {m.get('content','')}\n"
    async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=30.0)) as client:
        resp = await client.post(
            "https://api.deepai.org/api/text-generator",
            headers={"api-key": api_key},
            data={"text": prompt},
        )
        data = resp.json()
        job["buffer"] = data.get("output") or data.get("err") or "Erreur DeepAI"
        job["updated_at"] = time.time()


# --------------------------------------------------------------------------
# Recherche Web agentique — serveur MCP communautaire "victor/websearch" (HF Space)
# https://huggingface.co/spaces/victor/websearch
# --------------------------------------------------------------------------
#
# Flux : 1) un petit appel non-streamé au modèle choisi par l'utilisateur décide
# si une recherche est utile et propose 1 à 3 requêtes ; 2) ces requêtes sont
# envoyées EN PARALLÈLE (asyncio.gather) au serveur MCP, qui interroge lui-même
# plusieurs sites web par requête ; 3) les résultats sont injectés dans le
# system prompt avant l'appel réel au modèle. Tout ceci reste invisible pour
# l'utilisateur, qui ne voit qu'un court statut ("Je vais rechercher...").

WEBSEARCH_MCP_SSE_URL = "https://victor-websearch.hf.space/gradio_api/mcp/sse"


class MCPSSEClient:
    """Client MCP minimal (transport SSE, JSON-RPC 2.0) pour interroger un
    serveur MCP externe. Une instance = une session courte, ouverte pour la
    durée d'une recherche puis refermée (pas de session partagée entre
    utilisateurs)."""

    def __init__(self, sse_url: str):
        self.sse_url = sse_url
        self.post_url: Optional[str] = None
        self._client = httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=15.0))
        self._futures: dict[int, asyncio.Future] = {}
        self._id_counter = 0
        self._ready = asyncio.Event()
        self._stream_task: Optional[asyncio.Task] = None

    async def start(self):
        self._stream_task = asyncio.create_task(self._read_stream())
        await asyncio.wait_for(self._ready.wait(), timeout=15)
        await self._request("initialize", {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "chatia-multi-modeles", "version": "9.5"},
        })
        await self._notify("notifications/initialized", {})

    async def _read_stream(self):
        try:
            async with self._client.stream("GET", self.sse_url, headers={"Accept": "text/event-stream"}) as resp:
                event_name, data_lines = None, []
                async for line in resp.aiter_lines():
                    if line.startswith(":"):
                        continue
                    if line.startswith("event:"):
                        event_name = line[6:].strip()
                    elif line.startswith("data:"):
                        data_lines.append(line[5:].strip())
                    elif line.strip() == "":
                        if data_lines:
                            self._handle_event(event_name, "\n".join(data_lines))
                        event_name, data_lines = None, []
        except Exception:
            pass
        finally:
            for fut in self._futures.values():
                if not fut.done():
                    fut.set_exception(RuntimeError("Connexion MCP fermée prématurément"))

    def _handle_event(self, event_name, data):
        if event_name == "endpoint":
            self.post_url = urljoin(self.sse_url, data)
            self._ready.set()
            return
        try:
            msg = json.loads(data)
        except Exception:
            return
        mid = msg.get("id")
        if mid is not None and mid in self._futures and not self._futures[mid].done():
            self._futures[mid].set_result(msg)

    async def _request(self, method, params, timeout=25.0):
        self._id_counter += 1
        mid = self._id_counter
        fut = asyncio.get_event_loop().create_future()
        self._futures[mid] = fut
        try:
            resp = await self._client.post(self.post_url, json={"jsonrpc": "2.0", "id": mid, "method": method, "params": params})
            if resp.status_code >= 400:
                raise RuntimeError(f"MCP HTTP {resp.status_code}")
            msg = await asyncio.wait_for(fut, timeout=timeout)
        finally:
            self._futures.pop(mid, None)
        if "error" in msg:
            raise RuntimeError(f"Erreur MCP: {msg['error']}")
        return msg.get("result")

    async def _notify(self, method, params):
        await self._client.post(self.post_url, json={"jsonrpc": "2.0", "method": method, "params": params})

    async def list_tools(self):
        result = await self._request("tools/list", {})
        return (result or {}).get("tools", [])

    async def call_tool(self, name, arguments):
        return await self._request("tools/call", {"name": name, "arguments": arguments}, timeout=40.0)

    async def close(self):
        if self._stream_task:
            self._stream_task.cancel()
        await self._client.aclose()


def _pick_search_tool(tools: list) -> Optional[dict]:
    for t in tools:
        blob = ((t.get("name") or "") + " " + (t.get("description") or "")).lower()
        if "search" in blob:
            return t
    return tools[0] if tools else None


def _match_arg(props: dict, candidates: list) -> Optional[str]:
    lower_map = {k.lower(): k for k in props}
    for c in candidates:
        if c.lower() in lower_map:
            return lower_map[c.lower()]
    return None


def _extract_mcp_text(result) -> str:
    if not result:
        return ""
    content = result.get("content") if isinstance(result, dict) else None
    if not content:
        return json.dumps(result, ensure_ascii=False)[:4000]
    parts = [item.get("text", "") for item in content if isinstance(item, dict) and item.get("type") == "text"]
    return "\n".join(parts)[:6000]  # borne la taille injectée dans le contexte du modèle


async def run_web_search(queries: list, search_type: str = "search", num_results: int = 4) -> str:
    """Interroge le serveur MCP de victor avec plusieurs requêtes EN PARALLÈLE
    (donc potentiellement plusieurs sites différents par requête, sur plusieurs
    requêtes à la fois) et retourne un texte agrégé prêt à injecter dans le
    contexte du modèle."""
    client = MCPSSEClient(WEBSEARCH_MCP_SSE_URL)
    try:
        await client.start()
        tools = await client.list_tools()
        tool = _pick_search_tool(tools)
        if not tool:
            return ""
        tool_name = tool["name"]
        props = ((tool.get("inputSchema") or {}).get("properties")) or {}
        query_key = _match_arg(props, ["query", "q", "search_query", "text"]) or "query"
        type_key = _match_arg(props, ["search_type", "type", "mode"])
        num_key = _match_arg(props, ["num_results", "n", "count", "max_results", "limit"])

        async def one_search(q: str):
            args = {query_key: q}
            if type_key:
                args[type_key] = search_type
            if num_key:
                args[num_key] = num_results
            try:
                result = await client.call_tool(tool_name, args)
                return q, _extract_mcp_text(result)
            except Exception as e:
                return q, f"[Erreur de recherche pour « {q} » : {e}]"

        results = await asyncio.gather(*(one_search(q) for q in queries))
        return "\n\n".join(f"### Résultats pour : {q}\n{text}" for q, text in results if text)
    finally:
        await client.close()


async def run_single_completion(provider: str, model: str, api_key: str, base_url: str, system: str, user_text: str, max_tokens: int = 250) -> str:
    """Appel non-streamé, court, utilisé uniquement pour la décision agentique
    ("faut-il chercher sur le web, et quoi ?"). Ne consomme pas le buffer du job."""
    async with httpx.AsyncClient(timeout=httpx.Timeout(20.0, connect=10.0)) as client:
        if provider == "gemini":
            url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
            body = {
                "contents": [{"role": "user", "parts": [{"text": user_text}]}],
                "systemInstruction": {"parts": [{"text": system}]},
                "generationConfig": {"maxOutputTokens": max_tokens, "temperature": 0},
            }
            resp = await client.post(url, json=body)
            resp.raise_for_status()
            data = resp.json()
            return data["candidates"][0]["content"]["parts"][0]["text"]

        if provider == "openai":
            url = "https://api.openai.com/v1/chat/completions"
        elif provider == "groq":
            url = "https://api.groq.com/openai/v1/chat/completions"
        elif provider == "custom":
            url = base_url.rstrip("/") + "/chat/completions"
        else:
            raise ValueError(f"Provider non supporté pour la décision de recherche : {provider}")

        headers = {"Content-Type": "application/json"}
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"
        body = {
            "model": model,
            "messages": [{"role": "system", "content": system}, {"role": "user", "content": user_text}],
            "max_tokens": max_tokens, "temperature": 0, "stream": False,
        }
        resp = await client.post(url, headers=headers, json=body)
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"]


WEB_SEARCH_DECISION_PROMPT = (
    "Tu es un module de décision pour un assistant IA. Réponds UNIQUEMENT avec un objet JSON, "
    "sans aucun texte ni markdown autour, au format exact : "
    '{"search": true, "queries": ["...", "..."]} ou {"search": false}. '
    "Mets \"search\" à true seulement si une recherche web est vraiment utile pour bien répondre "
    "(actualité récente, informations qui changent dans le temps, faits précis et vérifiables). "
    "Si search est true, propose entre 1 et 3 requêtes de recherche courtes et précises, "
    "éventuellement sous des angles complémentaires (ex: sites différents, formulations différentes) "
    "pour couvrir plusieurs sources à la fois."
)


async def decide_search_queries(provider: str, model: str, api_key: str, base_url: str, user_text: str) -> list:
    if provider == "deepai" or not user_text.strip():
        return []
    try:
        raw = (await run_single_completion(provider, model, api_key, base_url, WEB_SEARCH_DECISION_PROMPT, user_text)).strip()
        if raw.startswith("```"):
            raw = raw.strip("`")
            if "\n" in raw:
                raw = raw.split("\n", 1)[1]
        data = json.loads(raw)
        if data.get("search") and isinstance(data.get("queries"), list):
            return [q.strip() for q in data["queries"] if isinstance(q, str) and q.strip()][:3]
    except Exception as e:
        print(f"[WARN] Décision de recherche web échouée (on continue sans recherche) : {e}")
    return []


async def maybe_run_agentic_web_search(job: dict, provider: str, model: str, api_key: str, base_url: str, system: str, messages: list) -> str:
    """Retourne le system prompt éventuellement enrichi des résultats de recherche.
    Met à jour job["status_text"] pour que le client affiche un statut ("Je vais
    rechercher...") pendant cette étape, sans jamais exposer le détail des tool calls."""
    last_user_text = ""
    for m in reversed(messages):
        if m.get("role") == "user":
            last_user_text = m.get("content", "")
            break
    if not last_user_text:
        return system

    job["status_text"] = "Je réfléchis à si une recherche web est nécessaire..."
    queries = await decide_search_queries(provider, model, api_key, base_url, last_user_text)
    if not queries:
        job["status_text"] = None
        return system

    job["status_text"] = "Je vais rechercher sur le web, pour vous fournir une réponse complète..."
    try:
        results_text = await run_web_search(queries)
    except Exception as e:
        print(f"[WARN] Recherche web échouée : {e}")
        results_text = ""
    job["status_text"] = None

    if not results_text:
        return system
    return (
        system
        + "\n\n--- Résultats de recherche web (obtenus juste avant ta réponse, plusieurs sites interrogés) ---\n"
        + results_text
        + "\n--- Fin des résultats de recherche web ---\n"
        + "Utilise ces informations pour répondre de façon complète et à jour. "
        + "Ne mentionne pas explicitement que tu as \"utilisé un outil\" ou \"MCP\" ; réponds naturellement."
    )


async def run_job(job_id: str, params: dict):
    job = JOBS[job_id]
    try:
        provider = params["provider"]
        model = params["model"]
        system = params.get("system", "")
        messages = params.get("messages", [])
        api_key = params.get("apiKey", "")
        base_url = params.get("baseUrl", "")

        if params.get("webSearchEnabled"):
            system = await maybe_run_agentic_web_search(job, provider, model, api_key, base_url, system, messages)

        if provider == "openai":
            await run_openai_like(job, "https://api.openai.com/v1/chat/completions", api_key, model, system, messages)
        elif provider == "groq":
            await run_openai_like(job, "https://api.groq.com/openai/v1/chat/completions", api_key, model, system, messages)
        elif provider == "custom":
            url = base_url.rstrip("/") + "/chat/completions"
            await run_openai_like(job, url, api_key, model, system, messages)
        elif provider == "gemini":
            await run_gemini(job, api_key, model, system, messages)
        elif provider == "deepai":
            await run_deepai(job, api_key, system, messages)
        else:
            raise ValueError(f"Provider inconnu: {provider}")

        job["status"] = "cancelled" if job["cancel_event"].is_set() else "done"
    except Exception as e:
        job["status"] = "error"
        job["error"] = str(e)
    finally:
        job["status_text"] = None
        job["updated_at"] = time.time()
        try:
            await persist_job_result(job)
        except Exception as e:
            print(f"[WARN] Échec de persistance du job {job_id}: {e}")


async def cleanup_jobs_loop():
    while True:
        await asyncio.sleep(300)
        now = time.time()
        to_delete = [
            jid for jid, j in JOBS.items()
            if j["status"] != "running" and (now - j["updated_at"]) > JOB_GRACE_SECONDS
        ]
        for jid in to_delete:
            JOBS.pop(jid, None)


# --------------------------------------------------------------------------
# FastAPI app
# --------------------------------------------------------------------------

app = FastAPI(title="Chat IA - Serveur de Sync V9.4")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
async def on_startup():
    asyncio.create_task(cleanup_jobs_loop())


# ---- Auth ----

@app.post("/register")
async def register(body: AuthBody):
    username = body.username.strip()
    if not username or not body.password:
        raise HTTPException(status_code=400, detail="Champs requis.")
    if len(body.password) < 4:
        raise HTTPException(status_code=400, detail="Mot de passe trop court (4 caractères min).")
    users = await load_users()
    if username in users:
        raise HTTPException(status_code=409, detail="Ce nom d'utilisateur existe déjà.")
    pw_hash = bcrypt.hashpw(body.password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")
    users[username] = {"password_hash": pw_hash, "created_at": time.time()}
    await save_users(users)
    await save_state(username, json.loads(json.dumps(EMPTY_STATE)))
    return {"ok": True}


@app.post("/login")
async def login(body: AuthBody):
    username = body.username.strip()
    users = await load_users()
    user = users.get(username)
    if not user or not bcrypt.checkpw(body.password.encode("utf-8"), user["password_hash"].encode("utf-8")):
        raise HTTPException(status_code=401, detail="Identifiants incorrects.")
    return {"token": make_token(username)}


# ---- Sync ----

@app.get("/sync")
async def get_sync(authorization: Optional[str] = Header(None)):
    username = verify_token(authorization)
    data = await load_state(username)
    return {"data": data}


@app.post("/sync")
async def post_sync(body: SyncBody, authorization: Optional[str] = Header(None)):
    username = verify_token(authorization)
    await save_state(username, body.data)
    return {"ok": True}


class BeaconSyncBody(BaseModel):
    token: str
    data: dict


@app.post("/sync/beacon")
async def post_sync_beacon(body: BeaconSyncBody):
    # navigator.sendBeacon() ne permet pas de fixer un header Authorization,
    # le token est donc transmis dans le corps de la requête pour ce cas précis
    # (utilisé uniquement comme filet de sécurité à la fermeture de l'onglet).
    username = verify_token(f"Bearer {body.token}")
    await save_state(username, body.data)
    return {"ok": True}


# ---- Génération (proxy streaming résilient) ----

class GenerateBody(BaseModel):
    chat_id: object
    chat_title: Optional[str] = ""
    message_id: str
    provider: str
    model: str
    apiKey: Optional[str] = ""
    baseUrl: Optional[str] = ""
    system: Optional[str] = ""
    messages: list
    recent_messages: Optional[list] = None
    webSearchEnabled: Optional[bool] = False


@app.post("/api/generate")
async def start_generate(body: GenerateBody, authorization: Optional[str] = Header(None)):
    username = verify_token(authorization)

    if body.recent_messages:
        await upsert_pending_message(username, body.chat_id, body.chat_title, body.recent_messages)

    job_id = new_job(username, body.chat_id, body.message_id, body.provider, body.model)
    params = {
        "provider": body.provider, "model": body.model, "system": body.system or "",
        "messages": body.messages, "apiKey": body.apiKey or "", "baseUrl": body.baseUrl or "",
        "webSearchEnabled": bool(body.webSearchEnabled),
    }
    asyncio.create_task(run_job(job_id, params))
    return {"job_id": job_id}


@app.post("/api/generate/cancel/{job_id}")
async def cancel_generate(job_id: str, authorization: Optional[str] = Header(None)):
    username = verify_token(authorization)
    job = JOBS.get(job_id)
    if not job or job["username"] != username:
        raise HTTPException(status_code=404, detail="Job introuvable")
    job["cancel_event"].set()
    return {"ok": True}


@app.get("/api/generate/active")
async def active_jobs(authorization: Optional[str] = Header(None)):
    username = verify_token(authorization)
    jobs = [
        {"job_id": j["id"], "chat_id": j["chat_id"], "message_id": j["message_id"],
         "provider": j["provider"], "model": j["model"]}
        for j in JOBS.values() if j["username"] == username and j["status"] == "running"
    ]
    return {"jobs": jobs}


@app.get("/api/generate/stream/{job_id}")
async def stream_generate(job_id: str, authorization: Optional[str] = Header(None)):
    username = verify_token(authorization)
    job = JOBS.get(job_id)
    if not job or job["username"] != username:
        raise HTTPException(status_code=404, detail="Job introuvable ou expiré")

    async def event_gen():
        # 1) Rattrapage immédiat de ce qui a déjà été généré.
        sent_len = len(job["buffer"])
        yield _sse("sync", {"content": job["buffer"], "tokens": job.get("tokens")})
        last_status_sent = None
        if job.get("status_text"):
            last_status_sent = job["status_text"]
            yield _sse("status", {"text": last_status_sent})

        # 2) Puis on suit les nouveaux morceaux au fur et à mesure (+ les statuts
        #    agentiques transitoires, ex: "Je vais rechercher sur le web...").
        while True:
            await asyncio.sleep(0.1)

            current_status = job.get("status_text")
            if current_status != last_status_sent:
                last_status_sent = current_status
                yield _sse("status", {"text": current_status or ""})

            current = job["buffer"]
            if len(current) > sent_len:
                new_text = current[sent_len:]
                sent_len = len(current)
                yield _sse("delta", {"text": new_text})

            if job["status"] != "running":
                # on vide le dernier reste éventuel avant de conclure
                current = job["buffer"]
                if len(current) > sent_len:
                    yield _sse("delta", {"text": current[sent_len:]})
                    sent_len = len(current)
                if job["status"] == "error":
                    yield _sse("error", {"message": job.get("error") or "Erreur inconnue"})
                else:
                    yield _sse("done", {
                        "tokens": job.get("tokens"),
                        "promptTokens": job.get("promptTokens"),
                        "completionTokens": job.get("completionTokens"),
                    })
                break

    return StreamingResponse(
        event_gen(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive"},
    )


# ---- Fichiers statiques : sert le client HTML directement depuis le Space ----

STATIC_DIR = Path(__file__).parent / "static"
if STATIC_DIR.exists():
    app.mount("/assets", StaticFiles(directory=str(STATIC_DIR)), name="assets")

    @app.get("/")
    async def serve_index():
        index_file = STATIC_DIR / "index.html"
        if index_file.exists():
            return FileResponse(str(index_file))
        return JSONResponse({"status": "ok", "info": "Placez V9_4.html dans server/static/index.html"})


@app.get("/health")
async def health():
    return {"status": "ok", "jobs_running": sum(1 for j in JOBS.values() if j["status"] == "running")}