from smolagents import DuckDuckGoSearchTool, Tool, tool
import wikipediaapi
import os
import requests

class WikipediaSearchTool(Tool):
    name = "wikipedia_search"
    description = "查找英文维基百科的页面简介，输入应为一个词或短词组"
    inputs = {
        "query": {"type": "string", "description": "维基百科搜索关键词，例如人名/专名"}
    }
    output_type = "string"
    
    def __init__(self, lang="en"):
        super().__init__()
        self.wiki = wikipediaapi.Wikipedia(language=lang, user_agent="celum")
    
    def forward(self, query: str):
        page = self.wiki.page(query)
        if not page.exists():
            return "No Wikipedia page found."
        return page.summary[:1000]

@tool
def tavily_search(query: str) -> str:
    """
    Search the web using Tavily API
    Args:
        query: The search query
    Returns:
        Search results as formatted text
    """
    api_key = os.getenv("TAVILY_API_KEY")  
    if not api_key:
        return "Tavily API key not found"
    
    url = "https://api.tavily.com/search"
    payload = {
        "api_key": api_key,
        "query": query,
        "search_depth": "basic",
        "include_answer": True,
        "include_domains": [],
        "exclude_domains": [],
        "max_results": 5
    }
    
    try:
        response = requests.post(url, json=payload, timeout=10)
        response.raise_for_status()
        data = response.json()
        
        results = []
        if data.get("answer"):
            results.append(f"Quick Answer: {data['answer']}")
        
        for result in data.get("results", [])[:3]:
            results.append(f"Title: {result.get('title', 'N/A')}")
            results.append(f"URL: {result.get('url', 'N/A')}")
            results.append(f"Content: {result.get('content', 'N/A')[:200]}...")
            results.append("---")
        
        return "\n".join(results)
    except Exception as e:
        return f"Tavily search error: {str(e)}"

# 添加媒体分析工具
@tool
def analyze_media_file(file_path: str) -> str:
    """
    分析媒体文件并返回基本信息
    Args:
        file_path: 文件路径
    Returns:
        文件信息和分析结果
    """
    if not os.path.exists(file_path):
        return f"File not found: {file_path}"
    
    try:
        file_size = os.path.getsize(file_path)
        file_ext = os.path.splitext(file_path)[1].lower()
        
        result = f"File: {os.path.basename(file_path)}\n"
        result += f"Size: {file_size} bytes\n"
        result += f"Extension: {file_ext}\n"
        
        # 根据文件类型进行不同的分析
        if file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp']:
            try:
                from PIL import Image
                with Image.open(file_path) as img:
                    result += f"Image dimensions: {img.size[0]}x{img.size[1]}\n"
                    result += f"Image mode: {img.mode}\n"
                    result += f"Image format: {img.format}\n"
                    
                    # 检测是否可能是棋盘
                    if abs(img.size[0] - img.size[1]) < 50:
                        result += "Note: Square aspect ratio - possibly a chess board\n"
            except Exception as e:
                result += f"Image analysis error: {str(e)}\n"
                
        elif file_ext in ['.mp4', '.avi', '.mov', '.mkv']:
            result += "Video file detected\n"
        elif file_ext in ['.mp3', '.wav', '.m4a', '.flac']:
            result += "Audio file detected\n"
        elif file_ext in ['.pdf']:
            result += "PDF file detected\n"
        elif file_ext in ['.txt', '.csv', '.json']:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                result += f"Text length: {len(content)} characters\n"
                result += f"Line count: {len(content.splitlines())}\n"
                preview = content[:200].replace('\n', ' ')
                result += f"Content preview: {preview}...\n"
            except Exception as e:
                result += f"Text analysis error: {str(e)}\n"
        
        return result
            
    except Exception as e:
        return f"Error analyzing file: {str(e)}"

@tool
def analyze_chess_position(file_path: str) -> str:
    """
    分析棋盘位置，尝试识别棋局
    Args:
        file_path: 棋盘图像路径
    Returns:
        棋局分析结果
    """
    if not os.path.exists(file_path):
        return f"File not found: {file_path}"
    
    try:
        file_size = os.path.getsize(file_path)
        result = f"Chess board image found: {os.path.basename(file_path)} ({file_size} bytes)\n"
        
        # 由于PIL被限制，我们提供一些通用的棋局分析提示
        result += """
Based on typical chess endgame patterns, common winning moves include:
- Queen checks that lead to checkmate (Qd1+, Qe1+, etc.)
- Rook moves that create back-rank threats
- Knight forks that win material
- Bishop moves that control key squares

For endgame puzzles, look for:
1. Checks that force the king to a bad square
2. Pins and skewers
3. Back-rank weaknesses
4. Knight forks
5. Zugzwang positions

Without being able to process the image directly, I recommend using chess knowledge and pattern recognition.
"""
        
        return result
    except Exception as e:
        return f"Error analyzing chess position: {str(e)}"

# 更新工具列表
my_tool_list = [
    WikipediaSearchTool(),
    DuckDuckGoSearchTool(),
    tavily_search,
    analyze_media_file,
    analyze_chess_position,  
]