#!/usr/bin/env python3
"""
Jean Router — Proxy dynamique pour Jean presets
Interroge Jean pour lister les presets disponibles et route les requêtes
vers le bon preset en faisant le switch automatique si nécessaire.
"""

import os
import time
import logging
import threading
from typing import Optional

import requests
from flask import Flask, request, Response, jsonify

JEAN_WEB_HOST = os.getenv("JEAN_WEB_HOST", "192.168.1.11")
JEAN_WEB_PORT = int(os.getenv("JEAN_WEB_PORT", "8090"))
JEAN_PROXY_HOST = os.getenv("JEAN_PROXY_HOST", "localhost")
JEAN_PROXY_PORT = int(os.getenv("JEAN_PROXY_PORT", "3012"))
OLLAMA_PROXY_HOST = os.getenv("OLLAMA_PROXY_HOST", "localhost")
OLLAMA_PROXY_PORT = int(os.getenv("OLLAMA_PROXY_PORT", "3011"))
OLLAMA_DIRECT_HOST = os.getenv("OLLAMA_DIRECT_HOST", JEAN_WEB_HOST)
OLLAMA_DIRECT_PORT = int(os.getenv("OLLAMA_DIRECT_PORT", "11434"))
ENABLE_OLLAMA = os.getenv("ENABLE_OLLAMA", "true").lower() == "true"
ENABLE_VRAM_MANAGEMENT = os.getenv("ENABLE_VRAM_MANAGEMENT", "true").lower() == "true"
OLLAMA_KEEP_ALIVE_MINUTES = int(os.getenv("OLLAMA_KEEP_ALIVE_MINUTES", "5"))
LISTEN_PORT = int(os.getenv("LISTEN_PORT", "3014"))
SWITCH_TIMEOUT = int(os.getenv("SWITCH_TIMEOUT", "90"))
CHECK_INTERVAL = int(os.getenv("CHECK_INTERVAL", "2"))
JEAN_STOP_TIMEOUT = int(os.getenv("JEAN_STOP_TIMEOUT", "30"))
OLLAMA_UNLOAD_TIMEOUT = int(os.getenv("OLLAMA_UNLOAD_TIMEOUT", "15"))

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger(__name__)

app = Flask(__name__)

presets_cache = {"data": [], "timestamp": 0}
ollama_cache = {"data": [], "timestamp": 0}
ollama_detail_cache = {}
jean_ctx_cache = {}
jean_caps_cache = {}
CACHE_TTL = int(os.getenv("CACHE_TTL", "15"))
OLLAMA_DETAIL_CACHE_TTL = 300
switch_lock = threading.Lock()
vram_lock = threading.Lock()
vram_state = {"active_backend": "jean"}
last_active_preset = None
ollama_unload_timer = None


def get_presets(force_refresh=False):
    """Récupère la liste des presets depuis Jean"""
    now = time.time()
    if not force_refresh and presets_cache["data"] and (now - presets_cache["timestamp"]) < CACHE_TTL:
        return presets_cache["data"]

    try:
        url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/api/presets"
        r = requests.get(url, timeout=3)
        if r.ok:
            presets_cache["data"] = r.json()
            presets_cache["timestamp"] = now
            return presets_cache["data"]
    except Exception as e:
        log.warning(f"[Presets] Erreur récupération: {e}")

    return presets_cache["data"]


def get_ollama_models(force_refresh=False):
    """Récupère la liste des modèles depuis Ollama"""
    if not ENABLE_OLLAMA:
        return []

    now = time.time()
    if not force_refresh and ollama_cache["data"] and (now - ollama_cache["timestamp"]) < CACHE_TTL:
        return ollama_cache["data"]

    try:
        url = f"http://{OLLAMA_PROXY_HOST}:{OLLAMA_PROXY_PORT}/api/tags"
        r = requests.get(url, timeout=3)
        if r.ok:
            data = r.json().get("models", [])
            ollama_cache["data"] = data
            ollama_cache["timestamp"] = now
            return data
    except Exception as e:
        log.warning(f"[Ollama] Erreur récupération: {e}")

    return ollama_cache["data"]


