"""TerminalSession - 终端会话状态管理

维护 cwd、env、history，支持跨命令保持状态。
通过在命令前注入 cd 和 export 语句实现跨命令状态持久化。
"""

from __future__ import annotations

import asyncio
import logging
import os
import re
import shlex
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Any

from agentkit.tools.output_parser import OutputParser, ParsedOutput

logger = logging.getLogger(__name__)

_ENV_KEY_PATTERN = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*$')


@dataclass
class CommandRecord:
    """命令执行记录

    Attributes:
        command: 执行的命令
        exit_code: 退出码
        output: 标准输出+错误输出
        cwd: 执行时的工作目录
        timestamp: 执行时间戳
        duration_ms: 执行耗时（毫秒）
    """

    command: str
    exit_code: int
    output: str
    cwd: str
    timestamp: float
    duration_ms: int


class TerminalSession:
    """终端会话 - 跨命令保持 cwd/env/history 状态

    通过在命令前注入 `cd {cwd} && ` 和 `export K=V && ` 实现跨命令状态持久化。
    每次命令执行后自动更新 cwd 和 env 状态。

    Usage:
        session = TerminalSession(session_id="build-01")
        result = await session.execute("cd /tmp")
        result = await session.execute("pwd")  # 输出 /tmp
    """

    def __init__(
        self,
        session_id: str,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        max_history: int = 1000,
    ):
        self.session_id = session_id
        self._cwd = cwd or os.getcwd()
        self._env: dict[str, str] = dict(env or os.environ)
        self._history: deque[CommandRecord] = deque(maxlen=max_history)
        self._output_parser = OutputParser()
        self._created_at = time.time()

    @property
    def cwd(self) -> str:
        """当前工作目录"""
        return self._cwd

    @property
    def env(self) -> dict[str, str]:
        """当前环境变量（副本）"""
        return dict(self._env)

    @property
    def history(self) -> list[CommandRecord]:
        """命令执行历史（副本）"""
        return list(self._history)

    @property
    def created_at(self) -> float:
        """会话创建时间戳"""
        return self._created_at

    def get_cwd(self) -> str:
        """获取当前工作目录"""
        return self._cwd

    def set_cwd(self, cwd: str) -> None:
        """手动设置当前工作目录"""
        self._cwd = cwd

    def get_env(self) -> dict[str, str]:
        """获取当前环境变量（副本）"""
        return dict(self._env)

    def set_env(self, key: str, value: str) -> None:
        """设置单个环境变量"""
        self._env[key] = value

    def update_env(self, env: dict[str, str]) -> None:
        """批量更新环境变量"""
        self._env.update(env)

    def get_history(self) -> list[CommandRecord]:
        """获取命令执行历史（副本）"""
        return list(self._history)

    async def execute(
        self,
        command: str,
        timeout: float | None = None,
    ) -> ParsedOutput:
        """在会话上下文中执行命令

        自动在命令前注入 cd 和 export 语句以保持会话状态。
        执行后自动更新 cwd 和 env。

        Args:
            command: 要执行的命令
            timeout: 超时时间（秒），None 表示不超时

        Returns:
            ParsedOutput 结构化解析结果
        """
        full_command = self._build_command(command)
        start = time.monotonic()

        try:
            proc = await asyncio.create_subprocess_shell(
                full_command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
                env=self._env,
            )
            try:
                stdout, _ = await asyncio.wait_for(
                    proc.communicate(),
                    timeout=timeout,
                )
            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()
                output = f"命令执行超时（{timeout}s）"
                exit_code = -1
            else:
                output = stdout.decode("utf-8", errors="replace") if stdout else ""
                exit_code = proc.returncode if proc.returncode is not None else 0
        except Exception as e:
            output = str(e)
            exit_code = -1

        duration_ms = int((time.monotonic() - start) * 1000)

        # 更新会话状态
        self._update_state_after_execution(command, output, exit_code)

        # 记录历史
        record = CommandRecord(
            command=command,
            exit_code=exit_code,
            output=output,
            cwd=self._cwd,
            timestamp=time.time(),
            duration_ms=duration_ms,
        )
        self._add_history(record)

        # 解析输出
        parsed = self._output_parser.parse(output, exit_code)
        logger.debug(
            "Session %s: command=%r exit_code=%d duration=%dms",
            self.session_id,
            command,
            exit_code,
            duration_ms,
        )
        return parsed

    def _build_command(self, command: str) -> str:
        """构建带会话状态的完整命令

        在原始命令前注入 cd 和 export 语句。
        """
        parts: list[str] = []

        # 注入 cd
        if self._cwd:
            parts.append(f"cd {shlex.quote(self._cwd)}")

        # 注入环境变量
        for key, value in self._env.items():
            if not _ENV_KEY_PATTERN.match(key):
                continue  # Skip invalid env key names
            parts.append(f"export {shlex.quote(key)}={shlex.quote(value)}")

        parts.append(command)
        return " && ".join(parts)

    def _update_state_after_execution(
        self, command: str, output: str, exit_code: int
    ) -> None:
        """命令执行后更新会话状态

        解析 cd 和 export 命令更新 cwd 和 env。
        """
        if exit_code != 0:
            return

        # 解析 cd 命令更新 cwd
        self._parse_cd_commands(command, output)

        # 解析 export 命令更新 env
        self._parse_export_commands(command)

    def _parse_cd_commands(self, command: str, output: str) -> None:
        """从命令中解析 cd 并更新 cwd

        支持:
        - cd /path
        - cd dir (相对路径，需要通过 pwd 获取实际路径)
        - cd - (切换到上一个目录)
        """
        import re

        # 匹配 cd 命令（可能出现在 && 链中）
        cd_pattern = re.compile(r"(?:^|\s|&&\s*)cd\s+(.+?)(?:\s*&&|\s*$)")
        matches = cd_pattern.findall(command)

        for target in matches:
            target = target.strip().strip("'\"")
            if not target:
                continue

            if target == "-":
                # cd - 切换到 OLDPWD
                old_pwd = self._env.get("OLDPWD")
                if old_pwd:
                    self._cwd = old_pwd
            elif os.path.isabs(target):
                self._cwd = target
            else:
                # 相对路径：拼接后规范化
                new_cwd = os.path.normpath(os.path.join(self._cwd, target))
                self._cwd = new_cwd

    def _parse_export_commands(self, command: str) -> None:
        """从命令中解析 export 并更新 env

        支持:
        - export KEY=VALUE
        - export KEY="VALUE WITH SPACES"
        """
        import re

        export_pattern = re.compile(
            r"(?:^|\s|&&\s*)export\s+(\w+)=(.+?)(?:\s*&&|\s*$)"
        )
        matches = export_pattern.findall(command)

        for key, value in matches:
            value = value.strip().strip("'\"")
            self._env[key] = value

    def _add_history(self, record: CommandRecord) -> None:
        """添加命令记录到历史，deque maxlen 自动淘汰最旧记录"""
        self._history.append(record)

    def close(self) -> None:
        """关闭会话，清理资源"""
        logger.info(
            "Session %s closed: %d commands executed",
            self.session_id,
            len(self._history),
        )


