"""执行 Shell 命令工具。"""

from __future__ import annotations

import asyncio
import json
import os
import signal
import subprocess
from dataclasses import dataclass, field
from tempfile import NamedTemporaryFile
from typing import Any, Literal, Optional, TextIO, Type

from pydantic import BaseModel, Field

from app.agent.tools.base import MoviePilotTool
from app.agent.tools.tags import ToolTag
from app.agent.tools.impl._terminal_session import (
    TERMINAL_DEFAULT_READ_BYTES,
    TERMINAL_MAX_READ_BYTES,
    TERMINAL_WAIT_DEFAULT_MS,
    terminal_session_manager,
)
from app.log import logger


DEFAULT_TIMEOUT_SECONDS = 60
MAX_TIMEOUT_SECONDS = 300
MAX_OUTPUT_PREVIEW_BYTES = 10 * 1024
READ_CHUNK_SIZE = 4096
KILL_GRACE_SECONDS = 3
COMMAND_CONCURRENCY_LIMIT = 2
COMMAND_FORBIDDEN_KEYWORDS = (
    ":(){ :|:& };:",
    "dd if=/dev/zero",
    "mkfs",
    "reboot",
    "shutdown",
)

_command_semaphore = asyncio.Semaphore(COMMAND_CONCURRENCY_LIMIT)


@dataclass
class _CommandOutput:
    """保存前 10KB 预览，并在超限时将完整输出写入临时文件。"""

    preview_limit_bytes: int
    preview_entries: list[tuple[str, str]] = field(default_factory=list)
    captured_bytes: int = 0
    preview_truncated: bool = False
    temp_file_path: Optional[str] = None
    temp_file_handle: Optional[TextIO] = None
    last_written_stream: Optional[str] = None

    @staticmethod
    def _clip_text_to_bytes(text: str, byte_limit: int) -> str:
        """按 UTF-8 字节数截断文本，避免截断后出现非法字符。"""
        if byte_limit <= 0:
            return ""
        return text.encode("utf-8")[:byte_limit].decode("utf-8", errors="ignore")

    def _write_chunk(self, stream_name: str, text: str) -> None:
        """把输出分片按 stdout/stderr 分段写入临时文件。"""
        if not self.temp_file_handle or not text:
            return

        if self.last_written_stream != stream_name:
            if self.temp_file_handle.tell() > 0:
                self.temp_file_handle.write("\n")
            title = "标准输出" if stream_name == "stdout" else "错误输出"
            self.temp_file_handle.write(f"[{title}]\n")
            self.last_written_stream = stream_name

        self.temp_file_handle.write(text)

    def _ensure_temp_file(self) -> None:
        """首次超出预览上限时创建临时文件并补写已缓存预览。"""
        if self.temp_file_handle:
            return

        temp_file = NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            suffix=".log",
            prefix="moviepilot-command-",
            delete=False,
        )
        self.temp_file_path = temp_file.name
        self.temp_file_handle = temp_file
        for stream_name, chunk in self.preview_entries:
            self._write_chunk(stream_name, chunk)

    def close(self) -> None:
        """关闭临时文件句柄，确保输出落盘。"""
        if not self.temp_file_handle:
            return
        self.temp_file_handle.flush()
        self.temp_file_handle.close()
        self.temp_file_handle = None

    def append(self, stream_name: str, text: str) -> None:
        """追加一段输出，超出预览上限后只保留完整日志文件。"""
        if not text:
            return

        if self.temp_file_handle:
            self._write_chunk(stream_name, text)
            return

        chunk_bytes = len(text.encode("utf-8"))
        remaining = self.preview_limit_bytes - self.captured_bytes
        if chunk_bytes <= remaining:
            self.preview_entries.append((stream_name, text))
            self.captured_bytes += chunk_bytes
            return

        self.preview_truncated = True
        self._ensure_temp_file()
        self._write_chunk(stream_name, text)

        preview = self._clip_text_to_bytes(text, remaining)
        if preview:
            self.preview_entries.append((stream_name, preview))
            self.captured_bytes += len(preview.encode("utf-8"))

    @property
    def stdout(self) -> str:
        """返回当前保留的 stdout 预览。"""
        return "".join(
            text for stream_name, text in self.preview_entries if stream_name == "stdout"
        ).strip()

    @property
    def stderr(self) -> str:
        """返回当前保留的 stderr 预览。"""
        return "".join(
            text for stream_name, text in self.preview_entries if stream_name == "stderr"
        ).strip()


