#!/usr/bin/env python3
"""
AI Studio GPU Worker.

Ce worker expose un serveur HTTP local. Le serveur Hugging Face Space l'appelle
via l'URL ngrok enregistrée dans l'interface. Les requêtes du Space vers ngrok
incluent le header `ngrok-skip-browser-warning: true`.

Installation minimale:
  pip install psutil pynvml

Pour l'entraînement Hugging Face:
  pip install torch transformers datasets peft accelerate evaluate scikit-learn

Lancement exemple:
  python worker.py --server-url https://ton-space.hf.space --user-token TON_TOKEN --worker-name "Colab T4" --port 8765
"""

import argparse
import json
import signal
import sys
import threading
import time
import traceback
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

try:
    import psutil
except Exception:
    psutil = None

try:
    import pynvml
except Exception:
    pynvml = None

STATE = {
    "server_url": "",
    "user_token": "",
    "worker_name": "",
    "ngrok_url": "",
    "current_job": None,
    "stop_event": None,
    "lock": threading.Lock(),
    "started_at": time.time()
}


def log(message):
    print(f"[worker] {message}", flush=True)


def post_json(path, payload, token=None, timeout=20):
    url = f"{STATE['server_url'].rstrip('/')}{path}"
    body = json.dumps(payload).encode("utf-8")
    headers = {
        "Content-Type": "application/json",
        "ngrok-skip-browser-warning": "true"
    }

    if token:
        headers["Authorization"] = f"Bearer {token}"

    request = Request(url, data=body, headers=headers, method="POST")

    try:
        with urlopen(request, timeout=timeout) as response:
            raw = response.read().decode("utf-8")
            return json.loads(raw) if raw else {}
    except HTTPError as error:
        raw = error.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"POST {path} -> HTTP {error.code}: {raw}") from error
    except URLError as error:
        raise RuntimeError(f"POST {path} -> {error}") from error


def send_job_log(job_id, level, message, progress=None):
    try:
        post_json(
            f"/api/jobs/{job_id}/logs",
            {
                "level": level,
                "message": message,
                "progress": progress
            },
            token=STATE["user_token"],
            timeout=15
        )
    except Exception as exc:
        log(f"Impossible d'envoyer le log: {exc}")


def send_job_status(job_id, status, progress=None, error=None):
    try:
        post_json(
            f"/api/jobs/{job_id}/status",
            {
                "status": status,
                "progress": progress,
                "error": error
            },
            token=STATE["user_token"],
            timeout=20
        )
    except Exception as exc:
        log(f"Impossible d'envoyer le statut: {exc}")


def decode_nvml_value(value):
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="ignore")
    return value


def get_gpu_metrics():
    metrics = {
        "gpu_count": 0,
        "gpu_name": None,
        "gpu_index": 0,
        "vram_total": 0,
        "vram_used": 0,
        "vram_free": 0,
        "gpu_util": 0,
        "gpu_temperature": None
    }

    if pynvml is not None:
        try:
            pynvml.nvmlInit()
            count = pynvml.nvmlDeviceGetCount()
            metrics["gpu_count"] = count

            if count > 0:
                handle = pynvml.nvmlDeviceGetHandleByIndex(0)
                name = decode_nvml_value(pynvml.nvmlDeviceGetName(handle))
                memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
                utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)

                metrics.update({
                    "gpu_name": name,
                    "gpu_index": 0,
                    "vram_total": int(memory.total),
                    "vram_used": int(memory.used),
                    "vram_free": int(memory.free),
                    "gpu_util": int(utilization.gpu),
                    "gpu_temperature": int(pynvml.nvmlDeviceGetTemperature(
                        handle,
                        pynvml.NVML_TEMPERATURE_GPU
                    ))
                })

            pynvml.nvmlShutdown()
            return metrics
        except Exception as exc:
            metrics["gpu_error"] = f"pynvml: {exc}"

    try:
        import subprocess

        command = [
            "nvidia-smi",
            "--query-gpu=name,memory.total,memory.used,memory.free,utilization.gpu,temperature.gpu",
            "--format=csv,noheader,nounits"
        ]
        result = subprocess.run(command, check=True, capture_output=True, text=True, timeout=10)
        lines = [line for line in result.stdout.splitlines() if line.strip()]

        if lines:
            parts = [part.strip() for part in lines[0].split(",")]
            if len(parts) >= 6:
                metrics.update({
                    "gpu_count": len(lines),
                    "gpu_name": parts[0],
                    "gpu_index": 0,
                    "vram_total": int(float(parts[1]) * 1024 * 1024),
                    "vram_used": int(float(parts[2]) * 1024 * 1024),
                    "vram_free": int(float(parts[3]) * 1024 * 1024),
                    "gpu_util": int(float(parts[4])),
                    "gpu_temperature": int(float(parts[5]))
                })
    except Exception as exc:
        if "gpu_error" not in metrics:
            metrics["gpu_error"] = f"nvidia-smi: {exc}"

    return metrics