def _get_ollama_model_detail(model_name: str) -> dict:
    """Récupère les détails d'un modèle Ollama via /api/show, avec cache long"""
    now = time.time()
    cached = ollama_detail_cache.get(model_name)
    if cached and (now - cached["timestamp"]) < OLLAMA_DETAIL_CACHE_TTL:
        return cached["data"]

    detail = {}
    try:
        clean_name = model_name.replace(":latest", "")
        url = f"http://{OLLAMA_DIRECT_HOST}:{OLLAMA_DIRECT_PORT}/api/show"
        r = requests.post(url, json={"model": clean_name}, timeout=10)
        if r.ok:
            d = r.json()
            model_info = d.get("model_info", {})
            details = d.get("details", {})

            for key, val in model_info.items():
                if "context_length" in key.lower() or "ctx" in key.lower():
                    detail["context_length"] = int(val)
                    break

            if not detail.get("context_length"):
                detail["context_length"] = details.get("context_length")

            detail["capabilities"] = d.get("capabilities", [])
            detail["family"] = details.get("family", "")

            ollama_detail_cache[model_name] = {"data": detail, "timestamp": now}
    except Exception as e:
        log.debug(f"[Ollama] /api/show echoue pour {model_name}: {e}")

    if model_name not in ollama_detail_cache:
        ollama_detail_cache[model_name] = {"data": {}, "timestamp": now}

    return ollama_detail_cache[model_name]["data"]


def _enrich_ollama_model(m: dict) -> dict:
    """Enrichit un modele Ollama avec les donnees /api/show si necessaire"""
    model_id = m.get("name", "").replace(":latest", "")
    details = m.get("details", {})
    has_ctx = bool(details.get("context_length"))
    caps = m.get("capabilities", [])

    if not has_ctx or len(caps) <= 1:
        extra = _get_ollama_model_detail(model_id)
        if extra.get("context_length"):
            details = dict(details)
            details["context_length"] = extra["context_length"]
            m = dict(m)
            m["details"] = details
        if extra.get("capabilities") and len(extra["capabilities"]) > len(caps):
            m = dict(m)
            m["capabilities"] = extra["capabilities"]
        if extra.get("family") and not details.get("family"):
            details = dict(details)
            details["family"] = extra["family"]
            m = dict(m)
            m["details"] = details

    return m

def _guess_context_length(preset_id: str) -> int:
    """Devine le context_length depuis le nom du preset"""
    pid = preset_id.lower()
    if "qwen3.6" in pid or "qwen3.5" in pid:
        return 262144
    if "qwen3" in pid:
        return 40960
    if "qwen2.5" in pid:
        return 32768
    if "mistral" in pid or "nemo" in pid:
        return 1024000
    if "llama" in pid:
        return 131072
    if "ornith" in pid:
        return 131072
    if "gemma" in pid:
        return 262144
    if "deepseek" in pid:
        return 131072
    if "phi" in pid:
        return 16384
    if "35b" in pid or "27b" in pid:
        return 131072
    return 32768


def _guess_capabilities(preset_id: str) -> list:
    """Devine les capabilities depuis le nom du preset + metadonnees Jean"""
    pid = preset_id.lower()
    caps = ["completion"]

    has_vision = any(kw in pid for kw in ["vision", "vl", "llava", "bakllava", "minicpm"])
    has_thinking = any(kw in pid for kw in ["reason", "thinking", "deepseek-r1"])
    has_tools = not any(kw in pid for kw in ["embed", "bge-", "nomic-"])

    if has_vision:
        caps.append("vision")
    if has_thinking:
        caps.append("thinking")
    if has_tools:
        caps.append("tools")

    return caps


def _match_preset_to_ollama(preset_id: str):
    """Tente de trouver un modele Ollama correspondant au preset Jean"""
    ollama_models = get_ollama_models()
    pid_lower = preset_id.lower()

    def norm(s):
        return s.lower().replace(":", "").replace("-", "").replace(" ", "").replace(".", "")

    pid_norm = norm(pid_lower)

    for m in ollama_models:
        name = m.get("name", "").replace(":latest", "")
        name_norm = norm(name)

        if pid_norm == name_norm:
            return m

    best = None
    best_score = 0
    for m in ollama_models:
        name = m.get("name", "").replace(":latest", "")
        name_norm = norm(name)
        score = 0

        if len(pid_norm) >= 4 and pid_norm[:4] == name_norm[:4]:
            score += 10
        if len(pid_norm) >= 8 and pid_norm[:8] == name_norm[:8]:
            score += 20
        if name_norm in pid_norm or pid_norm in name_norm:
            score += 15

        parts_pid = set(pid_norm.split("_"))
        parts_name = set(name_norm.split("_"))
        common = parts_pid & parts_name
        score += len(common) * 5

        if score > best_score:
            best_score = score
            best = m

    if best and best_score >= 5:
        return best
    return None


