"""工具运行状态 + 来源管理

职责：
- 管理每个工具的来源（source）：本地 uv、Docker、外部 Hub 等
- 管理运行状态：启动/停止工具进程
- 同一个 tool_id 可以有多个 source，路由层选择活跃的那个
"""

from __future__ import annotations

import json
import logging
import socket
import subprocess
import threading
import time
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field

from tool_agent.config import settings
from tool_agent.models import ToolMeta, ToolStatus
from tool_agent.registry.registry import ToolRegistry

logger = logging.getLogger(__name__)


# ---- 来源模型 ----

class SourceType(str, Enum):
    LOCAL = "local"       # 本地 uv 项目
    DOCKER = "docker"     # Docker 容器
    REMOTE = "remote"     # 远程 API / 云端服务
    HUB = "hub"           # 兼容旧数据（等价于 remote）


class ToolSource(BaseModel):
    """工具的一个来源"""
    type: SourceType
    # local
    host_dir: str = ""                     # 本地项目目录
    # docker
    container_id: str = ""
    image: str = ""
    # remote / hub
    remote_url: str = ""                   # 远程服务基础 URL
    remote_path: str = ""                  # 远程服务内工具路径
    remote_api_key: str = ""               # 远程服务 API Key
    hub_url: str = ""                      # 兼容旧数据
    hub_tool_path: str = ""                # 兼容旧数据
    hub_api_key: str = ""                  # 兼容旧数据
    # 通用
    endpoint_path: str = "/"               # HTTP API 路径
    http_method: str = "POST"
    internal_port: int = 0                 # 进程/容器内服务端口


class ProcessState(str, Enum):
    STOPPED = "stopped"
    STARTING = "starting"
    RUNNING = "running"
    ERROR = "error"


class ToolRoute(BaseModel):
    """工具的路由信息 = 来源列表 + 运行状态"""
    tool_id: str
    sources: list[ToolSource] = Field(default_factory=list)
    active_source: int = 0                 # 当前使用的 source 索引
    state: ProcessState = ProcessState.STOPPED
    pid: int | None = None
    port: int | None = None                # 宿主机上的实际服务端口
    started_at: datetime | None = None
    last_error: str | None = None


# ---- 持久化 ----

class SourceStore:
    """来源信息持久化 — data/sources.json"""

    def __init__(self, path: Path | None = None) -> None:
        self._path = path or (settings.data_dir / "sources.json")
        self._lock = threading.Lock()

    def load(self) -> dict[str, list[dict]]:
        if not self._path.exists():
            return {}
        data = json.loads(self._path.read_text(encoding="utf-8"))
        return data.get("sources", {})

    def save(self, sources: dict[str, list[dict]]) -> None:
        with self._lock:
            self._path.parent.mkdir(parents=True, exist_ok=True)
            self._path.write_text(
                json.dumps({"sources": sources}, indent=2, ensure_ascii=False, default=str),
                encoding="utf-8",
            )

    def add_source(self, tool_id: str, source: ToolSource) -> None:
        with self._lock:
            all_sources = self.load()
            tool_sources = all_sources.get(tool_id, [])
            tool_sources.append(source.model_dump(mode="json"))
            all_sources[tool_id] = tool_sources
            # 直接写入，不调用 save() 避免重复获取锁
            self._path.parent.mkdir(parents=True, exist_ok=True)
            self._path.write_text(
                json.dumps({"sources": all_sources}, indent=2, ensure_ascii=False, default=str),
                encoding="utf-8",
            )

    def get_sources(self, tool_id: str) -> list[ToolSource]:
        data = self.load()
        return [ToolSource(**s) for s in data.get(tool_id, [])]

    def remove_tool(self, tool_id: str) -> None:
        with self._lock:
            all_sources = self.load()
            if tool_id in all_sources:
                del all_sources[tool_id]
                self.save(all_sources)


# ---- 状态管理器 ----

