import os
import json
import time
import socket
import logging
import threading
from urllib.parse import urljoin

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

PC_HOST = os.getenv("PC_HOST", "192.168.1.11")
PC_MAC = os.getenv("PC_MAC", "")
BROADCAST_IP = os.getenv("BROADCAST_IP", "192.168.1.255")
WAKE_TIMEOUT = int(os.getenv("WAKE_TIMEOUT", "300"))
SERVICE_TIMEOUT = int(os.getenv("SERVICE_TIMEOUT", "120"))
CHECK_INTERVAL = int(os.getenv("CHECK_INTERVAL", "5"))
STATE_FILE = os.getenv("STATE_FILE", "/app/state/active.json")
LISTEN_UI = int(os.getenv("LISTEN_UI", "3017"))

MODEL_ALIAS = os.getenv("MODEL_ALIAS", "qwen3.8-27b")

LLM_INSTANCES = [
    {"port": int(os.getenv("LLAMACPP_A_PORT", "8080")),
     "real": os.getenv("LLAMACPP_A_MODEL", "qwen3.8-27b"),
     "label": os.getenv("LLAMACPP_A_LABEL", "Q4_K_M")},
    {"port": int(os.getenv("LLAMACPP_B_PORT", "8081")),
     "real": os.getenv("LLAMACPP_B_MODEL", "qwen3.8-27b-iq4xs"),
     "label": os.getenv("LLAMACPP_B_LABEL", "IQ4_XS")},
]
DEFAULT_ACTIVE = os.getenv("DEFAULT_ACTIVE", "qwen3.8-27b-iq4xs")

ROUTES = [
    {"name": "ollama", "port": int(os.getenv("OLLAMA_PORT", "11434")),
     "listen": int(os.getenv("LISTEN_OLLAMA", "3011")), "health": "/api/version", "wol": True},
    {"name": "llamacpp", "port": None,
     "listen": int(os.getenv("LISTEN_LLAMACPP", "3012")), "health": "/health", "wol": True,
     "alias": MODEL_ALIAS},
    {"name": "stt", "port": int(os.getenv("STT_PORT", "8200")),
     "listen": int(os.getenv("LISTEN_STT", "3015")), "health": "/health", "wol": True},
    {"name": "tts", "host": os.getenv("TTS_HOST", "127.0.0.1"),
     "port": int(os.getenv("TTS_PORT", "8100")),
     "listen": int(os.getenv("LISTEN_TTS", "3016")), "health": "/v1/models", "wol": False},
]

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

wol_lock = threading.Lock()
wol_state = {"last_sent": 0.0}
state_lock = threading.Lock()

METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]
EXCLUDED_HEADERS = {"content-encoding", "content-length", "transfer-encoding", "connection"}
ALIAS_PATHS = ("v1/chat/completions", "v1/completions", "v1/models", "v1/embeddings")


def load_active():
    try:
        with open(STATE_FILE) as f:
            data = json.load(f)
        a = data.get("active", DEFAULT_ACTIVE)
        if any(i["real"] == a for i in LLM_INSTANCES):
            return a
    except Exception:
        pass
    return DEFAULT_ACTIVE


def save_active(alias):
    if not any(i["real"] == alias for i in LLM_INSTANCES):
        raise ValueError("instance inconnue: %s" % alias)
    os.makedirs(os.path.dirname(STATE_FILE), exist_ok=True)
    tmp = STATE_FILE + ".tmp"
    with open(tmp, "w") as f:
        json.dump({"active": alias}, f)
    os.replace(tmp, STATE_FILE)


def active_instance():
    a = load_active()
    for i in LLM_INSTANCES:
        if i["real"] == a:
            return i
    return LLM_INSTANCES[0]


def send_wol():
    mac = PC_MAC.replace(":", "").replace("-", "")
    if len(mac) != 12:
        raise ValueError("MAC invalide: %s" % PC_MAC)
    data = bytes.fromhex("FF" * 6 + mac * 16)
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        for i in range(3):
            s.sendto(data, (BROADCAST_IP, 9))
            time.sleep(0.5)
    log.info(f"[WoL] 3 paquets magic envoyes vers {PC_MAC} via {BROADCAST_IP}:9")


def service_alive(route) -> bool:
    if route.get("port") is None:
        return False
    host = route.get("host", PC_HOST)
    try:
        r = requests.get(f"http://{host}:{route['port']}{route['health']}", timeout=3)
        return r.ok
    except Exception:
        return False


def pc_alive() -> bool:
    return any(service_alive(r) for r in ROUTES if r.get("wol"))


def wait_for(route, timeout: int) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        if service_alive(route):
            return True
        log.info(f"[{route['name']}] en attente du service... ({int(deadline - time.time())}s restantes)")
        time.sleep(CHECK_INTERVAL)
    return service_alive(route)