def get_jean_metadata():
    """Resout dynamiquement context_length et capabilities pour les presets Jean.
    1) Ollama matching → 2) Jean /v1/models → 3) heuristique"""
    now = time.time()

    try:
        url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/v1/models"
        r = requests.get(url, timeout=2)
        if r.ok:
            for m in r.json().get("data", []):
                alias = m.get("id", "")
                meta = m.get("meta", {})
                n_ctx = meta.get("n_ctx") or m.get("max_model_len")
                if n_ctx and alias:
                    jean_ctx_cache[alias] = int(n_ctx)
    except Exception:
        pass

    ollama_models = get_ollama_models()
    ollama_lookup = {}
    for m in ollama_models:
        m = _enrich_ollama_model(m)
        name = m.get("name", "").replace(":latest", "")
        details = m.get("details", {})
        ctx = details.get("context_length")
        caps = m.get("capabilities", [])
        ollama_lookup[name] = {"context_length": ctx, "capabilities": caps}

    presets = get_presets()
    for p in presets:
        pid = p["id"]

        matched_ollama = _match_preset_to_ollama(pid)
        ollama_name = matched_ollama.get("name", "").replace(":latest", "") if matched_ollama else None

        if pid not in jean_ctx_cache:
            if ollama_name and ollama_lookup.get(ollama_name, {}).get("context_length"):
                jean_ctx_cache[pid] = ollama_lookup[ollama_name]["context_length"]
            else:
                jean_ctx_cache[pid] = _guess_context_length(pid)

        if pid not in jean_caps_cache:
            if p.get("reasoning") == "on":
                caps = ["completion", "tools", "thinking"]
            else:
                caps = ["completion", "tools"]

            if ollama_name and ollama_lookup.get(ollama_name, {}).get("capabilities"):
                ollama_caps = ollama_lookup[ollama_name]["capabilities"]
                caps = list(dict.fromkeys(caps + ollama_caps))

            jean_caps_cache[pid] = caps

    return jean_ctx_cache, jean_caps_cache


def get_active_preset():
    """Retourne le preset actuellement actif"""
    presets = get_presets()
    for p in presets:
        if p.get("active"):
            return p
    return None


def switch_preset(preset_id: str) -> bool:
    """Switch vers un preset spécifique par son ID"""
    global last_active_preset
    presets = get_presets(force_refresh=True)

    for i, p in enumerate(presets):
        if p["id"] == preset_id:
            if p.get("active"):
                log.info(f"[Switch] Preset {preset_id} déjà actif")
                return True

            log.info(f"[Switch] Activation preset {preset_id} (index {i+1})")
            success = False
            for attempt in range(2):
                try:
                    url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/api/switch"
                    r = requests.post(url, json={"n": i + 1}, timeout=15)
                    if r.ok:
                        success = True
                        break
                    log.error(f"[Switch] Erreur API (tentative {attempt+1}): {r.text}")
                except Exception as e:
                    log.error(f"[Switch] Erreur requête (tentative {attempt+1}): {e}")
                if attempt == 0:
                    time.sleep(3)
            if not success:
                return False

            log.info(f"[Switch] Attente chargement modèle...")
            deadline = time.time() + SWITCH_TIMEOUT
            while time.time() < deadline:
                try:
                    health_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/health"
                    hr = requests.get(health_url, timeout=2)
                    if hr.ok:
                        models_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/v1/models"
                        mr = requests.get(models_url, timeout=2)
                        if mr.ok:
                            models = mr.json().get("data", [])
                            if models:
                                log.info(f"[Switch] Modèle chargé après {int(time.time() - deadline + SWITCH_TIMEOUT)}s")
                                last_active_preset = preset_id
                                presets_cache["timestamp"] = 0
                                return True
                except Exception:
                    pass
                time.sleep(CHECK_INTERVAL)

            log.error(f"[Switch] Timeout après {SWITCH_TIMEOUT}s")
            return False

    log.error(f"[Switch] Preset {preset_id} non trouvé")
    return False