class TerminalSessionManager:
    """终端会话管理器 - 按 ID 管理多个 TerminalSession

    Usage:
        manager = TerminalSessionManager()
        session = manager.get_or_create("build-01")
        session = manager.get("build-01")
        manager.remove("build-01")
    """

    def __init__(self, max_sessions: int = 100):
        self._sessions: dict[str, TerminalSession] = {}
        self._max_sessions = max_sessions

    def get_or_create(
        self,
        session_id: str,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
    ) -> TerminalSession:
        """获取或创建会话

        Args:
            session_id: 会话 ID
            cwd: 初始工作目录（仅创建时使用）
            env: 初始环境变量（仅创建时使用）

        Returns:
            TerminalSession 实例
        """
        if session_id not in self._sessions:
            if len(self._sessions) >= self._max_sessions:
                # 移除最旧的会话
                oldest_id = min(
                    self._sessions, key=lambda k: self._sessions[k].created_at
                )
                self.remove(oldest_id)
            self._sessions[session_id] = TerminalSession(
                session_id=session_id,
                cwd=cwd,
                env=env,
            )
            logger.info("Session created: %s", session_id)
        return self._sessions[session_id]

    def get(self, session_id: str) -> TerminalSession | None:
        """获取会话，不存在返回 None"""
        return self._sessions.get(session_id)

    def remove(self, session_id: str) -> None:
        """移除并关闭会话"""
        session = self._sessions.pop(session_id, None)
        if session:
            session.close()

    def list_sessions(self) -> list[str]:
        """列出所有会话 ID"""
        return list(self._sessions.keys())

    def has_session(self, session_id: str) -> bool:
        """检查会话是否存在"""
        return session_id in self._sessions

    def close_all(self) -> None:
        """关闭所有会话"""
        for session_id in list(self._sessions.keys()):
            self.remove(session_id)