def ensure_up(route):
    name = route["name"]
    if service_alive(route):
        return
    if not route["wol"]:
        raise RuntimeError(f"service {name} inatteignable ({route.get('host', PC_HOST)}:{route['port']})")
    with wol_lock:
        if pc_alive():
            fast = int(os.getenv("FAST_FAIL_TIMEOUT", "20"))
            log.info(f"[WoL] PC deja UP, attente du service {name} (fail fast {fast}s)")
            if not wait_for(route, fast):
                raise RuntimeError(f"service {name} ne repond pas apres {fast}s (instance arretee ?)")
            return
        now = time.time()
        if now - wol_state["last_sent"] > WAKE_TIMEOUT:
            log.info(f"[WoL] PC DOWN -> envoi WOL vers {PC_MAC}")
            send_wol()
            wol_state["last_sent"] = now
        else:
            log.info(f"[WoL] reveil deja en cours, attente sans re-envoi")
        deadline = time.time() + WAKE_TIMEOUT
        while time.time() < deadline:
            if pc_alive():
                log.info(f"[WoL] PC UP apres reveil")
                break
            log.info(f"[WoL] en attente du PC... ({int(deadline - time.time())}s restantes)")
            time.sleep(CHECK_INTERVAL)
        else:
            raise RuntimeError(f"PC ne repond pas apres {WAKE_TIMEOUT}s")
    if not wait_for(route, SERVICE_TIMEOUT):
        raise RuntimeError(f"service {route['name']} ne repond pas apres le reveil du PC")


def rewrite_request_body(body: bytes, real: str) -> bytes:
    try:
        obj = json.loads(body)
    except Exception:
        return body
    if isinstance(obj, dict) and obj.get("model") == MODEL_ALIAS:
        obj["model"] = real
        return json.dumps(obj, ensure_ascii=False).encode()
    return body


def rewrite_json(obj, real: str) -> bool:
    changed = False
    if isinstance(obj, dict):
        if obj.get("model") == real:
            obj["model"] = MODEL_ALIAS
            changed = True
        for key in ("data", "models"):
            items = obj.get(key)
            if not isinstance(items, list):
                continue
            for m in items:
                if not isinstance(m, dict):
                    continue
                if m.get("id") == real:
                    m["id"] = MODEL_ALIAS
                    changed = True
                for field in ("name", "model"):
                    if m.get(field) == real:
                        m[field] = MODEL_ALIAS
                        changed = True
                aliases = m.get("aliases")
                if isinstance(aliases, list) and real in aliases:
                    m["aliases"] = [MODEL_ALIAS if a == real else a for a in aliases]
                    changed = True
    return changed


def rewrite_sse_line(line: bytes, real: str) -> bytes:
    if line.startswith(b"data:"):
        payload = line[5:].lstrip(b" ")
        if payload and payload != b"[DONE]":
            try:
                obj = json.loads(payload)
                if isinstance(obj, dict) and obj.get("model") == real:
                    obj["model"] = MODEL_ALIAS
                    return b"data: " + json.dumps(obj, ensure_ascii=False).encode()
            except Exception:
                pass
    return line


def sse_response(resp, real: str):
    def generate():
        buf = b""
        for chunk in resp.iter_content(chunk_size=4096):
            if not chunk:
                continue
            buf += chunk
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                yield rewrite_sse_line(line, real) + b"\n"
        if buf:
            yield rewrite_sse_line(buf, real) + b"\n"

    headers = [(n, v) for n, v in resp.headers.items() if n.lower() not in EXCLUDED_HEADERS]
    return Response(generate(), status=resp.status_code, headers=headers,
                    content_type="text/event-stream; charset=utf-8")


def proxy_request(route, path: str):
    real = None
    alias_path = None
    if route.get("alias"):
        inst = active_instance()
        route = dict(route)
        route["port"] = inst["port"]
        real = inst["real"]
        norm = path.lstrip("/")
        if any(norm == p or norm.startswith(p + "/") for p in ALIAS_PATHS):
            alias_path = norm
    try:
        ensure_up(route)
    except Exception as e:
        return Response(str(e), status=503)
    host = route.get("host", PC_HOST)
    base = f"http://{host}:{route['port']}"
    upstream_url = urljoin(base + "/", path)
    headers = dict(request.headers)
    headers.pop("Host", None)
    body = request.get_data()
    want_stream = False
    if alias_path and request.method == "POST" and body:
        try:
            jd = json.loads(body)
            want_stream = bool(jd.get("stream"))
        except Exception:
            pass
        body = rewrite_request_body(body, real)
    try:
        resp = requests.request(
            method=request.method,
            url=upstream_url,
            headers=headers,
            data=body,
            cookies=request.cookies,
            params=request.args,
            allow_redirects=False,
            stream=True,
            timeout=None,
        )
    except Exception as e:
        log.warning(f"[{route['name']}] erreur upstream: {e}")
        return Response(str(e), status=502)

    response_headers = [(n, v) for n, v in resp.headers.items() if n.lower() not in EXCLUDED_HEADERS]

    if alias_path and resp.status_code == 200:
        ctype = resp.headers.get("content-type", "")
        if "application/json" in ctype:
            data = resp.content
            try:
                obj = json.loads(data)
                if rewrite_json(obj, real):
                    data = json.dumps(obj, ensure_ascii=False).encode()
            except Exception:
                pass
            return Response(data, status=resp.status_code, headers=response_headers,
                            content_type="application/json")
        if want_stream and "text/event-stream" in ctype:
            return sse_response(resp, real)

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