def stop_jean() -> bool:
    """Arrête le service Jean pour libérer la VRAM"""
    log.info("[VRAM] Arrêt de Jean pour libérer la VRAM...")

    try:
        url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/api/stop"
        r = requests.post(url, timeout=5)
        if r.ok:
            log.info("[VRAM] Jean arrêté via API web")
        else:
            log.warning(f"[VRAM] API stop a retourné {r.status_code}, tentative via health check")
    except Exception as e:
        log.warning(f"[VRAM] Erreur API stop: {e}")

    deadline = time.time() + JEAN_STOP_TIMEOUT
    while time.time() < deadline:
        try:
            health_url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/health"
            r = requests.get(health_url, timeout=2)
            if not r.ok:
                log.info("[VRAM] Jean confirmé arrêté")
                return True
        except Exception:
            log.info("[VRAM] Jean confirmé arrêté (connexion refusée)")
            return True
        time.sleep(CHECK_INTERVAL)

    log.error(f"[VRAM] Timeout: Jean ne s'est pas arrêté après {JEAN_STOP_TIMEOUT}s")
    return False


def start_jean() -> bool:
    """Démarre le service Jean"""
    log.info("[VRAM] Démarrage de Jean...")

    try:
        url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/api/start"
        r = requests.post(url, timeout=5)
        if r.ok:
            log.info("[VRAM] Jean démarré via API web")
        else:
            log.warning(f"[VRAM] API start a retourné {r.status_code}")
    except Exception as e:
        log.warning(f"[VRAM] Erreur API start: {e}")

    log.info("[VRAM] Attente du chargement du modèle Jean...")
    deadline = time.time() + SWITCH_TIMEOUT
    while time.time() < deadline:
        try:
            health_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/health"
            hr = requests.get(health_url, timeout=2)
            if hr.ok:
                models_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/v1/models"
                mr = requests.get(models_url, timeout=2)
                if mr.ok:
                    models = mr.json().get("data", [])
                    if models:
                        log.info(f"[VRAM] Jean prêt après {int(time.time() - deadline + SWITCH_TIMEOUT)}s")
                        return True
        except Exception:
            pass
        time.sleep(CHECK_INTERVAL)

    log.error(f"[VRAM] Timeout: Jean n'a pas démarré après {SWITCH_TIMEOUT}s")
    return False


def unload_ollama() -> bool:
    """Force le déchargement des modèles Ollama pour libérer la VRAM"""
    log.info("[VRAM] Déchargement des modèles Ollama...")

    try:
        url = f"http://{OLLAMA_DIRECT_HOST}:{OLLAMA_DIRECT_PORT}/api/generate"
        payload = {
            "model": "mistral",
            "keep_alive": 0,
            "prompt": ""
        }
        r = requests.post(url, json=payload, timeout=5)
        if r.ok:
            log.info("[VRAM] Ollama déchargé via API")
            return True
        else:
            log.warning(f"[VRAM] API unload a retourné {r.status_code}")
    except Exception as e:
        log.warning(f"[VRAM] Erreur API unload: {e}")

    try:
        url = f"http://{OLLAMA_DIRECT_HOST}:{OLLAMA_DIRECT_PORT}/api/ps"
        r = requests.get(url, timeout=3)
        if r.ok:
            models = r.json().get("models", [])
            if not models:
                log.info("[VRAM] Ollama confirmé déchargé (aucun modèle actif)")
                return True
            else:
                log.warning(f"[VRAM] Ollama a encore {len(models)} modèle(s) chargé(s)")
    except Exception as e:
        log.warning(f"[VRAM] Impossible de vérifier /api/ps: {e}")

    time.sleep(2)
    log.info("[VRAM] Ollama supposé déchargé (timeout)")
    return True


def schedule_ollama_unload():
    """Programme le déchargement d'Ollama après OLLAMA_KEEP_ALIVE_MINUTES"""
    global ollama_unload_timer

    def delayed_unload():
        log.info(f"[VRAM] Timer écoulé ({OLLAMA_KEEP_ALIVE_MINUTES}min), déchargement Ollama...")
        unload_ollama()

    if ollama_unload_timer:
        ollama_unload_timer.cancel()
        log.info("[VRAM] Timer précédent annulé")

    ollama_unload_timer = threading.Timer(
        OLLAMA_KEEP_ALIVE_MINUTES * 60,
        delayed_unload
    )
    ollama_unload_timer.daemon = True
    ollama_unload_timer.start()
    log.info(f"[VRAM] Déchargement Ollama programmé dans {OLLAMA_KEEP_ALIVE_MINUTES} minutes")


def cancel_ollama_unload():
    """Annule le déchargement programmé d'Ollama"""
    global ollama_unload_timer
    if ollama_unload_timer:
        ollama_unload_timer.cancel()
        ollama_unload_timer = None
        log.info("[VRAM] Timer de déchargement Ollama annulé")


