"""
Agent 模块 - 封装 browser-use Agent
"""

import asyncio
from typing import Optional, Type
from pydantic import BaseModel

from .config import ScraperConfig, LLMConfig, BrowserConfig


class NSFCAgent:
    """国自然爬虫 Agent 封装"""
    
    def __init__(self, config: Optional[ScraperConfig] = None):
        self.config = config or ScraperConfig.from_env()
        self._agent = None
        self._browser_profile = None
        self._llm = None
    
    @property
    def llm(self):
        """获取 LLM 实例（懒加载）"""
        if self._llm is None:
            self._llm = self._create_llm()
        return self._llm
    
    def _create_llm(self):
        """创建 LLM 实例"""
        llm_config = self.config.llm
        
        if llm_config.provider == "openai":
            from browser_use import ChatOpenAI
            return ChatOpenAI(
                model=llm_config.model,
                api_key=llm_config.api_key,
                base_url=llm_config.base_url,
                temperature=llm_config.temperature,
            )
        elif llm_config.provider == "anthropic":
            from browser_use import ChatAnthropic
            return ChatAnthropic(
                model=llm_config.model,
                api_key=llm_config.api_key,
            )
        elif llm_config.provider == "deepseek":
            from browser_use.llm.deepseek.chat import ChatDeepSeek
            return ChatDeepSeek(
                model=llm_config.model,
                api_key=llm_config.api_key,
            )
        elif llm_config.provider == "ollama":
            from browser_use import ChatOllama
            return ChatOllama(model=llm_config.model)
        else:
            raise ValueError(f"不支持的 LLM 提供商: {llm_config.provider}")
    
    def _create_browser_profile(self):
        """创建浏览器配置"""
        from browser_use import BrowserProfile
        
        browser_config = self.config.browser
        
        # 构建配置参数
        profile_kwargs = {
            "headless": browser_config.headless,
            "wait_between_actions": browser_config.wait_between_actions,
            "minimum_wait_page_load_time": browser_config.minimum_wait_page_load_time,
            "wait_for_network_idle_page_load_time": browser_config.wait_for_network_idle_page_load_time,
            "highlight_elements": browser_config.highlight_elements,
            "keep_alive": browser_config.keep_alive,
        }
        
        # 如果指定了 user_data_dir，使用它；否则使用内置 Chromium
        if browser_config.user_data_dir:
            profile_kwargs["user_data_dir"] = browser_config.user_data_dir
        else:
            # 使用 Playwright 内置 Chromium，避免 Chrome 版本兼容问题
            profile_kwargs["channel"] = "chromium"
        
        return BrowserProfile(**profile_kwargs)
    
    async def run(
        self,
        task: str,
        output_model: Optional[Type[BaseModel]] = None,
        max_steps: Optional[int] = None,
    ) -> dict:
        """
        执行爬取任务
        
        Args:
            task: 任务描述（自然语言）
            output_model: Pydantic 模型，用于结构化输出
            max_steps: 最大步数
            
        Returns:
            爬取结果
        """
        from browser_use import Agent
        
        if self._llm is None:
            self._llm = self._create_llm()
        
        if self._browser_profile is None:
            self._browser_profile = self._create_browser_profile()
        
        agent = Agent(
            task=task,
            llm=self._llm,
            browser_profile=self._browser_profile,
            output_model_schema=output_model,
        )
        
        steps = max_steps or self.config.max_steps
        history = await agent.run(max_steps=steps)
        
        result = {
            "success": True,
            "task": task,
            "raw_result": history.final_result(),
            "steps_taken": len(history.history) if history.history else 0,
        }
        
        # 如果有结构化输出模型，解析结果
        if output_model and result["raw_result"]:
            try:
                parsed = output_model.model_validate_json(result["raw_result"])
                result["parsed"] = parsed.model_dump()
            except Exception as e:
                result["parse_error"] = str(e)
        
        return result
    
    async def run_interactive(self, task: str) -> str:
        """
        交互式运行任务，返回原始结果字符串
        """
        result = await self.run(task)
        return result.get("raw_result", "")
