"""
元宝全自动代理 MVP
验证思路：通过 Playwright 拦截元宝对话，识别本地调用意图，
执行本地操作后将结果注入回对话。

核心机制：
1. 监听元宝输入框，检测特殊指令标记（如 /local 或 @agent）
2. 拦截用户输入，在本地执行
3. 将执行结果作为"用户消息"注入回对话
4. 元宝基于结果继续回复
"""

import asyncio
import json
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional, Callable

from playwright.async_api import async_playwright, Page, BrowserContext

YUANBAO_URL = "https://yuanbao.tencent.com/chat"

# 本地工具注册表
TOOLS: dict[str, Callable] = {}


def register_tool(name: str):
    """工具装饰器"""
    def decorator(func: Callable):
        TOOLS[name] = func
        return func
    return decorator


# ============ 本地工具实现 ============

@register_tool("read_file")
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) > 5000:
            content = content[:5000] + f"\n... (已截断，共 {len(content)} 字符)"
        return f"[文件内容: {path}]\n```\n{content}\n```"
    except Exception as e:
        return f"[错误] 读取文件失败: {e}"


@register_tool("list_dir")
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()):
            item_type = "📁" if item.is_dir() else "📄"
            items.append(f"{item_type} {item.name}")
        return f"[目录列表: {path}]\n" + "\n".join(items[:50])
    except Exception as e:
        return f"[错误] 列出目录失败: {e}"


@register_tool("run_command")
def tool_run_command(command: str) -> str:
    """执行本地命令"""
    try:
        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) > 5000:
            output = output[:5000] + f"\n... (已截断，共 {len(output)} 字符)"
        return f"[命令执行: {command}]\n返回码: {result.returncode}\n```\n{output}\n```"
    except subprocess.TimeoutExpired:
        return f"[错误] 命令执行超时: {command}"
    except Exception as e:
        return f"[错误] 执行命令失败: {e}"


@register_tool("web_search")
def tool_web_search(query: str) -> str:
    """模拟网络搜索（MVP 阶段返回提示）"""
    return f"[搜索: {query}]\nMVP 阶段暂未接入真实搜索，请直接询问元宝。"


# ============ 指令解析器 ============

class CommandParser:
    """解析用户输入中的本地调用指令"""
    
    # 指令模式：@agent:tool_name(args) 或 /local tool_name args
    PATTERNS = [
        # @agent:read_file("/path/to/file")
        r'@agent:(\w+)\s*\(\s*["\']?([^"\')]*)["\']?\s*\)',
        # /local read_file /path/to/file
        r'/local\s+(\w+)\s+(.+?)(?:\n|$)',
        # ```local
        # read_file /path/to/file
        # ```
        r'```local\s*\n(\w+)\s+(.+?)\n```',
    ]
    
    def parse(self, text: str) -> Optional[dict]:
        """解析指令，返回工具调用信息或 None"""
        text = text.strip()
        
        # 模式 1: @agent:tool_name(args)
        match = re.search(self.PATTERNS[0], text)
        if match:
            return {
                "tool": match.group(1),
                "args": [match.group(2)],
                "raw": match.group(0)
            }
        
        # 模式 2: /local tool_name args
        match = re.search(self.PATTERNS[1], text)
        if match:
            return {
                "tool": match.group(1),
                "args": [match.group(2).strip()],
                "raw": match.group(0)
            }
        
        # 模式 3: ```local ... ```
        match = re.search(self.PATTERNS[2], text, re.DOTALL)
        if match:
            return {
                "tool": match.group(1),
                "args": [match.group(2).strip()],
                "raw": match.group(0)
            }
        
        return None


# ============ 元宝页面操作 ============