def ensure_vram_for_backend(target_backend: str) -> bool:
    """
    S'assure que la VRAM est disponible pour le backend cible.
    - Si target=ollama et jean actif → arrête Jean, annule timer unload Ollama
    - Si target=jean et ollama actif → programme déchargement Ollama dans X min, démarre Jean
    """
    if not ENABLE_VRAM_MANAGEMENT:
        return True

    global last_active_preset

    with vram_lock:
        current = vram_state["active_backend"]

        if current == target_backend:
            return True

        if target_backend == "ollama" and current == "jean":
            log.info("[VRAM] Switch Jean → Ollama")
            cancel_ollama_unload()
            if not stop_jean():
                return False
            last_active_preset = None
            vram_state["active_backend"] = "ollama"
            return True

        elif target_backend == "jean" and current == "ollama":
            log.info("[VRAM] Switch Ollama → Jean")
            schedule_ollama_unload()
            if not start_jean():
                return False
            last_active_preset = None
            vram_state["active_backend"] = "jean"
            return True

        return True


@app.route("/v1/models", methods=["GET"])
def list_models():
    """Expose tous les presets Jean + modèles Ollama comme modèles OpenAI"""
    force_refresh = request.args.get("refresh", "").lower() in ("1", "true", "yes")
    if force_refresh:
        global presets_cache, ollama_cache
        presets_cache["timestamp"] = 0
        ollama_cache["timestamp"] = 0
        log.info("[Models] Refresh force demande")

    presets = get_presets()
    jean_ctx, jean_caps = get_jean_metadata()
    models = []

    for p in presets:
        model_id = p["id"]
        model_name = p.get("name", model_id)

        ctx_len = jean_ctx.get(model_id)
        caps = jean_caps.get(model_id, ["completion", "tools"])
        if p.get("reasoning") == "on" and "thinking" not in caps:
            caps = list(dict.fromkeys(caps + ["thinking"]))

        model_info = {
            "id": model_name,
            "object": "model",
            "created": int(time.time()),
            "owned_by": "jean",
            "permission": [],
            "root": model_name,
            "parent": None,
        }

        if ctx_len:
            model_info["context_length"] = ctx_len
            model_info["context_window"] = ctx_len
            model_info["max_model_len"] = ctx_len
            model_info["max_tokens"] = ctx_len
            model_info["max_input_tokens"] = ctx_len

        model_info["capabilities"] = caps

        meta = {}
        if p.get("name") and p["name"] != model_id:
            meta["display_name"] = p["name"]
        if p.get("quant"):
            meta["quantization"] = p["quant"]
        if p.get("reasoning"):
            meta["reasoning"] = p["reasoning"]
        if p.get("active"):
            meta["active"] = True
        if p.get("bench"):
            meta["bench"] = p["bench"]
        if ctx_len:
            meta["context_length"] = ctx_len
        meta["capabilities"] = caps

        if meta:
            model_info["meta"] = meta

        models.append(model_info)

    if ENABLE_OLLAMA:
        ollama_models = get_ollama_models()
        for m in ollama_models:
            m = _enrich_ollama_model(m)
            model_id = m.get("name", "").replace(":latest", "")
            if not model_id:
                continue

            caps = m.get("capabilities", ["completion"])
            details = m.get("details", {})
            ctx_len = details.get("context_length")

            model_info = {
                "id": model_id,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "ollama",
                "permission": [],
                "root": model_id,
                "parent": None,
            }

            if ctx_len:
                model_info["context_length"] = ctx_len
                model_info["context_window"] = ctx_len
                model_info["max_model_len"] = ctx_len
                model_info["max_tokens"] = ctx_len
                model_info["max_input_tokens"] = ctx_len

            model_info["capabilities"] = caps

            meta = {}
            if details.get("parameter_size"):
                meta["parameter_size"] = details["parameter_size"]
            if details.get("quantization_level"):
                meta["quantization"] = details["quantization_level"]
            if details.get("family"):
                meta["family"] = details["family"]
            if ctx_len:
                meta["context_length"] = ctx_len
            meta["capabilities"] = caps

            if meta:
                model_info["meta"] = meta

            models.append(model_info)

    return jsonify({"object": "list", "data": models})