class ExecuteCommandInput(BaseModel):
    """执行 Shell 命令工具的输入参数模型。"""

    explanation: Optional[str] = Field(None, description="Clear explanation of why this command action is needed")
    action: Optional[Literal["start", "read", "wait", "write", "kill", "run"]] = Field(
        "start",
        description=(
            "Command action. start launches a managed background session and returns "
            "session_id. read/wait/write/kill operate on that session. run executes "
            "once and waits until completion or timeout."
        ),
    )
    command: Optional[str] = Field(
        None,
        description="Shell command. Required for action=start or action=run.",
    )
    session_id: Optional[str] = Field(
        None,
        description="Command session id returned by action=start.",
    )
    input_text: Optional[str] = Field(
        None,
        description="Text to send to stdin for action=write. Use \\u0003 for Ctrl+C.",
    )
    signal_name: Optional[str] = Field(
        "TERM",
        description="Signal for action=kill, such as TERM, INT, KILL, or 15.",
    )
    cwd: Optional[str] = Field(
        None,
        description="Working directory for action=start or action=run.",
    )
    env: Optional[dict[str, Any]] = Field(
        None,
        description="Additional environment variables for action=start.",
    )
    use_pty: Optional[bool] = Field(
        True,
        description="Use a pseudo terminal for action=start when supported.",
    )
    since_seq: Optional[int] = Field(
        None,
        description="For action=read/wait, return output chunks after this seq.",
    )
    max_bytes: Optional[int] = Field(
        TERMINAL_DEFAULT_READ_BYTES,
        description="For action=read/wait, maximum output bytes to return.",
    )
    timeout_ms: Optional[int] = Field(
        TERMINAL_WAIT_DEFAULT_MS,
        description="For action=wait, maximum segmented wait time in milliseconds.",
    )
    timeout: Optional[int] = Field(
        60,
        description="For action=run, max execution time in seconds.",
    )