class ToolStatusManager:
    """工具进程生命周期管理 + 来源管理"""

    def __init__(self, registry: ToolRegistry) -> None:
        self._registry = registry
        self._source_store = SourceStore()
        self._routes: dict[str, ToolRoute] = {}
        self._popen_refs: dict[str, subprocess.Popen] = {}
        self._lock = threading.Lock()
        self._sync_from_registry()

    def _sync_from_registry(self) -> None:
        """从注册表 + 来源存储同步"""
        for tool in self._registry.list_all():
            tid = tool.tool_id
            if tid not in self._routes:
                sources = self._source_store.get_sources(tid)
                # 自动发现本地项目
                if not sources:
                    local_dir = settings.tools_dir / "local" / tid
                    if local_dir.exists() and (local_dir / "main.py").exists():
                        # 使用相对路径
                        relative_path = f"tools/local/{tid}"
                        src = ToolSource(type=SourceType.LOCAL, host_dir=relative_path)
                        self._source_store.add_source(tid, src)
                        sources = [src]
                self._routes[tid] = ToolRoute(tool_id=tid, sources=sources)

    def add_source(self, tool_id: str, source: ToolSource) -> None:
        """为工具添加来源"""
        self._source_store.add_source(tool_id, source)
        if tool_id in self._routes:
            self._routes[tool_id].sources.append(source)
        else:
            self._routes[tool_id] = ToolRoute(tool_id=tool_id, sources=[source])

    @staticmethod
    def _get_free_port() -> int:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]

    def start_tool(self, tool_id: str) -> ToolRoute:
        """启动工具"""
        route = self._routes.get(tool_id)
        if not route or not route.sources:
            return self._set_error(tool_id, "No source available")

        if route.state == ProcessState.RUNNING:
            return route

        source = route.sources[route.active_source]

        if source.type == SourceType.LOCAL:
            return self._start_local(tool_id, source)
        elif source.type == SourceType.DOCKER:
            return self._check_docker(tool_id, source)
        elif source.type == SourceType.HUB or source.type == SourceType.REMOTE:
            # 远程工具不需要启动本地进程，直接标记 running
            route.state = ProcessState.RUNNING
            with self._lock:
                self._routes[tool_id] = route
            return route
        else:
            return self._set_error(tool_id, f"Unsupported source type: {source.type}")

    def _start_local(self, tool_id: str, source: ToolSource) -> ToolRoute:
        # 将相对路径转为绝对路径
        host_dir_path = Path(source.host_dir)
        if not host_dir_path.is_absolute():
            tool_dir = (settings.base_dir / host_dir_path).resolve()
        else:
            tool_dir = host_dir_path

        main_py = tool_dir / "main.py"
        if not main_py.exists():
            return self._set_error(tool_id, f"main.py not found in {tool_dir}")

        port = self._get_free_port()
        try:
            log_dir = tool_dir / "tests"
            log_dir.mkdir(parents=True, exist_ok=True)

            with open(log_dir / "server.log", "w") as lf:
                proc = subprocess.Popen(
                    ["uv", "run", "--directory", str(tool_dir),
                     "python", "main.py", "--port", str(port)],
                    stdout=lf, stderr=subprocess.STDOUT, text=True,
                )

            if self._wait_for_port(port, timeout=10):
                route = ToolRoute(
                    tool_id=tool_id,
                    sources=self._routes.get(tool_id, ToolRoute(tool_id=tool_id)).sources,
                    state=ProcessState.RUNNING,
                    pid=proc.pid, port=port,
                    started_at=datetime.now(timezone.utc),
                )
                with self._lock:
                    self._routes[tool_id] = route
                    self._popen_refs[tool_id] = proc
                logger.info(f"Started '{tool_id}' on port {port} (PID {proc.pid})")
                return route
            else:
                proc.kill()
                return self._set_error(tool_id, f"Service failed to start. Check {log_dir / 'server.log'}")
        except Exception as e:
            return self._set_error(tool_id, str(e))

    def _check_docker(self, tool_id: str, source: ToolSource) -> ToolRoute:
        if source.internal_port:
            route = self._routes.get(tool_id, ToolRoute(tool_id=tool_id))
            route.state = ProcessState.RUNNING
            route.port = source.internal_port
            with self._lock:
                self._routes[tool_id] = route
            return route
        return self._set_error(tool_id, "No port mapping for Docker tool")

    @staticmethod
    def _wait_for_port(port: int, timeout: int = 10) -> bool:
        start = time.time()
        while time.time() - start < timeout:
            try:
                with socket.create_connection(("127.0.0.1", port), timeout=1):
                    return True
            except (ConnectionRefusedError, OSError):
                time.sleep(0.3)
        return False

    def stop_tool(self, tool_id: str) -> ToolRoute:
        with self._lock:
            proc = self._popen_refs.pop(tool_id, None)
            if proc:
                try:
                    proc.kill()
                    proc.wait(timeout=5)
                except Exception:
                    pass
                logger.info(f"Stopped '{tool_id}'")
            route = self._routes.get(tool_id, ToolRoute(tool_id=tool_id))
            route.state = ProcessState.STOPPED
            route.pid = None
            route.port = None
            self._routes[tool_id] = route
            return route

    def stop_all(self) -> int:
        count = 0
        for tool_id in list(self._popen_refs.keys()):
            self.stop_tool(tool_id)
            count += 1
        return count

    def get_status(self, tool_id: str) -> ToolRoute | None:
        with self._lock:
            proc = self._popen_refs.get(tool_id)
            if proc and proc.poll() is not None:
                route = self._routes.get(tool_id, ToolRoute(tool_id=tool_id))
                route.state = ProcessState.ERROR
                route.last_error = f"Process exited with code {proc.returncode}"
                self._routes[tool_id] = route
                del self._popen_refs[tool_id]
        return self._routes.get(tool_id)

    def get_primary_source(self, tool_id: str) -> ToolSource | None:
        """当前工具选用的来源（与 route.active_source 下标对应）。"""
        route = self._routes.get(tool_id)
        if not route or not route.sources:
            return None
        idx = route.active_source
        if idx < 0 or idx >= len(route.sources):
            return route.sources[0]
        return route.sources[idx]

    def get_active_endpoint(self, tool_id: str) -> dict | None:
        """获取工具当前活跃的调用端点"""
        route = self._routes.get(tool_id)
        if not route or route.state != ProcessState.RUNNING or not route.sources:
            return None
        source = route.sources[route.active_source]
        if source.type == SourceType.HUB:
            return {
                "type": "hub",
                "url": f"{source.hub_url.rstrip('/')}{source.hub_tool_path}",
                "http_method": source.http_method,
                "api_key": source.hub_api_key,
            }
        else:
            return {
                "type": source.type.value,
                "url": f"http://127.0.0.1:{route.port}{source.endpoint_path}",
                "http_method": source.http_method,
            }

    def list_status(self) -> list[ToolRoute]:
        self._sync_from_registry()
        for tool_id in list(self._popen_refs.keys()):
            self.get_status(tool_id)
        return list(self._routes.values())

    def _set_error(self, tool_id: str, error: str) -> ToolRoute:
        route = self._routes.get(tool_id, ToolRoute(tool_id=tool_id))
        route.state = ProcessState.ERROR
        route.last_error = error
        with self._lock:
            self._routes[tool_id] = route
        logger.error(f"Tool '{tool_id}' error: {error}")
        return route
