"""Per-user Docker container manager for main agent."""

from __future__ import annotations

import json
import re
import shutil
import subprocess
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Optional

from config import (
    MAX_ACTIVE_USER_CONTAINERS,
    OUTPUT_FORMATS,
    TERMINAL_SANDBOX_BIN,
    TERMINAL_SANDBOX_BINDS,
    TERMINAL_SANDBOX_CPUS,
    TERMINAL_SANDBOX_ENV,
    TERMINAL_SANDBOX_IMAGE,
    TERMINAL_SANDBOX_MEMORY,
    TERMINAL_SANDBOX_MODE,
    TERMINAL_SANDBOX_MOUNT_PATH,
    TERMINAL_SANDBOX_NAME_PREFIX,
    TERMINAL_SANDBOX_NETWORK,
    TERMINAL_SANDBOX_REQUIRE,
    LOGS_DIR,
)
from modules.container_monitor import collect_stats, inspect_state


@dataclass
class ContainerHandle:
    """Lightweight record describing a user workspace container."""

    username: str
    mode: str
    workspace_path: str
    mount_path: str
    container_name: Optional[str] = None
    container_id: Optional[str] = None
    sandbox_bin: Optional[str] = None
    created_at: float = field(default_factory=time.time)
    last_active: float = field(default_factory=time.time)

    def touch(self):
        self.last_active = time.time()

    def to_dict(self) -> Dict:
        return {
            "username": self.username,
            "mode": self.mode,
            "workspace_path": self.workspace_path,
            "mount_path": self.mount_path,
            "container_name": self.container_name,
            "container_id": self.container_id,
            "created_at": self.created_at,
            "last_active": self.last_active,
        }


