import os
import re
import json
import asyncio
from typing import Union, List, Any


async def search_serper_main(keywords: List[str]) -> List[str]:
    """
    Mocks a web search by returning predefined results based on input keywords.
    This function simulates the behavior of a Serper API call.
    """
    print(f"[MockSerper] Performing mock search for keywords: {keywords}")
    mock_results = []
    
    # Simple keyword-based response simulation
    if "Transformer architectures" in keywords or "neural network types" in keywords:
        mock_results.append("Found: Transformer architectures are a type of neural network. (mock data from Serper)")
    if "Large Language Models" in keywords or "LLM capabilities" in keywords:
        mock_results.append("Found: Large Language Models (LLMs) are deep learning models designed to understand and generate human-like text. (mock data from Serper)")
    if "current weather London" in keywords or "London temperature" in keywords:
        mock_results.append("Found: Current weather in London: 15°C, partly cloudy. (mock data from Serper)")
    
    if not mock_results:
        mock_results.append(f"Found: No specific mock results for keywords: {', '.join(keywords)}. (generic mock data from Serper)")
    
    # Simulate network latency
    await asyncio.sleep(0.5) 
    return mock_results

# It's good practice to import OpenAI client directly if using it
from openai import AsyncOpenAI # Using AsyncOpenAI for async methods

class Serper_Searcher(Tool_node):
    """
    A tool that performs web searches using Serper.dev by first generating search keywords.
    It leverages an LLM (DeepSeek in this case) to extract relevant keywords from a user's task/question
    and then uses those keywords to perform an actual web search via Serper.dev.
    """
    def __init__(self):
        super().__init__(
            name="Serper_Searcher",
            description="Performs web searches by generating keywords and querying Serper.dev."
        )
        # Load API keys and base URLs from environment variables for security and flexibility.
        # Fallback to hardcoded values for demonstration, but these should be replaced in production.
        self.deepseek_api_key: str = os.getenv("DEEPSEEK_API_KEY", "YOUR_DEEPSEEK_API_KEY_HERE")
        self.deepseek_base_url: str = os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com")

    async def _execute_tool(self, task: str) -> str:
        """
        Executes a web search. It first generates search keywords using an LLM based on the task,
        then uses these keywords to query Serper.dev, and finally returns a summary of the results.

        Args:
            task (str): The user's question or task for which to perform a web search.

        Returns:
            str: A formatted string containing the search results, or an error message.
        """
        if not self.deepseek_api_key or self.deepseek_api_key == "YOUR_DEEPSEEK_API_KEY_HERE":
             return "Error: DeepSeek API key not configured. Please set the 'DEEPSEEK_API_KEY' environment variable."

        # Generate search prompt for keyword extraction
        keyword_generation_prompt = self._get_websearch_prompt(question=task)

        try:
            # Initialize OpenAI client (compatible with DeepSeek's API)
            client = AsyncOpenAI(api_key=self.deepseek_api_key, base_url=self.deepseek_base_url)
            
            # Request keyword generation from the LLM
            response = await client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": "You are an expert in search keyword extraction. Your goal is to identify and list the most pertinent search terms for a given question."},
                    {"role": "user", "content": keyword_generation_prompt},
                ],
                stream=False
            )
            
            response_str = response.choices[0].message.content
            keywords_list = self._str_to_list(response_str)

            print("[SerperSearcher] Generated Keywords:", keywords_list)

            if not keywords_list:
                return f"Error: Failed to generate valid search keywords for task '{task}'. LLM output: {response_str}"

            # Execute web search using the generated keywords
            serper_results = await search_serper_main(keywords_list)

            print("[SerperSearcher] Web Search Results:", serper_results)

            if serper_results:
                # Combine search results into a single string for return
                serper_summary = "\n".join(serper_results)
                return f"Web search for '{task}' yielded results:\n{serper_summary}"
            else:
                return f"No relevant web search results found for task: '{task}' with keywords: {', '.join(keywords_list)}."

        except Exception as e:
            return f"Error during web search execution: {type(e).__name__}: {e}"

    def _format_result_to_natural_language(self, raw_result: Any, task_description: str) -> str:
        """
        Formats the raw search results into a natural language description.

        Args:
            raw_result (Any): The raw output from `_execute_tool`, typically a string containing search results.
            task_description (str): The original task description provided to the tool.

        Returns:
            str: A human-readable summary of the search results.
        """
        if raw_result.startswith("Error"):
            return f"The web search for '{task_description}' failed: {raw_result}"
        
        # If _execute_tool already returns a formatted string, we can directly use it.
        # This method would parse raw JSON if _execute_tool returned raw data.
        return f"Web search for '{task_description}' completed with the following information:\n{raw_result}"

    def _get_websearch_prompt(self, question: str) -> str:
        """
        Generates a prompt for an LLM to extract search keywords from a given question.

        Args:
            question (str): The user's original question or task.

        Returns:
            str: The formatted prompt for keyword generation.
        """
        return (
            f"## Original Question: \n---\n{question}\n---\n\n"
            f"## Search Keyword Generation Objective:\n"
            f"Generate exactly three specific search keywords or phrases to find information relevant to the question. "
            f"Each keyword should focus on key concepts, terms, or formulas from the question. "
            f"Format the output as a comma-separated list of exactly three keywords.\n\n"
            f"## Examples:\n"
            f"Q: A rectangular garden has a length that is 5 meters longer than its width. If the area of the garden is 84 square meters, what are the dimensions of the garden?\n"
            f"A: 'quadratic equation, solve polynomial, area of rectangle'\n\n"
            f"Q: If a car travels at a constant speed of 60 km/h for 2.5 hours, how far does it travel?\n"
            f"A: 'distance formula, speed time distance, constant speed'\n\n"
            f"Output only the three keywords in the format: 'keyword1, keyword2, keyword3', with no other text or explanation."
        )

    def _str_to_list(self, keyword_string: str) -> List[str]:
        """
        Converts a comma-separated string of keywords into a list of strings.
        It expects exactly three keywords.

        Args:
            keyword_string (str): A string containing comma-separated keywords (e.g., "keyword1, keyword2, keyword3").

        Returns:
            List[str]: A list of cleaned keywords. Returns an empty list if not exactly three valid keywords are found.
        """
        if not isinstance(keyword_string, str):
            return []
            
        # Split by comma and strip whitespace from each part
        keywords = [k.strip() for k in keyword_string.split(",")]
        
        # Filter out any empty strings that might result from extra commas (e.g., "a,,b")
        cleaned_keywords = [k for k in keywords if k] 
        
        # Ensure exactly three keywords are present as per the prompt's instruction
        if len(cleaned_keywords) == 3:
            return cleaned_keywords
        else:
            print(f"[SerperSearcher] Warning: Expected 3 keywords, but received {len(cleaned_keywords)} from LLM: '{keyword_string}'")
            return []