class YuanbaoController:
    """控制元宝网页的输入输出"""
    
    def __init__(self, page: Page):
        self.page = page
        self.parser = CommandParser()
    
    async def find_input_box(self):
        """查找输入框"""
        selectors = [
            "div[contenteditable='true'][data-placeholder]",
            "div[contenteditable='true']",
            "textarea[placeholder]",
            "textarea",
        ]
        for sel in selectors:
            try:
                el = await self.page.wait_for_selector(sel, timeout=3000)
                if el:
                    return el
            except Exception:
                continue
        return None
    
    async def inject_message(self, content: str):
        """将内容注入到输入框并发送"""
        input_box = await self.find_input_box()
        if not input_box:
            print("[MVP] 未找到输入框")
            return False
        
        # 点击输入框
        await input_box.click()
        await asyncio.sleep(0.3)
        
        # 清空现有内容
        await self.page.keyboard.press("Control+a")
        await self.page.keyboard.press("Delete")
        await asyncio.sleep(0.2)
        
        # 输入内容
        # 分段输入避免过长
        max_chunk = 500
        for i in range(0, len(content), max_chunk):
            chunk = content[i:i+max_chunk]
            await input_box.type(chunk, delay=5)
        
        await asyncio.sleep(0.5)
        
        # 发送
        await self._click_send(input_box)
        return True
    
    async def _click_send(self, input_box):
        """点击发送按钮"""
        send_selectors = [
            "button[class*='send']",
            "button[aria-label*='发送']",
            "button[title*='发送']",
            "div[class*='send-btn']",
        ]
        for sel in send_selectors:
            try:
                btn = self.page.locator(sel).last
                if await btn.is_visible():
                    await btn.click()
                    return
            except Exception:
                continue
        # 回退到 Enter
        await input_box.press("Enter")
    
    async def get_last_user_message(self) -> Optional[str]:
        """获取最后一条用户消息"""
        try:
            # 查找用户消息
            messages = await self.page.evaluate("""
                () => {
                    const items = document.querySelectorAll('.agent-chat__list__item--user, [class*="user"]');
                    const last = items[items.length - 1];
                    return last ? last.innerText?.trim() : null;
                }
            """)
            return messages
        except Exception as e:
            print(f"[MVP] 获取消息失败: {e}")
            return None
    
    async def intercept_and_execute(self, user_input: str) -> Optional[str]:
        """
        拦截用户输入，如果包含本地指令则执行并返回结果
        返回: 要注入回对话的内容，或 None（无需拦截）
        """
        cmd = self.parser.parse(user_input)
        if not cmd:
            return None
        
        tool_name = cmd["tool"]
        args = cmd["args"]
        
        print(f"[MVP] 检测到本地指令: {tool_name}({args})")
        
        # 执行工具
        if tool_name not in TOOLS:
            return f"[错误] 未知工具: {tool_name}。可用工具: {list(TOOLS.keys())}"
        
        try:
            result = TOOLS[tool_name](*args)
            # 包装成自然对话格式
            return f"[本地执行结果]\n{result}\n\n请基于以上结果继续回答。"
        except Exception as e:
            return f"[错误] 工具执行失败: {e}"


# ============ 主循环 ============

class YuanbaoAgentMVP:
    """元宝全自动代理 MVP"""
    
    def __init__(self):
        self.context: Optional[BrowserContext] = None
        self.page: Optional[Page] = None
        self.controller: Optional[YuanbaoController] = None
        self._last_processed: Optional[str] = None
    
    async def run(self):
        """启动 MVP"""
        print("=" * 60)
        print("  元宝全自动代理 MVP")
        print("=" * 60)
        print("\n使用说明:")
        print("  1. 在元宝对话中输入特殊指令触发本地调用:")
        print("     - @agent:read_file('/path/to/file')")
        print("     - /local list_dir /path/to/dir")
        print("     - ```local\nrun_command ls -la\n```")
        print("  2. 系统将拦截指令，执行本地操作")
        print("  3. 执行结果将注入回对话，元宝基于结果继续")
        print("\n按 Ctrl+C 退出\n")
        
        async with async_playwright() as p:
            # 启动浏览器
            print("[MVP] 启动浏览器...")
            user_data_dir = str(Path.home() / ".yuanbao_agent_profile")
            self.context = await p.chromium.launch_persistent_context(
                user_data_dir=user_data_dir,
                channel="chrome",
                headless=False,
                args=["--start-maximized"],
                no_viewport=True,
            )
            
            self.page = self.context.pages[0] if self.context.pages else await self.context.new_page()
            await self.page.goto(YUANBAO_URL)
            
            print(f"[MVP] 已打开元宝: {YUANBAO_URL}")
            print("[MVP] 请手动登录元宝，然后按回车继续...")
            
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(None, input)
            
            self.controller = YuanbaoController(self.page)
            
            print("[MVP] 开始监听对话...")
            await self._monitor_loop()
    
    async def _monitor_loop(self):
        """监听循环"""
        while True:
            try:
                await asyncio.sleep(2)  # 每 2 秒检查一次
                
                # 获取最后一条用户消息
                last_msg = await self.controller.get_last_user_message()
                if not last_msg:
                    continue
                
                # 避免重复处理
                if last_msg == self._last_processed:
                    continue
                
                # 检查是否包含本地指令
                result = await self.controller.intercept_and_execute(last_msg)
                if result:
                    print(f"[MVP] 准备注入执行结果...")
                    # 等待元宝开始回复（可选）
                    await asyncio.sleep(1)
                    
                    # 注入结果
                    success = await self.controller.inject_message(result)
                    if success:
                        print(f"[MVP] 结果已注入对话")
                        self._last_processed = last_msg
                    else:
                        print(f"[MVP] 注入失败")
                else:
                    # 无指令，只是记录
                    self._last_processed = last_msg
                    
            except KeyboardInterrupt:
                print("\n[MVP] 用户中断")
                break
            except Exception as e:
                print(f"[MVP] 监听错误: {e}")
                await asyncio.sleep(5)


async def main():
    agent = YuanbaoAgentMVP()
    try:
        await agent.run()
    except Exception as e:
        print(f"[MVP] 错误: {e}")
        raise


if __name__ == "__main__":
    asyncio.run(main())