class UserContainerManager:
    """Create and track long-lived containers for each logged-in user."""

    def __init__(
        self,
        sandbox_mode: Optional[str] = None,
        max_containers: int = MAX_ACTIVE_USER_CONTAINERS,
    ):
        self.sandbox_mode = (sandbox_mode or TERMINAL_SANDBOX_MODE or "host").lower()
        self.max_containers = max_containers
        self.image = TERMINAL_SANDBOX_IMAGE
        self.mount_path = TERMINAL_SANDBOX_MOUNT_PATH or "/workspace"
        self.network = TERMINAL_SANDBOX_NETWORK
        self.cpus = TERMINAL_SANDBOX_CPUS
        self.memory = TERMINAL_SANDBOX_MEMORY
        self.binds = list(TERMINAL_SANDBOX_BINDS)
        self.sandbox_bin = TERMINAL_SANDBOX_BIN or "docker"
        self.name_prefix = TERMINAL_SANDBOX_NAME_PREFIX or "agent-user"
        self.require = bool(TERMINAL_SANDBOX_REQUIRE)
        self.extra_env = dict(TERMINAL_SANDBOX_ENV)
        self._containers: Dict[str, ContainerHandle] = {}
        self._lock = threading.Lock()
        self._stats_log_path = Path(LOGS_DIR).expanduser().resolve() / "container_stats.log"
        self._stats_log_path.parent.mkdir(parents=True, exist_ok=True)
        if not self._stats_log_path.exists():
            self._stats_log_path.touch()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def ensure_container(self, username: str, workspace_path: str) -> ContainerHandle:
        username = self._normalize_username(username)
        workspace = str(Path(workspace_path).expanduser().resolve())
        Path(workspace).mkdir(parents=True, exist_ok=True)

        with self._lock:
            handle = self._containers.get(username)
            if handle:
                if handle.mode == "docker" and not self._is_container_running(handle):
                    self._containers.pop(username, None)
                    self._kill_container(handle.container_name, handle.sandbox_bin)
                    handle = None
                else:
                    handle.workspace_path = workspace
                    handle.touch()
                    return handle

            if not self._has_capacity(username):
                raise RuntimeError("资源繁忙：容器配额已用尽，请稍候再试。")

            handle = self._create_handle(username, workspace)
            self._containers[username] = handle
            return handle

    def release_container(self, username: str, reason: str = "logout"):
        username = self._normalize_username(username)
        with self._lock:
            handle = self._containers.pop(username, None)
        if not handle:
            return
        if handle.mode == "docker" and handle.container_name:
            self._kill_container(handle.container_name, handle.sandbox_bin)
            print(f"{OUTPUT_FORMATS['info']} 容器已释放: {handle.container_name} ({reason})")

    def has_capacity(self, username: Optional[str] = None) -> bool:
        username = self._normalize_username(username) if username else None
        with self._lock:
            if username and username in self._containers:
                return True
            if self.max_containers <= 0:
                return True
            return len(self._containers) < self.max_containers

    def get_handle(self, username: str) -> Optional[ContainerHandle]:
        username = self._normalize_username(username)
        with self._lock:
            handle = self._containers.get(username)
            if handle:
                handle.touch()
            return handle

    def list_containers(self) -> Dict[str, Dict]:
        with self._lock:
            return {user: handle.to_dict() for user, handle in self._containers.items()}

    def get_container_status(self, username: str, include_stats: bool = True) -> Dict:
        username = self._normalize_username(username)
        with self._lock:
            handle = self._containers.get(username)
        if not handle:
            return {"username": username, "mode": "host"}

        info = {
            "username": username,
            "mode": handle.mode,
            "workspace_path": handle.workspace_path,
            "mount_path": handle.mount_path,
            "container_name": handle.container_name,
            "created_at": handle.created_at,
            "last_active": handle.last_active,
        }

        if handle.mode == "docker" and include_stats:
            stats = collect_stats(handle.container_name, handle.sandbox_bin)
            state = inspect_state(handle.container_name, handle.sandbox_bin)
            if stats:
                info["stats"] = stats
                self._log_stats(username, stats)
            if state:
                info["state"] = state

        return info

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _has_capacity(self, username: str) -> bool:
        if self.max_containers <= 0:
            return True
        existing = 1 if username in self._containers else 0
        return (len(self._containers) - existing) < self.max_containers

    def _create_handle(self, username: str, workspace: str) -> ContainerHandle:
        if self.sandbox_mode != "docker":
            return self._host_handle(username, workspace)

        docker_path = shutil.which(self.sandbox_bin or "docker")
        if not docker_path:
            message = f"未找到容器运行时 {self.sandbox_bin}"
            if self.require:
                raise RuntimeError(message)
            print(f"{OUTPUT_FORMATS['warning']} {message}，回退到宿主机执行。")
            return self._host_handle(username, workspace)

        if not self.image:
            raise RuntimeError("TERMINAL_SANDBOX_IMAGE 未配置，无法启动容器。")

        container_name = self._build_container_name(username)
        self._kill_container(container_name, docker_path)

        cmd = [
            docker_path,
            "run",
            "-d",
            "--name",
            container_name,
            "-w",
            self.mount_path,
            "-v",
            f"{workspace}:{self.mount_path}",
        ]

        if self.network:
            cmd += ["--network", self.network]
        if self.cpus:
            cmd += ["--cpus", str(self.cpus)]
        if self.memory:
            cmd += ["--memory", str(self.memory)]
        for bind in self.binds:
            chunk = bind.strip()
            if chunk:
                cmd += ["-v", chunk]

        envs = {
            "PYTHONIOENCODING": "utf-8",
            "TERM": "xterm-256color",
        }
        envs.update({k: v for k, v in self.extra_env.items() if v is not None})
        for key, value in envs.items():
            cmd += ["-e", f"{key}={value}"]

        cmd.append(self.image)
        cmd += ["tail", "-f", "/dev/null"]

        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        if result.returncode != 0:
            message = result.stderr.strip() or result.stdout.strip() or "容器启动失败"
            if self.require:
                raise RuntimeError(message)
            print(f"{OUTPUT_FORMATS['warning']} {message}，回退到宿主机。")
            return self._host_handle(username, workspace)

        container_id = result.stdout.strip() or None
        print(f"{OUTPUT_FORMATS['success']} 启动用户容器: {container_name} ({username})")
        return ContainerHandle(
            username=username,
            mode="docker",
            workspace_path=workspace,
            mount_path=self.mount_path,
            container_name=container_name,
            container_id=container_id,
            sandbox_bin=docker_path,
        )

    def _host_handle(self, username: str, workspace: str) -> ContainerHandle:
        return ContainerHandle(
            username=username,
            mode="host",
            workspace_path=workspace,
            mount_path=workspace,
        )

    def _kill_container(self, container_name: Optional[str], docker_bin: Optional[str]):
        if not container_name or not docker_bin:
            return
        subprocess.run(
            [docker_bin, "rm", "-f", container_name],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )

    def _is_container_running(self, handle: ContainerHandle) -> bool:
        if handle.mode != "docker" or not handle.container_name or not handle.sandbox_bin:
            return True
        try:
            result = subprocess.run(
                [
                    handle.sandbox_bin,
                    "inspect",
                    "-f",
                    "{{.State.Running}}",
                    handle.container_name,
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.DEVNULL,
                text=True,
                timeout=3,
                check=False,
            )
        except (OSError, subprocess.SubprocessError):
            return False
        return result.returncode == 0 and result.stdout.strip().lower() == "true"

    def _build_container_name(self, username: str) -> str:
        slug = re.sub(r"[^a-z0-9\-]", "-", username.lower()).strip("-")
        if not slug:
            slug = "user"
        return f"{self.name_prefix}-{slug}"

    def _log_stats(self, username: str, stats: Dict):
        try:
            record = {
                "username": username,
                "timestamp": time.time(),
                "stats": stats,
            }
            with self._stats_log_path.open('a', encoding='utf-8') as fh:
                fh.write(json.dumps(record, ensure_ascii=False) + "\n")
        except Exception:
            pass

    @staticmethod
    def _normalize_username(username: Optional[str]) -> str:
        return (username or "").strip().lower()