#!/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"))
ENABLE_OLLAMA = os.getenv("ENABLE_OLLAMA", "true").lower() == "true"
LISTEN_PORT = int(os.getenv("LISTEN_PORT", "3014"))
SWITCH_TIMEOUT = int(os.getenv("SWITCH_TIMEOUT", "90"))
CHECK_INTERVAL = int(os.getenv("CHECK_INTERVAL", "2"))

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()


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


@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:
        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:
            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)