import asyncio
import subprocess
import os
import sys
import pty
import select
import fcntl
import signal
from typing import Optional, Tuple
from pydantic import BaseModel, Field
from .base import BaseTool


def is_textual_app(command: str) -> bool:
    """
    检测命令是否可能是 Textual 应用
    
    检测方式：
    - 命令中包含 textual/Textual
    - 运行 python 脚本且脚本内容包含 textual
    """
    command_lower = command.lower()
    
    # 检查命令是否直接运行 textual
    if "textual" in command_lower:
        return True
    
    # 检查是否运行 python 脚本
    if "python" in command_lower or "python3" in command_lower:
        # 提取脚本路径
        parts = command.split()
        for i, part in enumerate(parts):
            if part in ("python", "python3", "python3.10", "python3.11", "python3.12"):
                if i + 1 < len(parts):
                    script_path = parts[i + 1]
                    # 如果是文件，检查内容是否包含 textual
                    if os.path.isfile(script_path):
                        try:
                            with open(script_path, 'r') as f:
                                content = f.read()
                                if 'textual' in content.lower() or 'Textual' in content:
                                    return True
                        except:
                            pass
                    # 如果是 -m 模块，检查是否与 textual 相关
                    elif script_path == "-m" and i + 2 < len(parts):
                        module = parts[i + 2]
                        if "textual" in module.lower():
                            return True
    
    return False