class ExecuteCommandTool(MoviePilotTool):
    """统一执行和管理 Shell 命令的 Agent 工具。"""

    name: str = "execute_command"
    tags: list[str] = [
        ToolTag.Read,
        ToolTag.Command,
        ToolTag.Admin,
    ]
    description: str = (
        "Start and manage shell commands on the server. By default action=start "
        "launches a background session and immediately returns session_id/status/"
        "last_seq/output_until_seq. Call the same tool with action=read, wait, "
        "write, or kill to poll output, wait in short segments, send stdin, or "
        "terminate it. Use action=run only when a one-shot bounded command result "
        "is preferred."
    )
    args_schema: Type[BaseModel] = ExecuteCommandInput
    require_admin: bool = True
    result_max_chars = TERMINAL_MAX_READ_BYTES + 4096

    def get_tool_message(self, **kwargs) -> Optional[str]:
        """根据命令动作生成友好的提示消息。"""
        action = kwargs.get("action") or "start"
        command = kwargs.get("command")
        session_id = kwargs.get("session_id")
        if action in {"start", "run"}:
            return f"执行系统命令: {command or ''}"
        if action == "read":
            return f"读取命令输出: {session_id or ''}"
        if action == "wait":
            return f"等待命令会话: {session_id or ''}"
        if action == "write":
            return f"写入命令输入: {session_id or ''}"
        if action == "kill":
            return f"终止命令会话: {session_id or ''}"
        return f"处理命令会话: {session_id or command or ''}"

    @staticmethod
    def _dump(payload: dict[str, Any]) -> str:
        """把结构化命令会话结果转换为 Agent 容易解析的 JSON 字符串。"""
        return json.dumps(payload, ensure_ascii=False, indent=2)

    @staticmethod
    def _require_session_id(session_id: Optional[str]) -> str:
        """校验会话型 action 必须传入 session_id。"""
        if not session_id:
            raise ValueError("action 需要传入 session_id")
        return session_id

    @staticmethod
    def _require_command(command: Optional[str]) -> str:
        """校验启动型 action 必须传入 command。"""
        if not command or not command.strip():
            raise ValueError("action 需要传入 command")
        return command

    @staticmethod
    def _validate_command(command: str) -> None:
        """复用旧工具的基础危险命令过滤，避免明显破坏性命令进入 shell。"""
        for keyword in COMMAND_FORBIDDEN_KEYWORDS:
            if keyword in command:
                raise ValueError(f"命令包含禁止使用的关键字 '{keyword}'")

        # 检查是否使用了 rm -r/R 删除根目录或一级目录，防止误杀多级目录
        import re
        import os.path
        tokens = re.split(r'\s+', command.strip())
        if any(t == "rm" or t.endswith("/rm") for t in tokens):
            has_r = False
            for token in tokens:
                if token.startswith("-") and ("r" in token or "R" in token):
                    has_r = True
                    break
            
            if has_r:
                for token in tokens:
                    # 提取可能包含目标路径的部分（去除重定向、管道、分号等末尾干扰）
                    m = re.match(r'^([^;\|&><]+)', token)
                    if m:
                        clean_token = m.group(1).strip('"\'')
                        # 仅对绝对路径进行一级目录限制
                        if clean_token.startswith('/'):
                            norm_path = os.path.normpath(clean_token)
                            if re.match(r'^/[^/]*$', norm_path) or re.match(r'^/[^/]*/$', norm_path):
                                raise ValueError(f"不允许使用 rm 命令删除根目录或一级目录: {clean_token}")

    @staticmethod
    def _normalize_timeout(timeout: Optional[int]) -> tuple[int, Optional[str]]:
        """限制一次性执行命令的最长运行时间。"""
        try:
            normalized = int(timeout or DEFAULT_TIMEOUT_SECONDS)
        except (TypeError, ValueError):
            normalized = DEFAULT_TIMEOUT_SECONDS

        if normalized <= 0:
            return DEFAULT_TIMEOUT_SECONDS, "timeout 参数无效，已使用默认 60 秒"
        if normalized > MAX_TIMEOUT_SECONDS:
            return (
                MAX_TIMEOUT_SECONDS,
                f"timeout 参数超过上限，已从 {normalized} 秒限制为 {MAX_TIMEOUT_SECONDS} 秒",
            )
        return normalized, None

    @staticmethod
    def _subprocess_kwargs() -> dict:
        """为一次性命令创建独立进程组，便于超时清理整棵子进程。"""
        kwargs = {
            "stdin": subprocess.DEVNULL,
            "stdout": asyncio.subprocess.PIPE,
            "stderr": asyncio.subprocess.PIPE,
        }
        if os.name == "posix":
            kwargs["start_new_session"] = True
        elif os.name == "nt":
            kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
        return kwargs

    @staticmethod
    async def _read_stream(
        stream: asyncio.StreamReader,
        stream_name: str,
        output: _CommandOutput,
    ) -> None:
        """按块读取一次性命令输出，只把前 10KB 保留在返回结果中。"""
        while True:
            chunk = await stream.read(READ_CHUNK_SIZE)
            if not chunk:
                break
            output.append(stream_name, chunk.decode("utf-8", errors="replace"))

    @staticmethod
    def _terminate_process(process: Any, sig: int) -> None:
        """向进程组发送终止信号，不支持进程组的平台回退为单进程终止。"""
        try:
            if os.name == "posix":
                os.killpg(process.pid, sig)
            elif sig == getattr(signal, "SIGKILL", None):
                process.kill()
            else:
                process.terminate()
        except ProcessLookupError:
            pass

    @classmethod
    async def _cleanup_process(
        cls,
        process: Any,
        wait_task: asyncio.Task,
    ) -> None:
        """先温和终止，失败后强杀，避免超时 shell 遗留子进程。"""
        if wait_task.done():
            return

        cls._terminate_process(process, signal.SIGTERM)
        try:
            await asyncio.wait_for(
                asyncio.shield(wait_task), timeout=KILL_GRACE_SECONDS
            )
            return
        except asyncio.TimeoutError:
            pass

        kill_signal = getattr(signal, "SIGKILL", signal.SIGTERM)
        cls._terminate_process(process, kill_signal)
        try:
            await asyncio.wait_for(
                asyncio.shield(wait_task), timeout=KILL_GRACE_SECONDS
            )
        except asyncio.TimeoutError:
            logger.warning("命令进程强制清理超时: pid=%s", process.pid)

    @staticmethod
    async def _finish_reader_tasks(reader_tasks: list[asyncio.Task]) -> None:
        """等待一次性命令输出读取任务退出，异常只记录不影响工具返回。"""
        if not reader_tasks:
            return
        done, pending = await asyncio.wait(reader_tasks, timeout=1)
        for task in pending:
            task.cancel()
        results = await asyncio.gather(*done, *pending, return_exceptions=True)
        for result in results:
            if isinstance(result, Exception) and not isinstance(
                result, asyncio.CancelledError
            ):
                logger.debug("命令输出读取任务异常: %s", result)

    @staticmethod
    def _format_run_result(
        *,
        exit_code: Optional[int],
        output: _CommandOutput,
        timeout: int,
        timed_out: bool,
        timeout_note: Optional[str],
    ) -> str:
        """格式化 action=run 的兼容文本结果。"""
        if timed_out:
            result = f"命令执行超时 (限制: {timeout}秒，已终止进程)"
        else:
            result = f"命令执行完成 (退出码: {exit_code})"

        if timeout_note:
            result += f"\n\n提示:\n{timeout_note}"
        if output.temp_file_path:
            file_note = "截至命令终止前的完整输出" if timed_out else "完整输出"
            result += (
                "\n\n提示:\n"
                f"命令输出超过 10KB，仅返回前 {MAX_OUTPUT_PREVIEW_BYTES} 字节内容。\n"
                f"{file_note}已写入临时文件: {output.temp_file_path}\n"
                "如需完整内容，请继续读取该文件。"
            )
        if output.stdout:
            result += f"\n\n标准输出:\n{output.stdout}"
        if output.stderr:
            result += f"\n\n错误输出:\n{output.stderr}"
        if output.preview_truncated:
            result += "\n\n...(仅展示前 10KB 内容)"
        if not output.stdout and not output.stderr:
            result += "\n\n(无输出内容)"
        return result

    async def _run_once(
        self,
        *,
        command: str,
        timeout: Optional[int],
        cwd: Optional[str] = None,
    ) -> str:
        """按旧模式一次性执行命令，等待完成或超时后返回文本结果。"""
        self._validate_command(command)
        normalized_timeout, timeout_note = self._normalize_timeout(timeout)

        async with _command_semaphore:
            process = await asyncio.create_subprocess_shell(
                command,
                cwd=cwd,
                **self._subprocess_kwargs(),
            )
            output = _CommandOutput(preview_limit_bytes=MAX_OUTPUT_PREVIEW_BYTES)
            wait_task = asyncio.create_task(process.wait())
            reader_tasks = [
                asyncio.create_task(self._read_stream(process.stdout, "stdout", output)),
                asyncio.create_task(self._read_stream(process.stderr, "stderr", output)),
            ]

            timed_out = False
            try:
                await asyncio.wait_for(
                    asyncio.shield(wait_task), timeout=normalized_timeout
                )
            except asyncio.TimeoutError:
                timed_out = True
                await self._cleanup_process(process, wait_task)

            try:
                await self._finish_reader_tasks(reader_tasks)
            finally:
                output.close()

        return self._format_run_result(
            exit_code=process.returncode,
            output=output,
            timeout=normalized_timeout,
            timed_out=timed_out,
            timeout_note=timeout_note,
        )

    async def run(
        self,
        action: Optional[str] = "start",
        command: Optional[str] = None,
        session_id: Optional[str] = None,
        input_text: Optional[str] = None,
        signal_name: Optional[str] = "TERM",
        cwd: Optional[str] = None,
        env: Optional[dict[str, Any]] = None,
        use_pty: Optional[bool] = True,
        since_seq: Optional[int] = None,
        max_bytes: Optional[int] = TERMINAL_DEFAULT_READ_BYTES,
        timeout_ms: Optional[int] = TERMINAL_WAIT_DEFAULT_MS,
        timeout: Optional[int] = 60,
        **kwargs,
    ) -> str:
        """执行命令动作：默认后台启动，也支持读取、等待、写入、终止和一次性执行。"""
        normalized_action = (action or "start").strip().lower()
        logger.info(
            "执行工具: %s, action=%s, command=%s, session_id=%s",
            self.name,
            normalized_action,
            command,
            session_id,
        )

        try:
            if normalized_action == "start":
                start_command = self._require_command(command)
                self._validate_command(start_command)
                payload = await terminal_session_manager.start(
                    command=start_command,
                    cwd=cwd,
                    env=env,
                    use_pty=use_pty,
                )
                return self._dump(payload)

            if normalized_action == "read":
                payload = await terminal_session_manager.read(
                    session_id=self._require_session_id(session_id),
                    since_seq=since_seq,
                    max_bytes=max_bytes,
                )
                return self._dump(payload)

            if normalized_action == "wait":
                payload = await terminal_session_manager.wait(
                    session_id=self._require_session_id(session_id),
                    timeout_ms=timeout_ms,
                    since_seq=since_seq,
                    max_bytes=max_bytes,
                )
                return self._dump(payload)

            if normalized_action == "write":
                payload = await terminal_session_manager.write(
                    session_id=self._require_session_id(session_id),
                    input_text=input_text or "",
                )
                return self._dump(payload)

            if normalized_action == "kill":
                payload = await terminal_session_manager.kill(
                    session_id=self._require_session_id(session_id),
                    sig=signal_name,
                )
                return self._dump(payload)

            if normalized_action == "run":
                return await self._run_once(
                    command=self._require_command(command),
                    timeout=timeout,
                    cwd=cwd,
                )

            raise ValueError(f"不支持的 action: {action}")
        except Exception as err:
            logger.error("执行命令 action 失败: %s", err, exc_info=True)
            return self._dump({"error": str(err), "status": "error", "action": normalized_action})
