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

import requests
from flask import Flask, request, Response

TARGET_HOST = os.getenv("TARGET_HOST", "192.168.1.10")
TARGET_PORT = int(os.getenv("TARGET_PORT", "11434"))
TARGET_MAC = os.getenv("TARGET_MAC", "02:76:C6:01:48:C8")
LISTEN_PORT = int(os.getenv("LISTEN_PORT", "3010"))
WAKE_TIMEOUT = int(os.getenv("WAKE_TIMEOUT", "60"))
CHECK_INTERVAL = int(os.getenv("CHECK_INTERVAL", "5"))
BROADCAST_IP = os.getenv("BROADCAST_IP", "192.168.1.255")  # ← broadcast dirigé

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

app = Flask(__name__)

def send_wol(mac_address: str):
    mac = mac_address.replace(":", "").replace("-", "")
    if len(mac) != 12:
        raise ValueError("MAC invalide: %s" % mac_address)
    data = bytes.fromhex("FF" * 6 + mac * 16)
    # Envoyer 3 paquets pour fiabilité
    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))
            log.info(f"[WoL] Paquet magic #{i+1} envoyé → MAC={mac_address} via {BROADCAST_IP}:9")

def ollama_alive() -> bool:
    try:
        url = f"http://{TARGET_HOST}:{TARGET_PORT}/api/version"
        r = requests.get(url, timeout=2)
        return r.ok
    except Exception:
        return False

def wait_for_ollama():
    if ollama_alive():
        log.info("[WoL] Ollama déjà UP, skip WoL")
        return
    log.info(f"[WoL] Ollama DOWN → envoi WoL vers {TARGET_MAC}")
    send_wol(TARGET_MAC)
    deadline = time.time() + WAKE_TIMEOUT
    while time.time() < deadline:
        remaining = int(deadline - time.time())
        if ollama_alive():
            log.info(f"[WoL] Ollama UP après réveil (reste {remaining}s)")
            return
        log.info(f"[WoL] En attente d'Ollama... ({remaining}s restantes)")
        time.sleep(CHECK_INTERVAL)
    raise RuntimeError("Ollama ne répond pas après WoL")

def proxy_request(path: str):
    wait_for_ollama()
    base = f"http://{TARGET_HOST}:{TARGET_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,
            allow_redirects=False,
            stream=True,
            timeout=60,
        )
    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,
    )

@app.route("/", defaults={"path": ""}, methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
def handle(path):
    return proxy_request(path)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=LISTEN_PORT)