"""llama-server 进程管理 & 模型下载。

职责：
- 启动/停止本地 llama-server 子进程（Windows 兼容）
- 探测已有 .gguf 模型文件
- 从 HuggingFace URL 下载模型（支持 huggingface.co 路径别名）
- 下载进度可通过 SSE /llama/download/stream 订阅

用法：
  from gateway.llama_manager import get_llama_manager
  lm = get_llama_manager()
  await lm.start(model="models/qwen3.5-4b-q4_k_m.gguf")
  await lm.stop()
"""
from __future__ import annotations

import asyncio
import json
import os
import re
import subprocess
import sys
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, AsyncGenerator, Optional

# ─────────────────────────────────────────────────────────────────────────────
# 路径配置（与 config.yaml runtime.llama_server 段保持一致）
# ─────────────────────────────────────────────────────────────────────────────
_ROOT = Path(__file__).resolve().parent.parent          # E:\projectAIpopular
BIN_DIR = _ROOT / "bin"
MODELS_DIR = _ROOT / "models"
PID_FILE = _ROOT / "data" / "llama-server.pid"
LOG_FILE = _ROOT / "data" / "llama-server.log"

# 确保目录存在
BIN_DIR.mkdir(parents=True, exist_ok=True)
MODELS_DIR.mkdir(parents=True, exist_ok=True)
PID_FILE.parent.mkdir(parents=True, exist_ok=True)


# ─────────────────────────────────────────────────────────────────────────────
# 数据模型
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class LlamaServerStatus:
    running: bool
    pid: Optional[int] = None
    model: Optional[str] = None
    port: Optional[int] = None
    base_url: Optional[str] = None
    started_at: Optional[float] = None
    error: Optional[str] = None


@dataclass
class DownloadProgress:
    url: str
    dest: str
    total_bytes: Optional[int] = None
    downloaded_bytes: int = 0
    progress_pct: float = 0.0
    speed: str = ""
    eta: str = ""
    done: bool = False
    error: Optional[str] = None


# ─────────────────────────────────────────────────────────────────────────────
# llama_manager 单例
# ─────────────────────────────────────────────────────────────────────────────

