from __future__ import annotations

import json
import os
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest

from app.config import (
    NODE_INVOKE_TIMEOUT_SECONDS,
    NODE_READY_TIMEOUT_SECONDS,
    NODE_REAP_INTERVAL_SECONDS,
    SDK_SRC,
    WORKSPACE_ROOT,
)
from app.db import Database


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _idle_seconds(last_used_at: str) -> float:
    if not last_used_at:
        return 0.0
    try:
        last = datetime.fromisoformat(last_used_at)
        return max(0.0, (datetime.now(timezone.utc) - last).total_seconds())
    except ValueError:
        return 0.0


def _cuda_library_dirs(repo_dir: Path) -> list[Path]:
    venv = repo_dir / ".venv"
    site_packages: list[Path] = []
    windows_site = venv / "Lib" / "site-packages"
    if windows_site.is_dir():
        site_packages.append(windows_site)
    site_packages.extend((venv / "lib").glob("python3*/site-packages"))

    dirs: list[Path] = []
    for site in site_packages:
        for vendor in ("cublas", "cudnn"):
            for subdir in ("lib", "bin"):
                lib_dir = site / "nvidia" / vendor / subdir
                if lib_dir.is_dir():
                    dirs.append(lib_dir)
    return dirs


def _with_cuda_library_path(env: dict[str, str], repo_dir: Path) -> dict[str, str]:
    dirs = [str(path) for path in _cuda_library_dirs(repo_dir)]
    if not dirs:
        return env
    var = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH"
    existing = [item for item in env.get(var, "").split(os.pathsep) if item]
    additions = [path for path in dirs if path not in existing]
    if additions:
        env[var] = os.pathsep.join(additions + existing)
    return env


@dataclass
class NodeRuntime:
    instance_id: str
    node_id: str
    manifest: NodeManifest
    status: str = "starting"
    process: subprocess.Popen | None = None
    pid: int | None = None
    address: str | None = None
    started_at: str = field(default_factory=_now_iso)
    last_used_at: str = field(default_factory=_now_iso)
    busy_count: int = 0
    error: str | None = None
    stderr_tail: list[str] = field(default_factory=list)