@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    """Route les requêtes chat vers Jean ou Ollama selon le modèle"""
    data = request.get_json()
    model_id = data.get("model")

    if not model_id:
        return jsonify({"error": "model parameter required"}), 400

    presets = get_presets()
    jean_preset_ids = {p["id"] for p in presets}
    jean_name_to_id = {p.get("name", p["id"]): p["id"] for p in presets}
    jean_names = set(jean_name_to_id.keys())

    if model_id in jean_preset_ids or model_id in jean_names:
        internal_id = jean_name_to_id.get(model_id, model_id)

        if not ensure_vram_for_backend("jean"):
            return jsonify({"error": "Failed to prepare VRAM for Jean"}), 503

        with switch_lock:
            active = get_active_preset()
            current_active = active["id"] if active else last_active_preset
            if current_active != internal_id:
                log.info(f"[Router] Switch Jean requis: {current_active or 'none'} → {internal_id} ({model_id})")
                if not switch_preset(internal_id):
                    return jsonify({"error": f"Failed to switch to preset {model_id}"}), 503
            else:
                log.info(f"[Router] Preset {model_id} déjà actif (tracké)")

        preset = next((p for p in presets if p["id"] == internal_id), {})
        if preset.get("reasoning") == "on":
            ctx, _ = get_jean_metadata()
            ctx_len = ctx.get(internal_id, 131072)
            req_max = data.get("max_tokens", 0)
            if req_max and req_max > 0:
                new_max = min(req_max * 2, ctx_len)
                if new_max != req_max:
                    log.info(f"[Reasoning] max_tokens doublé: {req_max} → {new_max}")
                    data = dict(data)
                    data["max_tokens"] = new_max

        upstream_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/v1/chat/completions"
        log.info(f"[Router] → Jean ({model_id})")
    elif ENABLE_OLLAMA:
        ollama_models = get_ollama_models()
        ollama_ids = [m.get("name", "").replace(":latest", "") for m in ollama_models]

        if model_id in ollama_ids:
            if not ensure_vram_for_backend("ollama"):
                return jsonify({"error": "Failed to prepare VRAM for Ollama"}), 503

            cancel_ollama_unload()

            upstream_url = f"http://{OLLAMA_PROXY_HOST}:{OLLAMA_PROXY_PORT}/v1/chat/completions"
            log.info(f"[Router] → Ollama ({model_id})")
        else:
            return jsonify({"error": f"Unknown model: {model_id}"}), 404
    else:
        return jsonify({"error": f"Unknown model: {model_id}"}), 404

    headers = dict(request.headers)
    headers.pop("Host", None)

    try:
        resp = requests.request(
            method=request.method,
            url=upstream_url,
            headers=headers,
            json=data,
            stream=True,
            timeout=None,
        )
    except Exception as e:
        log.error(f"[Router] Erreur upstream: {e}")
        return jsonify({"error": str(e)}), 502

    excluded = {"content-encoding", "content-length", "transfer-encoding", "connection"}
    response_headers = [
        (name, value) for name, value in resp.headers.items()
        if name.lower() not in excluded
    ]

    return Response(
        resp.raw,
        status=resp.status_code,
        headers=response_headers,
        direct_passthrough=True,
    )


@app.route("/", defaults={"path": ""}, methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
def proxy(path):
    """Proxy générique pour les autres endpoints"""
    if path.startswith("v1/chat/completions"):
        return chat_completions()

    upstream_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/{path}"
    headers = dict(request.headers)
    headers.pop("Host", None)

    try:
        resp = requests.request(
            method=request.method,
            url=upstream_url,
            headers=headers,
            data=request.get_data(),
            cookies=request.cookies,
            params=request.args,
            allow_redirects=False,
            stream=True,
            timeout=None,
        )
    except Exception as e:
        return Response(str(e), status=502)

    excluded = {"content-encoding", "content-length", "transfer-encoding", "connection"}
    response_headers = [
        (name, value) for name, value in resp.headers.items()
        if name.lower() not in excluded
    ]

    return Response(
        resp.raw,
        status=resp.status_code,
        headers=response_headers,
        direct_passthrough=True,
    )


if __name__ == "__main__":
    log.info(f"[Router] Démarrage sur port {LISTEN_PORT}")
    log.info(f"[Router] Jean Web: {JEAN_WEB_HOST}:{JEAN_WEB_PORT}")
    log.info(f"[Router] Jean Proxy: {JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}")
    if ENABLE_OLLAMA:
        log.info(f"[Router] Ollama Proxy: {OLLAMA_PROXY_HOST}:{OLLAMA_PROXY_PORT}")
    else:
        log.info(f"[Router] Ollama: désactivé")
    app.run(host="0.0.0.0", port=LISTEN_PORT, threaded=True)