"""
元宝 Agent 工具测试脚本
测试本地工具调用功能，包含 conda 命令
"""

import subprocess
import sys
from pathlib import Path

# 测试工具注册表
class ToolTester:
    def __init__(self):
        self.results = []
    
    def test(self, name: str, func, *args):
        """运行单个测试"""
        print(f"\n{'='*50}")
        print(f"测试: {name}")
        print(f"参数: {args}")
        print(f"{'='*50}")
        
        try:
            result = func(*args)
            success = "[错误]" not in result and "失败" not in result
            status = "[PASS]" if success else "[FAIL]"
            
            print(f"结果:\n{result[:500]}{'...' if len(result) > 500 else ''}")
            print(f"状态: {status}")
            
            self.results.append({
                "name": name,
                "success": success,
                "result": result
            })
            return success
            
        except Exception as e:
            print(f"异常: {e}")
            self.results.append({
                "name": name,
                "success": False,
                "result": str(e)
            })
            return False
    
    def report(self):
        """输出测试报告"""
        print(f"\n{'='*50}")
        print("测试报告")
        print(f"{'='*50}")
        
        passed = sum(1 for r in self.results if r["success"])
        total = len(self.results)
        
        for r in self.results:
            icon = "[OK]" if r["success"] else "[FAIL]"
            print(f"{icon} {r['name']}")
        
        print(f"\n总计: {passed}/{total} 通过")
        return passed == total


# ============ 工具实现 ============

def tool_read_file(path: str) -> str:
    """读取本地文件内容"""
    try:
        file_path = Path(path).expanduser().resolve()
        if not file_path.exists():
            return f"[错误] 文件不存在: {path}"
        content = file_path.read_text(encoding="utf-8", errors="ignore")
        if len(content) > 3000:
            content = content[:3000] + f"\n... (已截断，共 {len(content)} 字符)"
        return content
    except Exception as e:
        return f"[错误] {e}"


def tool_list_dir(path: str = ".") -> str:
    """列出目录内容"""
    try:
        dir_path = Path(path).expanduser().resolve()
        if not dir_path.exists():
            return f"[错误] 目录不存在: {path}"
        
        items = []
        for item in sorted(dir_path.iterdir()):
            prefix = "[DIR]" if item.is_dir() else "[FILE]"
            size = ""
            if item.is_file():
                try:
                    size_bytes = item.stat().st_size
                    size = f" ({size_bytes:,} bytes)"
                except:
                    pass
            items.append(f"{prefix} {item.name}{size}")
        
        return f"目录: {dir_path}\n" + "\n".join(items[:30])
    except Exception as e:
        return f"[错误] {e}"


def tool_run_command(command: str) -> str:
    """执行本地命令"""
    try:
        # 使用 shell=True 以支持 conda 等需要 shell 环境的命令
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=30,
            encoding="utf-8",
            errors="ignore"
        )
        
        output = result.stdout
        if result.stderr:
            output += f"\n[stderr]\n{result.stderr}"
        
        if len(output) > 3000:
            output = output[:3000] + "\n... (已截断)"
        
        return f"命令: {command}\n退出码: {result.returncode}\n\n{output}"
    except subprocess.TimeoutExpired:
        return f"[错误] 命令执行超时: {command}"
    except Exception as e:
        return f"[错误] {e}"


def tool_conda_env_list() -> str:
    """列出所有 conda 环境"""
    return tool_run_command("conda env list")


def tool_conda_info() -> str:
    """显示 conda 信息"""
    return tool_run_command("conda info")


def tool_conda_list(env: str = "base") -> str:
    """列出指定环境的包"""
    return tool_run_command(f"conda list -n {env}")


def tool_python_version() -> str:
    """显示 Python 版本"""
    return tool_run_command("python --version")


def tool_pip_list() -> str:
    """列出 pip 包"""
    return tool_run_command("pip list")


# ============ 主测试 ============

def main():
    print("=" * 60)
    print("元宝 Agent 工具测试")
    print("=" * 60)
    
    tester = ToolTester()
    
    # 1. 基础文件操作测试
    tester.test("读取当前脚本", tool_read_file, __file__)
    tester.test("列出当前目录", tool_list_dir, ".")
    tester.test("列出用户目录", tool_list_dir, "~")
    
    # 2. 基础命令测试
    tester.test("Python 版本", tool_python_version)
    tester.test("系统信息 (systeminfo)", tool_run_command, "systeminfo | findstr /B /C:\"OS\" | head -5")
    tester.test("当前目录", tool_run_command, "cd")
    
    # 3. Conda 相关测试
    print("\n" + "=" * 60)
    print("Conda 测试")
    print("=" * 60)
    
    # 先检查 conda 是否可用
    conda_check = tool_run_command("conda --version")
    print(f"\n检查 conda:\n{conda_check}")
    
    if "错误" not in conda_check and "conda" in conda_check.lower():
        tester.test("conda 环境列表", tool_conda_env_list)
        tester.test("conda 信息", tool_conda_info)
        tester.test("base 环境包列表 (前20行)", tool_run_command, "conda list -n base | head -20")
    else:
        print("⚠️ 未检测到 conda，跳过 conda 测试")
        tester.results.append({
            "name": "conda 环境列表",
            "success": False,
            "result": "conda 未安装或未添加到 PATH"
        })
    
    # 4. pip 测试
    tester.test("pip 包列表 (前20行)", tool_run_command, "pip list | head -20")
    
    # 5. 复杂命令测试
    tester.test("查找 Python 文件", tool_run_command, "dir /s /b *.py 2>nul | head -10")
    
    # 输出报告
    all_passed = tester.report()
    
    print("\n" + "=" * 60)
    if all_passed:
        print("[SUCCESS] 所有测试通过！")
    else:
        print("[WARNING] 部分测试失败，请检查环境配置")
    print("=" * 60)
    
    return 0 if all_passed else 1


if __name__ == "__main__":
    sys.exit(main())