def get_metrics():
    metrics = {
        "worker_url": STATE["ngrok_url"],
        "worker_name": STATE["worker_name"],
        "uptime_seconds": int(time.time() - STATE["started_at"]),
        "current_job": STATE["current_job"],
        "python_version": sys.version.split()[0]
    }

    if psutil is not None:
        try:
            cpu = psutil.cpu_percent(interval=0.2)
            memory = psutil.virtual_memory()
            disk = psutil.disk_usage("/")

            metrics.update({
                "cpu_percent": cpu,
                "cpu_count": psutil.cpu_count(logical=True),
                "ram_total": int(memory.total),
                "ram_used": int(memory.used),
                "ram_free": int(memory.free),
                "ram_percent": memory.percent,
                "disk_total": int(disk.total),
                "disk_used": int(disk.used),
                "disk_free": int(disk.free),
                "disk_percent": disk.percent
            })
        except Exception as exc:
            metrics["system_error"] = f"psutil: {exc}"
    else:
        metrics["system_error"] = "psutil non installé"

    metrics.update(get_gpu_metrics())
    return metrics


def register_worker():
    try:
        with STATE["lock"]:
            current_job = STATE["current_job"]

        post_json(
            "/api/workers/register",
            {
                "userToken": STATE["user_token"],
                "workerName": STATE["worker_name"],
                "status": "busy" if current_job else "online",
                "metrics": get_metrics()
            },
            timeout=15
        )
        log("Worker enregistré auprès du serveur.")
    except Exception as exc:
        log(f"Échec de l'enregistrement: {exc}")


def heartbeat_loop(interval):
    while True:
        time.sleep(interval)
        register_worker()


def run_job_thread(job, stop_event):
    job_id = job.get("jobId")
    log(f"Démarrage du job {job_id}")
    send_job_status(job_id, "running", progress=0)
    send_job_log(job_id, "info", "Démarrage de l'entraînement sur le worker GPU.", progress=0)

    try:
        from train_hf import run_training

        result = run_training(
            job=job,
            log_callback=lambda level, message, progress=None: send_job_log(job_id, level, message, progress),
            stop_event=stop_event
        )

        if stop_event.is_set():
            send_job_status(job_id, "cancelled", progress=100)
            send_job_log(job_id, "warning", "Job arrêté par l'utilisateur.", progress=100)
        else:
            send_job_status(job_id, "completed", progress=100)
            send_job_log(
                job_id,
                "success",
                f"Entraînement terminé. Sortie: {result.get('output_dir', 'inconnue')}",
                progress=100
            )
    except KeyboardInterrupt:
        send_job_status(job_id, "cancelled", progress=100)
        send_job_log(job_id, "warning", "Job annulé.", progress=100)
    except Exception as exc:
        log(traceback.format_exc())
        send_job_status(job_id, "failed", error=str(exc))
        send_job_log(job_id, "error", f"Erreur pendant l'entraînement: {exc}")
    finally:
        with STATE["lock"]:
            STATE["current_job"] = None
            STATE["stop_event"] = None
        log(f"Job {job_id} terminé.")