class LlamaManager:
    _instance: Optional["LlamaManager"] = None

    def __init__(self):
        self._proc: Optional[subprocess.Popen] = None
        self._pid: Optional[int] = None
        self._model: Optional[str] = None
        self._port: Optional[int] = None
        self._started_at: Optional[float] = None
        self._downloading: dict[str, DownloadProgress] = {}   # url -> progress
        self._dl_lock = threading.Lock()
        # 加载已有进程
        self._load_pid()

    # ── 进程持久化 ─────────────────────────────────────────────────────────

    def _load_pid(self) -> None:
        """从 pid 文件恢复进程引用（进程仍在运行时）。"""
        if not PID_FILE.exists():
            return
        try:
            pid = int(PID_FILE.read_text().strip())
            os.kill(pid, 0)   # 检查进程是否存活
            # 进程还在，尝试接管（通过 cmdline 判断是否是 llama-server）
            self._pid = pid
            self._proc = self._attach_to_process(pid)
        except (ValueError, FileNotFoundError, OSError):
            PID_FILE.unlink(missing_ok=True)

    def _attach_to_process(self, pid: int) -> Optional[subprocess.Popen]:
        """通过 pid 重新关联到 Popen（仅作状态恢复，不拥有 stdout）。"""
        try:
            return subprocess.Popen(
                [sys.executable, "-c",
                 f"import os; os.kill({pid}, 0)"],  # 存活检查
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except Exception:
            return None

    def _write_pid(self, pid: int) -> None:
        PID_FILE.write_text(str(pid), encoding="utf-8")

    def _clear_pid(self) -> None:
        PID_FILE.unlink(missing_ok=True)

    # ── 进程管理 ─────────────────────────────────────────────────────────

    def find_binary(self) -> Optional[Path]:
        """查找 llama-server 可执行文件。"""
        candidates = [
            BIN_DIR / "llama-server.exe",
            BIN_DIR / "llama-server",
            _ROOT / "llama-server.exe",
            _ROOT / "llama-server",
        ]
        for p in candidates:
            if p.exists():
                return p
        # PATH 中查找
        import shutil
        found = shutil.which("llama-server") or shutil.which("llama-server.exe")
        if found:
            return Path(found)
        return None

    def status(self) -> LlamaServerStatus:
        """返回当前服务状态。"""
        if self._proc is None or self._pid is None:
            return LlamaServerStatus(running=False)
        try:
            # 检查进程是否存活
            os.kill(self._pid, 0)
        except OSError:
            # 进程已死
            self._proc = None
            self._pid = None
            self._model = None
            self._port = None
            self._started_at = None
            self._clear_pid()
            return LlamaServerStatus(running=False)
        return LlamaServerStatus(
            running=True,
            pid=self._pid,
            model=self._model,
            port=self._port,
            base_url=f"http://127.0.0.1:{self._port}/v1",
            started_at=self._started_at,
        )

    async def start(
        self,
        model: str,
        port: int = 8901,
        ngl: int = 99,
        ctx: int = 4096,
        extra_args: Optional[list] = None,
    ) -> LlamaServerStatus:
        """启动 llama-server，阻塞直到监听就绪或超时。"""
        if self.status().running:
            s = self.status()
            if s.model == model and s.port == port:
                return s   # 已是同一模型，无需重启
            await self.stop()

        binary = self.find_binary()
        if binary is None:
            return LlamaServerStatus(
                running=False,
                error="未找到 llama-server 可执行文件。"
                        "请将 llama-server.exe 放入 bin/ 目录，"
                        "或从 https://github.com/ggerganov/llama.cpp/releases 下载。",
            )

        model_path = Path(model)
        if not model_path.is_absolute():
            model_path = MODELS_DIR / model
        if not model_path.exists():
            return LlamaServerStatus(
                running=False,
                error=f"模型文件不存在：{model_path}。"
                        "请先下载模型，或在设置页填写 HuggingFace URL 下载。",
            )

        args = [
            str(binary),
            "-m", str(model_path),
            "-c", str(ctx),
            "-ngl", str(ngl),
            "--port", str(port),
            "--host", "127.0.0.1",
        ]
        if extra_args:
            args.extend(extra_args)

        LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
        log_f = open(LOG_FILE, "w", encoding="utf-8", buffering=1)

        try:
            self._proc = subprocess.Popen(
                args,
                stdout=log_f,
                stderr=subprocess.STDOUT,
                cwd=str(_ROOT),
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0,
            )
        except OSError as e:
            log_f.close()
            return LlamaServerStatus(running=False, error=f"启动失败：{e}")

        self._pid = self._proc.pid
        self._model = str(model_path)
        self._port = port
        self._started_at = time.time()
        self._write_pid(self._pid)

        # 等待服务就绪
        ok = await self._wait_until_ready(port, timeout=30)
        if not ok:
            await self.stop()
            return LlamaServerStatus(
                running=False,
                error=f"llama-server 启动后 {port} 端口在 30 秒内未响应",
            )

        return self.status()

    async def stop(self) -> None:
        """优雅停止 llama-server。"""
        if self._pid is None:
            self._proc = None
            return

        try:
            if sys.platform == "win32":
                # Windows: CTRL_BREAK_EVENT 或 taskkill
                subprocess.run(
                    ["taskkill", "/PID", str(self._pid), "/T", "/F"],
                    capture_output=True,
                    timeout=10,
                )
            else:
                os.kill(self._pid, 15)   # SIGTERM
                time.sleep(1)
                try:
                    os.kill(self._pid, 0)
                    os.kill(self._pid, 9)
                except OSError:
                    pass
        except Exception:
            pass
        finally:
            self._proc = None
            self._pid = None
            self._model = None
            self._port = None
            self._started_at = None
            self._clear_pid()

    async def _wait_until_ready(self, port: int, timeout: float = 30) -> bool:
        """轮询检查端口是否开始监听。"""
        import httpx
        url = f"http://127.0.0.1:{port}/v1/models"
        deadline = time.time() + timeout
        async with httpx.AsyncClient(timeout=5.0) as client:
            while time.time() < deadline:
                try:
                    r = await client.get(url)
                    if r.status_code < 500:
                        return True
                except (httpx.ConnectError, httpx.ReadTimeout, OSError):
                    pass
                await asyncio.sleep(0.5)
        return False

    # ── 模型列表 ─────────────────────────────────────────────────────────

    def list_local_models(self) -> list[dict[str, str]]:
        """列出 models/ 目录下所有 .gguf 文件。"""
        models = []
        for p in MODELS_DIR.glob("*.gguf"):
            size_mb = p.stat().st_size // (1024 * 1024)
            models.append({
                "id": p.name,
                "name": p.name,
                "size_mb": size_mb,
                "path": str(p),
            })
        return sorted(models, key=lambda m: m["name"])

    # ── 模型下载 ─────────────────────────────────────────────────────────

    async def download_model(
        self,
        url: str,
        dest: Optional[str] = None,
    ) -> DownloadProgress:
        """从 HuggingFace 或直链下载 .gguf 模型文件。

        HuggingFace 路径别名：用户输入 "Qwen/Qwen3-4B-GGUF/Qwen3-4B-Q4_K_M.gguf"
        自动转换为 "https://huggingface.co/<repo>/resolve/main/<file>"

        支持断点续传（Content-Range）。

        返回 DownloadProgress 对象（含当前进度），进度通过 get_download_progress() 查询。
        """
        import httpx

        # URL 协议白名单：只允许 http/https（file://、ftp:// 等一律拒绝）。
        # 必须先于 HF 别名转换判定，否则 ftp:// 会被误拼成 HF 地址。
        if "://" in url:
            scheme = url.split("://", 1)[0].lower()
            if scheme not in ("http", "https"):
                prog = DownloadProgress(url=url, dest=str(dest or ""),
                                        error=f"仅允许 http/https 下载地址（收到 {scheme}）")
                return prog

        # 路径别名转换
        if not url.startswith("http"):
            url = f"https://huggingface.co/{url}/resolve/main"

        # 解析文件名
        filename = url.rstrip("/").split("/")[-1]
        if not filename.endswith(".gguf"):
            filename += ".gguf"

        if dest:
            dest_path = Path(dest)
            # 目标关押：自定义 dest 必须仍位于 models/ 目录内（防 ../ 越界写盘）
            models_root = MODELS_DIR.resolve()
            resolved = (models_root / dest_path).resolve() if not dest_path.is_absolute() \
                else dest_path.resolve()
            if resolved != models_root and models_root not in resolved.parents:
                prog = DownloadProgress(url=url, dest=str(dest_path),
                                        error=f"下载目标必须在 models/ 目录内: {dest}")
                return prog
            dest_path = resolved
        else:
            dest_path = MODELS_DIR / filename

        # 构造 HTTP 头
        headers = {}
        resume_bytes = 0
        if dest_path.exists():
            resume_bytes = dest_path.stat().st_size
            headers["Range"] = f"bytes={resume_bytes}-"

        # 获取文件大小
        total_bytes: Optional[int] = None
        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=10.0, read=60.0, write=30.0, pool=10.0), follow_redirects=True) as client:
                head = await client.head(url, headers={"Range": "bytes=0-0"})
                total_raw = head.headers.get("Content-Length")
                if total_raw:
                    total_bytes = int(total_raw)
                # Content-Range 响应时 total_bytes 在 Content-Range 头里
                cr = head.headers.get("Content-Range", "")
                m = re.search(r"/(\d+)", cr)
                if m:
                    total_bytes = int(m.group(1))
        except Exception:
            pass

        prog = DownloadProgress(
            url=url,
            dest=str(dest_path),
            total_bytes=total_bytes,
            downloaded_bytes=resume_bytes,
        )
        with self._dl_lock:
            self._downloading[url] = prog

        try:
            mode = "ab" if resume_bytes > 0 else "wb"
            async with httpx.AsyncClient(
                timeout=httpx.Timeout(300.0, connect=10.0, read=300.0, write=30.0, pool=10.0),
                follow_redirects=True,
            ) as client:
                t0 = time.time()
                last_bytes = resume_bytes
                async with client.stream("GET", url, headers=headers) as resp:
                    if resp.status_code not in (200, 206):
                        raise RuntimeError(f"HTTP {resp.status_code}")
                    with open(dest_path, mode) as f:
                        async for chunk in resp.aiter_bytes(chunk_size=8192):
                            f.write(chunk)
                            prog.downloaded_bytes += len(chunk)

                            # 速度 & ETA
                            elapsed = time.time() - t0
                            if elapsed > 0.5:
                                speed_bps = (prog.downloaded_bytes - last_bytes) / elapsed
                                speed_str = _format_speed(speed_bps)
                                if prog.total_bytes and speed_bps > 0:
                                    remain = prog.total_bytes - prog.downloaded_bytes
                                    eta_s = remain / speed_bps
                                    prog.eta = _format_eta(eta_s)
                                else:
                                    prog.eta = ""
                                prog.speed = speed_str
                                last_bytes = prog.downloaded_bytes
                                t0 = time.time()

                            if prog.total_bytes:
                                prog.progress_pct = min(prog.downloaded_bytes / prog.total_bytes * 100, 100)
        except Exception as e:
            prog.error = str(e)
        finally:
            prog.done = True
            with self._dl_lock:
                self._downloading[url] = prog

        return prog

    def get_download_progress(self, url: str) -> Optional[DownloadProgress]:
        """查询下载进度。"""
        with self._dl_lock:
            return self._downloading.get(url)

    def list_downloads(self) -> list[DownloadProgress]:
        """列出所有活跃下载。"""
        with self._dl_lock:
            return list(self._downloading.values())


def _format_speed(bps: float) -> str:
    if bps >= 1e9:
        return f"{bps/1e9:.1f} GB/s"
    if bps >= 1e6:
        return f"{bps/1e6:.1f} MB/s"
    if bps >= 1e3:
        return f"{bps/1e3:.1f} KB/s"
    return f"{bps:.0f} B/s"


def _format_eta(seconds: float) -> str:
    if seconds < 60:
        return f"{seconds:.0f}s"
    if seconds < 3600:
        return f"{seconds/60:.0f}m"
    return f"{seconds/3600:.1f}h"


# ─────────────────────────────────────────────────────────────────────────────
# 全局单例
# ─────────────────────────────────────────────────────────────────────────────
_lm: Optional[LlamaManager] = None


def get_llama_manager() -> LlamaManager:
    global _lm
    if _lm is None:
        _lm = LlamaManager()
    return _lm