def make_proxy_app(route):
    app = Flask(route["name"])

    @app.route("/", defaults={"path": ""}, methods=METHODS)
    @app.route("/<path:path>", methods=METHODS)
    def handle(path):
        return proxy_request(route, path)

    return app


def instance_up(inst) -> bool:
    try:
        r = requests.get(f"http://{PC_HOST}:{inst['port']}/health", timeout=3)
        return r.ok
    except Exception:
        return False


def instance_info(inst):
    info = {"port": inst["port"], "real": inst["real"], "label": inst["label"],
            "up": instance_up(inst), "model": None}
    if info["up"]:
        try:
            r = requests.get(f"http://{PC_HOST}:{inst['port']}/v1/models", timeout=5)
            r.raise_for_status()
            m = (r.json().get("data") or [{}])[0]
            d = m.get("details", {}) or {}
            meta = m.get("meta", {}) or {}
            info["model"] = {"id": m.get("id"),
                             "n_ctx": d.get("n_ctx") or meta.get("n_ctx"),
                             "size": d.get("size") or meta.get("size"),
                             "ftype": d.get("ftype") or meta.get("ftype")}
        except Exception:
            pass
    return info


INDEX_HTML = """<!doctype html>
<html lang="fr">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>llm-proxy — instances llama.cpp</title>
<style>
body{font-family:system-ui,sans-serif;background:#0f172a;color:#e2e8f0;margin:0;padding:24px}
h1{font-size:20px;margin:0 0 4px}
.small{font-size:12px;color:#94a3b8}
.card{background:#1e293b;border:1px solid #334155;border-radius:10px;padding:16px}
.card.active{border-color:#22c55e;box-shadow:0 0 0 1px #22c55e}
.row{display:flex;gap:16px;flex-wrap:wrap;margin-top:12px}
.card{flex:1;min-width:260px}
.tag{font-size:12px;padding:2px 8px;border-radius:99px;background:#334155;display:inline-block}
.tag.up{background:#14532d;color:#86efac}
.tag.down{background:#7f1d1d;color:#fca5a5}
button{background:#2563eb;border:0;color:#fff;padding:8px 14px;border-radius:8px;cursor:pointer;font-size:13px;margin-top:10px}
button:hover{background:#1d4ed8}
button.warn{background:#7c2d12}
button.warn:hover{background:#9a3412}
.model{font-size:13px;color:#cbd5e1;margin-top:8px;min-height:16px}
.warnbox{background:#451a03;border:1px solid #92400e;border-radius:8px;padding:10px 14px;margin-top:16px;font-size:13px}
.result{margin-top:8px;font-size:12px;white-space:pre-wrap;background:#0f172a;border-radius:6px;padding:8px;display:none}
</style>
</head>
<body>
<h1>llm-proxy — switch d'instances llama.cpp</h1>
<div class="small">Route :3012 &rarr; instance active &middot; alias mod&egrave;le <b>qwen3.8-27b</b> &middot; <span id="upd"></span></div>
<div class="row" id="instances"></div>
<div class="warnbox">&#9888;&#65039; Une seule instance &agrave; la fois : chaque instance occupe ~toute la VRAM (32 Go).
D&eacute;marrer l'instance choisie sur le PC (et arr&ecirc;ter l'autre) avant de l'actiquer ici.</div>
<script>
async function refresh(){
  try{
    const s = await (await fetch('/api/status')).json();
    const box = document.getElementById('instances');
    box.innerHTML = '';
    s.instances.forEach(i=>{
      const active = i.real === s.active;
      const m = i.model || {};
      const d = document.createElement('div');
      d.className = 'card' + (active ? ' active' : '');
      const ctx = m.n_ctx ? 'ctx ' + m.n_ctx : '';
      const size = m.size ? Math.round(m.size/1e9) + ' GB' : '';
      const detail = i.up ? (m.id + ' &middot; ' + [ctx,size].filter(Boolean).join(' &middot; ')) : 'ne repond pas';
      d.innerHTML =
        '<div><b>' + i.label + '</b> <span class="tag">:' + i.port + '</span> ' +
        '<span class="tag ' + (i.up ? 'up' : 'down') + '">' + (i.up ? 'UP' : 'DOWN') + '</span> ' +
        (active ? '<span class="tag up">ACTIVE</span>' : '') + '</div>' +
        '<div class="model">' + detail + '</div>' +
        '<div>' + (active ? '' : '<button onclick="activate(\\'' + i.real + '\\')">Activer</button>') +
        '<button class="warn" onclick="doTest(\\'' + i.real + '\\',' + i.port + ')">Tester</button></div>' +
        '<div class="result" id="res-' + i.port + '"></div>';
      box.appendChild(d);
    });
    document.getElementById('upd').textContent = 'maj ' + new Date().toLocaleTimeString();
  }catch(e){
    document.getElementById('upd').textContent = 'erreur: ' + e.message;
  }
}
async function activate(real){
  await fetch('/api/activate',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({real})});
  refresh();
}
async function doTest(real, port){
  const el = document.getElementById('res-' + port);
  el.style.display = 'block';
  el.textContent = 'test en cours...';
  try{
    const r = await (await fetch('/api/test',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({real})})).json();
    el.textContent = r.ok ? ('OK en ' + r.latency_ms + ' ms \\u2192 ' + r.content) : ('ERREUR: ' + (r.error || '?'));
  }catch(e){
    el.textContent = 'ERREUR: ' + e.message;
  }
}
setInterval(refresh, 3000);
refresh();
</script>
</body>
</html>
"""


