from google.adk.agents import LlmAgent
from google.adk.models.lite_llm import LiteLlm
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
from google.adk.tools.function_tool import FunctionTool
from mcp.client.stdio import StdioServerParameters
import os
import asyncio
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# --- 使用 vLLM 端点的智能体 ---
api_base_url = os.getenv("VLLM_API_BASE", "https://9router.aqroid.cn/v1")
model_name = os.getenv("VLLM_MODEL", "")
api_key = os.getenv("VLLM_API_KEY", "")

# Agent 可访问的工作目录
WORKSPACE_DIR = os.getenv("AGENT_WORKSPACE_DIR", r"D:\nzy\workspace_git")

# --- 文件系统 MCP 工具 ---
filesystem_mcp = McpToolset(
    connection_params=StdioConnectionParams(
        server_params=StdioServerParameters(
            command="npx",
            args=[
                "-y",
                "@modelcontextprotocol/server-filesystem",
                os.path.abspath(WORKSPACE_DIR),
            ],
        ),
        timeout=300000,
    ),
)

# --- 网页抓取 MCP 工具（Fetch）---
# 暂时注释：请求体过大导致 413 错误，后续需要时再启用
# fetch_mcp = McpToolset(
#     connection_params=StdioConnectionParams(
#         server_params=StdioServerParameters(
#             command="npx",
#             args=[
#                 "-y",
#                 "@modelcontextprotocol/server-fetch",
#             ],
#         ),
#         timeout=300000,
#     ),
# )

# --- Tavily 搜索 MCP 工具 ---
# 暂时注释：请求体过大导致 413 错误，后续需要时再启用
# tavily_api_key = os.getenv("TAVILY_API_KEY", "")
# tavily_mcp = McpToolset(
#     connection_params=StdioConnectionParams(
#         server_params=StdioServerParameters(
#             command="npx",
#             args=[
#                 "-y",
#                 "tavily-mcp@latest",
#             ],
#             env={
#                 "TAVILY_API_KEY": tavily_api_key,
#             },
#         ),
#         timeout=300000,
#     ),
# )


# --- 终端命令执行工具（Python 原生，绕开 MCP 通信问题）---
async def run_command(command: str, cwd: str = None, timeout: int = 300000000) -> str:
    """
    在终端中执行一条命令，返回输出结果。

    Args:
        command: 要执行的命令，如 'npm run build'、'python -m pytest' 等
        cwd: 命令执行的工作目录，默认使用 AGENT_WORKSPACE_DIR
        timeout: 超时时间（秒），默认 300

    Returns:
        命令执行结果（stdout + stderr + 状态）
    """
    if not command:
        return "错误：命令不能为空"

    work_dir = cwd or os.path.abspath(WORKSPACE_DIR)
    if not os.path.isdir(work_dir):
        return f"错误：工作目录不存在 {work_dir}"

    try:
        proc = await asyncio.create_subprocess_shell(
            command,
            cwd=work_dir,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout_bytes, stderr_bytes = await asyncio.wait_for(
            proc.communicate(), timeout=timeout
        )
    except asyncio.TimeoutError:
        proc.kill()
        await proc.wait()
        return f"命令执行超时（{timeout}秒）: {command}"
    except Exception as e:
        return f"命令执行出错: {e}"

    stdout = stdout_bytes.decode("utf-8", errors="replace")
    stderr = stderr_bytes.decode("utf-8", errors="replace")

    parts = []
    if stdout:
        parts.append(f"[stdout]\n{stdout}")
    if stderr:
        parts.append(f"[stderr]\n{stderr}")

    output = "\n".join(parts) if parts else "(无输出)"

    max_len = 10000
    if len(output) > max_len:
        output = output[:max_len] + f"\n\n...（输出已截断，共 {len(output)} 字符）"

    status = "成功" if proc.returncode == 0 else f"失败 (退出码 {proc.returncode})"
    return f"命令执行{status}\n{output}"


# 注册为 ADK 工具
run_command_tool = FunctionTool(run_command)


root_agent = LlmAgent(
    model=LiteLlm(
        model=model_name,
        api_base=api_base_url,
        api_key=api_key if api_key else None,
        custom_llm_provider="openai",
    ),
    name="luna_agent",
    description="全栈开发子 Agent（gpt-5.6-luna），可以读写文件、浏览目录、执行开发任务。",
    instruction=(
        "你是luna，一个使用gpt5.6-luna的前后端开发子agent\n"
        "\n"
        "## 记忆能力\n"
        "- 你拥有长期记忆，之前和用户的对话中提到的项目信息、技术偏好、任务历史都会被记住\n"
        "- 系统会自动从记忆中检索与当前任务相关的历史上下文，注入到对话中\n"
        "- 重要的项目信息（技术栈、目录结构、编码规范等）会自动沉淀到记忆里\n"
        "\n"
        "## 工作流程\n"
        "1. 先理解任务需求和项目上下文\n"
        "2. 使用文件系统工具浏览项目结构、读取相关文件\n"
        "3. 编写或修改代码\n"
        "4. 使用 run_command 工具运行编译/构建/测试，确保代码可正常工作\n"
        "5. 验证结果后，按指定格式报告完成情况\n"
        "\n"
        "## 工作边界\n"
        "- 所有文件操作限定在分配的工作目录范围内\n"
        "- 你拥有的工具：文件系统操作（读/写/列目录）、终端命令执行\n"
        "- 你可以自主完成：代码编写、bug 修复、样式调整、接口修改、简单重构\n"
        "- 遇到不熟悉的技术或 API，先查阅项目内的现有代码和文档参考\n"
        "- 需要上报的情况：\n"
        "  • 架构设计或重大技术选型决策\n"
        "  • 依赖包版本不兼容导致的编译/运行时错误（需要升级/降级依赖时）\n"
        "  • 工具调用异常、环境配置问题、命令超时等非代码问题\n"
        "  • 超出你能力范围或不确定的问题\n"
        "\n"
        "## 编译/构建守则\n"
        "- 写完代码后，优先运行编译或构建命令验证\n"
        "- 编译报错时，先判断错误类型：\n"
        "  • 代码语法/逻辑错误 → 自行修复后重试\n"
        "  • 依赖缺失或版本不兼容 → 上报，由主控决定处理方式\n"
        "  • 环境/工具问题 → 上报\n"
        "- 连续修复 3 次仍无法通过编译时，上报当前状态和所有错误信息\n"
        "- 只有编译通过后才算任务完成\n"
        "\n"
        "## 报告格式\n"
        "完成任务后，结构化报告：\n"
        "**状态**：成功 / 部分完成 / 失败（需上报）\n"
        "**修改的文件**：列出所有修改的文件路径\n"
        "**变更摘要**：简述做了什么\n"
        "**验证结果**：编译/测试是否通过，如有警告需列出\n"
        "**需要主控关注**：如有需要上报的问题，详细说明"
    ),
    tools=[filesystem_mcp, run_command_tool],
)