#!/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}
CACHE_TTL = 30
switch_lock = threading.Lock()
vram_lock = threading.Lock()
vram_state = {"active_backend": "jean"}
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_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"""
    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})")
            try:
                url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/api/switch"
                r = requests.post(url, json={"n": i + 1}, timeout=5)
                if not r.ok:
                    log.error(f"[Switch] Erreur API: {r.text}")
                    return False
            except Exception as e:
                log.error(f"[Switch] Erreur requête: {e}")
                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")
                                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

    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
            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
            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"""
    presets = get_presets()
    models = []

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

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

        meta = {}
        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 meta:
            model_info["meta"] = meta

        models.append(model_info)

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

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

            details = m.get("details", {})
            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 m.get("capabilities"):
                meta["capabilities"] = m["capabilities"]

            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]

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

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

        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)