def make_ui_app():
    app = Flask("ui")

    @app.route("/")
    def index():
        return Response(INDEX_HTML, content_type="text/html; charset=utf-8")

    @app.route("/api/status")
    def status():
        return jsonify({"active": active_instance()["real"],
                        "instances": [instance_info(i) for i in LLM_INSTANCES]})

    @app.route("/api/activate", methods=["POST"])
    def activate():
        data = request.get_json(silent=True) or {}
        real = data.get("real")
        if not any(i["real"] == real for i in LLM_INSTANCES):
            return jsonify({"ok": False, "error": "instance inconnue"}), 400
        with state_lock:
            save_active(real)
            log.info(f"[UI] instance active changee -> {real}")
        return jsonify({"ok": True, "active": real})

    @app.route("/api/test", methods=["POST"])
    def test_chat():
        data = request.get_json(silent=True) or {}
        real = data.get("real")
        inst = next((i for i in LLM_INSTANCES if i["real"] == real), None)
        if not inst:
            return jsonify({"ok": False, "error": "instance inconnue"}), 400
        t0 = time.time()
        try:
            r = requests.post(
                f"http://{PC_HOST}:{inst['port']}/v1/chat/completions",
                json={"model": real,
                      "messages": [{"role": "user", "content": "Reponds uniquement : OK"}],
                      "max_tokens": 16, "stream": False,
                      "chat_template_kwargs": {"enable_thinking": False}},
                timeout=SERVICE_TIMEOUT)
            latency = int((time.time() - t0) * 1000)
            if not r.ok:
                return jsonify({"ok": False, "latency_ms": latency,
                                "error": f"HTTP {r.status_code}: {r.text[:200]}"})
            obj = r.json()
            content = obj.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
            return jsonify({"ok": True, "latency_ms": latency, "content": content[:200]})
        except Exception as e:
            return jsonify({"ok": False, "error": str(e)})

    @app.route("/api/wake", methods=["POST"])
    def wake():
        with wol_lock:
            if pc_alive():
                return jsonify({"ok": True, "message": "PC deja UP"})
            now = time.time()
            if now - wol_state["last_sent"] > WAKE_TIMEOUT:
                send_wol()
                wol_state["last_sent"] = now
                return jsonify({"ok": True, "message": "WOL envoye"})
            return jsonify({"ok": True, "message": "reveil deja en cours"})

    return app


def serve(app, port):
    app.run(host="0.0.0.0", port=port, threaded=True)


if __name__ == "__main__":
    for r in ROUTES:
        target = f"{r.get('host', PC_HOST)}:{r['port']}" if r.get("port") is not None \
            else f"instance active (etat: {STATE_FILE})"
        log.info(f"[{r['name']}] :{r['listen']} -> {target} (wol={r['wol']})")
        threading.Thread(target=serve, args=(make_proxy_app(r), r["listen"]), daemon=True).start()
    log.info(f"[ui] :{LISTEN_UI} -> interface de gestion")
    threading.Thread(target=serve, args=(make_ui_app(), LISTEN_UI), daemon=True).start()
    while True:
        time.sleep(60)