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

import requests
from flask import Flask, Response, request

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

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": int(os.getenv("LLAMACPP_PORT", "8080")),
     "listen": int(os.getenv("LISTEN_LLAMACPP", "3012")), "health": "/health", "wol": True},
    {"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}

METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]


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:
    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["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():
            log.info(f"[WoL] PC deja UP, attente du service {name}")
            if not wait_for(route, SERVICE_TIMEOUT):
                raise RuntimeError(f"service {name} ne repond pas apres {SERVICE_TIMEOUT}s")
            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 {name} ne repond pas apres le reveil du PC")


def proxy_request(route, path: str):
    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)
    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:
        log.warning(f"[{route['name']}] erreur upstream: {e}")
        return Response(str(e), status=502)

    excluded = {"content-encoding", "content-length", "transfer-encoding", "connection"}
    response_headers = [
        (n, v) for n, v in resp.headers.items() if n.lower() not in excluded
    ]
    return Response(
        resp.raw,
        status=resp.status_code,
        headers=response_headers,
        direct_passthrough=True,
    )


def make_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 serve(route):
    make_app(route).run(host="0.0.0.0", port=route["listen"], threaded=True)


if __name__ == "__main__":
    for r in ROUTES:
        log.info(f"[{r['name']}] :{r['listen']} -> {r.get('host', PC_HOST)}:{r['port']} (wol={r['wol']})")
        threading.Thread(target=serve, args=(r,), daemon=True).start()
    while True:
        time.sleep(60)