"""GPU model manager — ensures only one model subprocess runs at a time."""

from __future__ import annotations

import asyncio
import logging
import os
import subprocess
import time
from pathlib import Path
from typing import Optional

import httpx

ROOT_DIR = Path(__file__).parent.parent


def _load_dotenv(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    result: dict[str, str] = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        key = key.strip()
        value = value.strip().strip("\"'")
        if key and value:
            result[key] = value
    return result

ACESTEP_PORT = 8001
FOUNDATION1_PORT = 8002
HEALTH_POLL_INTERVAL = 2.0
HEALTH_TIMEOUT_SECONDS = 180  # model loading can be slow on first run

logger = logging.getLogger(__name__)


def _python_exe(venv_dir: Path) -> Path:
    """Return the Python executable for a given venv directory."""
    win = venv_dir / "Scripts" / "python.exe"
    if win.exists():
        return win
    return venv_dir / "bin" / "python"


class ModelNotActiveError(Exception):
    def __init__(self, requested: str, active: Optional[str]) -> None:
        self.requested = requested
        self.active = active
        super().__init__(f"Model '{requested}' is not loaded (active: {active!r})")


class ModelManager:
    def __init__(self) -> None:
        self.active_model: Optional[str] = None
        self._lock = asyncio.Lock()
        self._process: Optional[subprocess.Popen] = None
        self.acestep_port: int = ACESTEP_PORT
        self.foundation1_port: int = FOUNDATION1_PORT

    async def load(self, model: str) -> None:
        """Load a model, unloading the currently active one first if needed."""
        async with self._lock:
            if self.active_model == model:
                logger.info("Model '%s' is already active.", model)
                return
            if self.active_model is not None:
                logger.info("Unloading model '%s'...", self.active_model)
                await asyncio.get_event_loop().run_in_executor(None, self._stop_process)
            logger.info("Starting model '%s'...", model)
            await asyncio.get_event_loop().run_in_executor(None, self._start_process, model)
            self.active_model = model
            logger.info("Model '%s' is ready.", model)

    async def unload(self) -> None:
        """Unload the active model."""
        async with self._lock:
            if self.active_model is None:
                return
            logger.info("Unloading model '%s'...", self.active_model)
            await asyncio.get_event_loop().run_in_executor(None, self._stop_process)

    def _stop_process(self) -> None:
        if self._process is None:
            return
        try:
            self._process.terminate()
            self._process.wait(timeout=30)
        except subprocess.TimeoutExpired:
            logger.warning("Process did not terminate gracefully, killing.")
            self._process.kill()
            self._process.wait()
        except Exception as exc:
            logger.error("Error stopping process: %s", exc)
        finally:
            self._process = None
        self.active_model = None

    def _start_process(self, model: str) -> None:
        if model == "fullsong":
            self._process = self._launch_acestep()
            self._wait_healthy(f"http://127.0.0.1:{self.acestep_port}/health")
        elif model == "clip":
            self._process = self._launch_foundation1()
            self._wait_healthy(f"http://127.0.0.1:{self.foundation1_port}/health")
        elif model == "separator":
            pass  # no persistent process; previous model already stopped by load(), VRAM is free
        else:
            raise ValueError(f"Unknown model: '{model}'. Must be 'fullsong', 'clip', or 'separator'.")

    def _launch_acestep(self) -> subprocess.Popen:
        venv_python = _python_exe(ROOT_DIR / "ace-step" / ".venv")
        cmd = [str(venv_python), "-c", "from acestep.api_server import main; main()"]
        dotenv_vars = _load_dotenv(ROOT_DIR / ".env")
        env = {**os.environ, **dotenv_vars, "ACESTEP_API_PORT": str(self.acestep_port)}
        if dotenv_vars:
            logger.info("ACE-Step env vars from .env:")
            for k, v in dotenv_vars.items():
                logger.info("  %s=%s", k, v)
        else:
            logger.info("No .env file found; no extra env vars applied to ACE-Step.")
        logger.info("Launching ACE-Step on port %d: %s", self.acestep_port, " ".join(cmd))
        return subprocess.Popen(cmd, cwd=str(ROOT_DIR / "ace-step"), env=env)

    def _launch_foundation1(self) -> subprocess.Popen:
        venv_python = _python_exe(ROOT_DIR / "foundation1" / ".venv")
        server_script = str(ROOT_DIR / "foundation1_server" / "server.py")
        cmd = [str(venv_python), server_script]
        env = {**os.environ, "FOUNDATION1_SERVER_PORT": str(self.foundation1_port)}
        logger.info("Launching Foundation-1 server on port %d: %s", self.foundation1_port, " ".join(cmd))
        return subprocess.Popen(cmd, cwd=str(ROOT_DIR), env=env)

    def _wait_healthy(self, url: str) -> None:
        deadline = time.monotonic() + HEALTH_TIMEOUT_SECONDS
        logger.info("Waiting for service at %s ...", url)
        while time.monotonic() < deadline:
            try:
                with httpx.Client(timeout=5.0) as client:
                    resp = client.get(url)
                    if resp.status_code == 200:
                        logger.info("Service healthy: %s", url)
                        return
            except Exception:
                pass
            time.sleep(HEALTH_POLL_INTERVAL)
        raise TimeoutError(
            f"Service at {url} did not become healthy within {HEALTH_TIMEOUT_SECONDS}s"
        )


model_manager = ModelManager()