"""
Search tool implementation for simulating internet searches
"""

import asyncio
import time
import random
from typing import Dict, List, Any, Optional
import os

from langchain_text_splitters import RecursiveCharacterTextSplitter
import requests

from meta_researcher.tool.tools.search_engine import Retriever, SearchEngine
from meta_researcher.tool.tools.search_engine.base_search import SearchConfig
from meta_researcher.tool.tools.search_engine.website_crawler.simple_crawler import SimpleContentFetch
from meta_researcher.tool.base import BaseTool

# from txtai.embeddings import Embeddings
import json
max_step_num = 5

class PlanTool(BaseTool):
    name = "Plan"
    description = f"""Task Planning Tool. It can formulate a detailed task plan for the input problem, with specific planning principles as follows:
    1. For any problem, you must first conduct task planning for the problem in your first-round reply.
    2. A successful task plan needs to decompose the research problem into multiple sub-problems, and take into account both the breadth of the problem (covering basic aspects) and the depth (detailed information of each aspect).
    3. The task plan ensures that the most critical perspectives are covered within {max_step_num} steps.
    4. The task plan needs to be output in the following format:
        Step 1: Collect information about xxx
            - Collect information about xxx
            ...
        Step 2: Collect information about xxx
            - Collect information about xxx
            ...
        ...
    """
    parameters = {
        "type": "object",
        "properties": {
            "plan": {
                "type": "string", 
                "description": "The specific content of task planning must fully cover the breadth and depth of the original problem.", 
                "examples": ["""Step 1: Collect information about Xiaomi YU7
                                    - Collect configuration information of Xiaomi YU7
                                    - Collect price information of Xiaomi YU7
                                Step 2: Collect information about user reviews of Xiaomi YU7
                                    - Collect real user review information of Xiaomi YU7"""],
            }
        },
        "required": ["plan"]
    }

    def __init__(self):
        super().__init__()
        print("[DEBUG] EMBEDDINGS LOADING")
    
    def execute(self, args: Dict) -> str:
        """
        Execute search query
        
        Args:
            args: Tool parameters, containing:
                - "query": search query string
                - "limit": optional int to limit number of results
            
        Returns:
            Formatted search results
        """
        pass
    
    def batch_execute(self, args_list: List[Dict]) -> List[str]:
        try:
            plans = [x["plan"] for x in args_list]
            return [{"content": "<plan>\n" + str(content) + "\n</plan>", "success": True} for content in plans]
        except Exception as e:
            return [{"content": str(e), "success": False} for _ in args_list]
    
    def calculate_reward(self, args: Dict, result: str) -> float:
        """
        Calculate reward for search action
        
        Args:
            args: Tool parameters
            result: Tool execution result
            
        Returns:
            Reward value
        """
        # valid tool call
        if "results" in result:
            return 0.1
        else:
            return 0.0