class NodeManager:
    def __init__(self, db: Database) -> None:
        self.db = db
        self._runtimes: dict[str, NodeRuntime] = {}
        self._lock = threading.RLock()
        self._reaper_thread: threading.Thread | None = None
        self._stopping = False

    def start_reaper(self) -> None:
        self._stopping = False
        self._reaper_thread = threading.Thread(
            target=self._reaper_loop,
            name="wov-node-reaper",
            daemon=True,
        )
        self._reaper_thread.start()

    def shutdown(self) -> None:
        self._stopping = True
        with self._lock:
            for runtime in list(self._runtimes.values()):
                self._stop_locked(runtime)

    def _reaper_loop(self) -> None:
        while not self._stopping:
            time.sleep(NODE_REAP_INTERVAL_SECONDS)
            with self._lock:
                for runtime in list(self._runtimes.values()):
                    if runtime.status != "ready" or runtime.busy_count > 0:
                        continue
                    if runtime.manifest.keep_warm:
                        continue
                    ttl = runtime.manifest.idle_ttl_seconds
                    if ttl >= 0 and _idle_seconds(runtime.last_used_at) >= ttl:
                        runtime.status = "stopping"
                        self._stop_locked(runtime)

    def _resolve_command(self, manifest: NodeManifest) -> list[str]:
        command = list(manifest.command)
        if command and command[0].lower() in {"python", "python3"}:
            repo_dir = WORKSPACE_ROOT / manifest.repo_dir
            windows_python = repo_dir / ".venv" / "Scripts" / "python.exe"
            unix_python = repo_dir / ".venv" / "bin" / "python"
            if windows_python.is_file():
                command[0] = str(windows_python)
            elif unix_python.is_file():
                command[0] = str(unix_python)
            else:
                command[0] = sys.executable
        return command

    def _node_env(self, manifest: NodeManifest) -> dict[str, str]:
        env = os.environ.copy()
        current_pythonpath = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = os.pathsep.join(
            item for item in [str(SDK_SRC), current_pythonpath] if item
        )
        env.update(manifest.env)
        repo_dir = WORKSPACE_ROOT / manifest.repo_dir
        return _with_cuda_library_path(env, repo_dir)

    def _start_locked(self, manifest: NodeManifest) -> NodeRuntime:
        runtime = NodeRuntime(
            instance_id=f"ni_{uuid.uuid4().hex[:12]}",
            node_id=manifest.id,
            manifest=manifest,
        )
        repo_dir = WORKSPACE_ROOT / manifest.repo_dir
        if not repo_dir.is_dir():
            runtime.status = "error"
            runtime.error = f"repo_dir not found: {repo_dir}"
            self.db.upsert_instance(self._instance_row(runtime))
            return runtime

        ready_event = threading.Event()
        ready_port: dict[str, int] = {}

        try:
            process = subprocess.Popen(
                self._resolve_command(manifest),
                cwd=str(repo_dir),
                env=self._node_env(manifest),
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                encoding="utf-8",
                bufsize=1,
            )
        except Exception as exc:  # noqa: BLE001
            runtime.status = "error"
            runtime.error = str(exc)
            self.db.upsert_instance(self._instance_row(runtime))
            return runtime

        runtime.process = process
        runtime.pid = process.pid

        def read_stdout() -> None:
            assert process.stdout is not None
            for line in iter(process.stdout.readline, ""):
                line = line.strip()
                if line.startswith("WOV_NODE_READY"):
                    try:
                        ready_port["port"] = int(line.split("port=", 1)[1])
                    except (IndexError, ValueError):
                        ready_port["port"] = 0
                    ready_event.set()

        def read_stderr() -> None:
            assert process.stderr is not None
            for line in iter(process.stderr.readline, ""):
                line = line.strip()
                if line:
                    runtime.stderr_tail.append(line)
                    runtime.stderr_tail = runtime.stderr_tail[-200:]

        threading.Thread(target=read_stdout, daemon=True).start()
        threading.Thread(target=read_stderr, daemon=True).start()

        ready = ready_event.wait(timeout=NODE_READY_TIMEOUT_SECONDS)
        if not ready or "port" not in ready_port or not ready_port["port"]:
            self._stop_process(process)
            runtime.status = "error"
            runtime.error = "node did not report readiness"
            self.db.upsert_instance(self._instance_row(runtime))
            return runtime

        address = f"http://127.0.0.1:{ready_port['port']}"
        try:
            with urllib.request.urlopen(f"{address}/health", timeout=2) as response:
                if response.status != 200:
                    raise RuntimeError(f"health check returned {response.status}")
        except Exception as exc:  # noqa: BLE001
            self._stop_process(process)
            runtime.status = "error"
            runtime.error = f"health check failed: {exc}"
            self.db.upsert_instance(self._instance_row(runtime))
            return runtime

        runtime.address = address
        runtime.status = "ready"
        runtime.started_at = _now_iso()
        runtime.last_used_at = _now_iso()
        self._runtimes[runtime.instance_id] = runtime
        self.db.upsert_instance(self._instance_row(runtime))
        return runtime

    def _stop_process(self, process: subprocess.Popen) -> None:
        if process.poll() is not None:
            return
        process.terminate()
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=5)

    def _stop_locked(self, runtime: NodeRuntime) -> None:
        if runtime.process is not None:
            self._stop_process(runtime.process)
        runtime.status = "stopped"
        runtime.pid = None
        runtime.address = None
        self._runtimes.pop(runtime.instance_id, None)
        self.db.upsert_instance(self._instance_row(runtime))

    def _instance_row(self, runtime: NodeRuntime) -> dict:
        return {
            "id": runtime.instance_id,
            "node_id": runtime.node_id,
            "status": runtime.status,
            "pid": runtime.pid,
            "address": runtime.address,
            "started_at": runtime.started_at,
            "last_used_at": runtime.last_used_at,
            "busy_since": None,
            "error": runtime.error,
        }

    def acquire(self, node_id: str) -> tuple[NodeRuntime, str]:
        manifest = self.db.get_node(node_id)
        if manifest is None:
            raise ValueError(f"node not registered: {node_id}")

        with self._lock:
            for runtime in self._runtimes.values():
                if (
                    runtime.node_id == node_id
                    and runtime.status == "ready"
                    and runtime.busy_count < manifest.max_concurrency
                ):
                    runtime.busy_count += 1
                    runtime.last_used_at = _now_iso()
                    return runtime, runtime.address or ""

            runtime = self._start_locked(manifest)
            if runtime.status != "ready":
                raise RuntimeError(runtime.error or "node failed to start")
            runtime.busy_count += 1
            runtime.last_used_at = _now_iso()
            return runtime, runtime.address or ""

    def release(self, instance_id: str) -> None:
        with self._lock:
            runtime = self._runtimes.get(instance_id)
            if runtime is None:
                return
            runtime.busy_count = max(0, runtime.busy_count - 1)
            runtime.last_used_at = _now_iso()
            self.db.upsert_instance(self._instance_row(runtime))

    def invoke(self, node_id: str, request: InvokeRequest) -> InvokeResponse:
        runtime, address = self.acquire(node_id)
        request.node_instance_id = runtime.instance_id
        body = json.dumps(request.to_dict(), ensure_ascii=False).encode("utf-8")
        http_request = urllib.request.Request(
            f"{address}/invoke",
            data=body,
            headers={"Content-Type": "application/json; charset=utf-8"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(http_request, timeout=NODE_INVOKE_TIMEOUT_SECONDS) as response:
                payload = json.loads(response.read().decode("utf-8"))
            return InvokeResponse.from_dict(payload)
        except urllib.error.HTTPError as exc:
            try:
                payload = json.loads(exc.read().decode("utf-8"))
                return InvokeResponse.from_dict(payload)
            except Exception:  # noqa: BLE001
                return InvokeResponse(status="failed", error=f"node returned {exc.code}")
        except Exception as exc:  # noqa: BLE001
            return InvokeResponse(status="failed", error=str(exc))
        finally:
            self.release(request.node_instance_id)

    def stop_instance(self, instance_id: str) -> None:
        with self._lock:
            runtime = self._runtimes.get(instance_id)
            if runtime is None:
                return
            self._stop_locked(runtime)

    def stop_all_for_node(self, node_id: str) -> None:
        with self._lock:
            for runtime in list(self._runtimes.values()):
                if runtime.node_id == node_id:
                    self._stop_locked(runtime)