class WorkerHandler(BaseHTTPRequestHandler):
    server_version = "AIStudioWorker/1.0"

    def log_message(self, format, *args):
        log(f"HTTP {args[0] if args else ''}")

    def send_json(self, status, payload):
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization, ngrok-skip-browser-warning")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.end_headers()
        self.wfile.write(body)

    def read_json(self):
        length = int(self.headers.get("Content-Length", "0") or 0)
        if not length:
            return {}

        raw = self.rfile.read(length).decode("utf-8")
        return json.loads(raw) if raw else {}

    def do_OPTIONS(self):
        self.send_response(204)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization, ngrok-skip-browser-warning")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.end_headers()

    def do_GET(self):
        if self.path == "/health":
            with STATE["lock"]:
                current_job = STATE["current_job"]
                stop_requested = bool(STATE["stop_event"].is_set() if STATE["stop_event"] else False)

            self.send_json(200, {
                "ok": True,
                "status": "busy" if current_job else "idle",
                "current_job": current_job,
                "stop_requested": stop_requested,
                "ngrok_url": STATE["ngrok_url"],
                "worker_name": STATE["worker_name"],
                "metrics": get_metrics()
            })
            return

        if self.path == "/metrics":
            self.send_json(200, get_metrics())
            return

        self.send_json(404, {"error": "Endpoint inconnu."})

    def do_POST(self):
        try:
            payload = self.read_json()

            if self.path == "/start-job":
                with STATE["lock"]:
                    if STATE["current_job"]:
                        self.send_json(409, {
                            "ok": False,
                            "error": "Un job est déjà en cours sur ce worker."
                        })
                        return

                    STATE["current_job"] = payload.get("jobId")
                    STATE["stop_event"] = threading.Event()
                    stop_event = STATE["stop_event"]

                thread = threading.Thread(
                    target=run_job_thread,
                    args=(payload, stop_event),
                    daemon=True
                )
                thread.start()
                self.send_json(200, {"ok": True, "jobId": payload.get("jobId")})
                return

            if self.path == "/stop-job":
                job_id = payload.get("jobId")
                with STATE["lock"]:
                    current_job = STATE["current_job"]
                    if current_job and current_job == job_id and STATE["stop_event"]:
                        STATE["stop_event"].set()
                        self.send_json(200, {"ok": True, "jobId": job_id, "stop_requested": True})
                    else:
                        self.send_json(404, {"ok": False, "error": "Aucun job correspondant n'est en cours."})
                return

            self.send_json(404, {"error": "Endpoint inconnu."})
        except Exception as exc:
            self.send_json(500, {"error": str(exc)})


def parse_args():
    parser = argparse.ArgumentParser(description="AI Studio GPU Worker")
    parser.add_argument("--server-url", required=True, help="URL du serveur AI Studio, par exemple https://ton-space.hf.space")
    parser.add_argument("--user-token", required=True, help="Token utilisateur pour s'authentifier auprès du serveur")
    parser.add_argument("--worker-name", default="GPU Worker", help="Nom de la machine worker")
    parser.add_argument("--ngrok-url", default="", help="URL publique ngrok, utile pour l'affichage et le diagnostic")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8765)
    parser.add_argument("--heartbeat-interval", type=int, default=15)
    return parser.parse_args()


def main():
    args = parse_args()
    STATE["server_url"] = args.server_url.rstrip("/")
    STATE["user_token"] = args.user_token
    STATE["worker_name"] = args.worker_name
    STATE["ngrok_url"] = args.ngrok_url

    def shutdown(signum, frame):
        log("Signal reçu, arrêt du worker...")
        with STATE["lock"]:
            if STATE["stop_event"]:
                STATE["stop_event"].set()
        server.shutdown()

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    register_worker()

    heartbeat = threading.Thread(
        target=heartbeat_loop,
        args=(args.heartbeat_interval,),
        daemon=True
    )
    heartbeat.start()

    global server
    server = ThreadingHTTPServer((args.host, args.port), WorkerHandler)
    log(f"Worker HTTP démarré sur {args.host}:{args.port}")
    log(f"Worker name: {STATE['worker_name']}")
    log("Le serveur Space doit utiliser l'URL ngrok enregistrée dans l'interface.")
    server.serve_forever()


if __name__ == "__main__":
    main()