async def run_with_pty(
    command: str,
    workdir: Optional[str] = None,
    timeout: int = 120,
    on_output: Optional[callable] = None,
) -> Tuple[str, int]:
    """
    使用 PTY 运行命令，隔离终端状态
    
    Args:
        command: 要执行的命令
        workdir: 工作目录
        timeout: 超时时间（秒）
        on_output: 输出回调函数，接收 (is_stdout: bool, data: str)
        
    Returns:
        (output: str, returncode: int)
    """
    import time
    import fcntl
    
    home = os.path.expanduser("~")
    shell = os.environ.get("SHELL", "/bin/bash")
    
    # 获取 shell 初始化命令
    init_cmd = ""
    if "zsh" in shell:
        zshrc = os.path.join(home, ".zshrc")
        if os.path.exists(zshrc):
            init_cmd = f"source {zshrc} 2>/dev/null; "
    else:
        bash_profile = os.path.join(home, ".bash_profile")
        bashrc = os.path.join(home, ".bashrc")
        if os.path.exists(bash_profile):
            init_cmd = f"source {bash_profile} 2>/dev/null; "
        elif os.path.exists(bashrc):
            init_cmd = f"source {bashrc} 2>/dev/null; "
    
    full_command = f"{init_cmd}{command}"
    
    # 创建 PTY
    master_fd, slave_fd = pty.openpty()
    
    # 设置为非阻塞模式
    flags = fcntl.fcntl(master_fd, fcntl.F_GETFL)
    fcntl.fcntl(master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
    
    # 启动子进程，使用 [shell, "-c", full_command] 替代 shell=True
    # 这样可以确保使用用户指定的 shell（如 zsh），避免 /bin/sh 导致加载用户配置失败
    process = subprocess.Popen(
        [shell, "-c", full_command],
        cwd=workdir,
        stdin=slave_fd,
        stdout=slave_fd,
        stderr=slave_fd,
        preexec_fn=os.setsid,  # 创建新的进程组
        env={**os.environ, "TERM": "xterm-256color", **(env_override or {})},
    )
    
    # 关闭 slave 端
    os.close(slave_fd)
    
    # 收集输出
    output_parts = []
    start_time = time.time()
    
    while True:
        # 检查超时
        if time.time() - start_time > timeout:
            # 杀死进程组
            try:
                os.killpg(os.getpgid(process.pid), signal.SIGTERM)
            except:
                pass
            output_parts.append(f"\n[Error: Command timed out after {timeout} seconds]")
            break
        
        # 使用 select 监听 PTY（带超时）
        ready, _, _ = select.select([master_fd], [], [], 0.5)
        
        if ready:
            try:
                data = os.read(master_fd, 4096)
                if data:
                    # 解码并存储
                    text = data.decode('utf-8', errors='replace')
                    output_parts.append(text)
                    
                    # 如果有回调，调用它
                    if on_output:
                        on_output(text)
            except OSError as e:
                if e.errno == 11:  # EAGAIN
                    pass  # 没有数据，继续
                else:
                    break
        
        # 检查进程是否结束
        ret = process.poll()
        if ret is not None:
            # 读取剩余数据（带超时）
            remaining_start = time.time()
            while time.time() - remaining_start < 1:  # 最多等1秒
                ready, _, _ = select.select([master_fd], [], [], 0.2)
                if ready:
                    try:
                        data = os.read(master_fd, 4096)
                        if data:
                            text = data.decode('utf-8', errors='replace')
                            output_parts.append(text)
                            if on_output:
                                on_output(text)
                    except OSError:
                        break
                else:
                    break
            break
    
    os.close(master_fd)
    process.wait()
    
    output = "".join(output_parts)
    return output, process.returncode


class BashToolArgs(BaseModel):
    command: str = Field(..., description="Shell command to execute")
    workdir: Optional[str] = Field(
        None, description="Working directory for the command (avoid using cd in command)"
    )
    timeout: int = Field(
        120, description="Timeout in seconds (default: 120)"
    )


class BashTool(BaseTool):
    name = "bash"
    description = """Executes a bash command in a shell session with optional timeout.

IMPORTANT: This tool is for terminal operations like git, npm, pytest, etc. DO NOT use it for file operations - use specialized tools instead.

Usage:
- Use workdir parameter to change directories (AVOID using 'cd dir && command')
- Always quote file paths with spaces: "path with spaces/file.txt"
- Output is truncated if too long
- Timeout default: 120 seconds

Avoid using bash for:
- File search: Use Glob (NOT find or ls)
- Content search: Use Grep (NOT grep)
- Read files: Use Read (NOT cat/head/tail)
- Edit files: Use Edit (NOT sed/awk)
- Write files: Use Write (NOT echo >)

Good example: workdir="/foo/bar", command="pytest tests"
Bad example: command="cd /foo/bar && pytest tests"
"""
    args_schema = BashToolArgs

    def _get_shell_init_command(self) -> str:
        """获取 shell 初始化命令，加载用户配置"""
        home = os.path.expanduser("~")
        shell = os.environ.get("SHELL", "/bin/bash")

        # 根据 shell 类型选择配置文件
        if "zsh" in shell:
            # zsh: 加载 .zshrc
            zshrc = os.path.join(home, ".zshrc")
            if os.path.exists(zshrc):
                return f"source {zshrc} 2>/dev/null; "
        else:
            # bash: 加载 .bash_profile 或 .bashrc
            bash_profile = os.path.join(home, ".bash_profile")
            bashrc = os.path.join(home, ".bashrc")
            if os.path.exists(bash_profile):
                return f"source {bash_profile} 2>/dev/null; "
            elif os.path.exists(bashrc):
                return f"source {bashrc} 2>/dev/null; "

        return ""

    async def run(
        self, command: str, workdir: Optional[str] = None, timeout: int = 120
    ) -> str:
        try:
            # 检测是否是 Textual 应用
            if is_textual_app(command):
                # 使用 PTY 模式运行，隔离终端状态
                # 设置 TERM=dumb 和 NO_COLOR=1 强制非交互/纯文本模式
                # 避免 ANSI 转义序列干扰客户端 TUI
                output, returncode = await run_with_pty(
                    command, workdir, timeout,
                    env_override={"TERM": "dumb", "NO_COLOR": "1"}
                )
                
                if returncode != 0:
                    output += f"\n[Exit code: {returncode}]"
                
                return output
            
            # 普通命令，使用原有方式运行
            # 获取用户 shell
            shell = os.environ.get("SHELL", "/bin/bash")

            # 构建完整命令，加载用户配置
            init_cmd = self._get_shell_init_command()
            full_command = f"{init_cmd}{command}"

            process = await asyncio.create_subprocess_shell(
                full_command,
                cwd=workdir,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                executable=shell,  # 使用用户的 shell
            )

            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(), timeout=timeout
                )
            except asyncio.TimeoutError:
                process.kill()
                return "Error: Command timed out."

            output = stdout.decode().strip()
            error = stderr.decode().strip()

            result_str = output
            if error:
                result_str += f"\nSTDERR:\n{error}"

            if len(result_str) > 5000:
                result_str = (
                    result_str[:2500]
                    + "\n... [Output Truncated] ...\n"
                    + result_str[-2500:]
                )

            return result_str
        except Exception as e:
            return f"Error executing command: {str(e)}"
