from smolagents import WebSearchTool, Tool
import requests
import time
from typing import List, Dict, Any, Optional
from duckduckgo_search import DDGS


class MySearchTool(WebSearchTool):
    name = "my_search_tool"
    description = "It receives a question and returns a list of search results from various sources. It uses DuckDuckGo as the primary source"
    inputs = {
        "query": {
            "type": "string",
            "description": "The question or topic to search in internet.",
        },
    }
    output_type = "string"

    def __init__(self, max_results=5):
        super().__init__()
        self.max_results = max_results

    def search_duckduckgo(self, query: str) -> list:
        """
        Free search using DuckDuckGo Instant Answer API

        Args:
            query: Search query
            max_results: Maximum number of results

        Returns:
            List of search results
        """
        try:
            # DuckDuckGo Instant Answer API (free)
            url = "https://api.duckduckgo.com/"
            params = {
                "q": query,
                "format": "json",
                "no_html": "1",
                "skip_disambig": "1",
            }

            response = requests.get(url, params=params, timeout=10)
            response.raise_for_status()
            data = response.json()
            time.sleep(2)  # 2-second delay

            results = []

            # Process abstract
            if data.get("Abstract"):
                results.append(
                    {
                        "title": data.get("Heading", "DuckDuckGo Result"),
                        "url": data.get("AbstractURL", ""),
                        "content": data.get("Abstract", ""),
                        "source": "DuckDuckGo",
                    }
                )

            # Process related topics
            for topic in data.get("RelatedTopics", [])[
                : self.max_results - len(results)
            ]:
                if isinstance(topic, dict) and "Text" in topic:
                    results.append(
                        {
                            "title": topic.get("Text", "")[:100],
                            "url": topic.get("FirstURL", ""),
                            "content": topic.get("Text", ""),
                            "source": "DuckDuckGo",
                        }
                    )

            return results[: self.max_results]

        except Exception as e:
            print(f"DuckDuckGo search failed: {str(e)}")
            return []


class DuckDuckGoSearchTool(Tool):
    name = "web_search"
    description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
    inputs = {
        "query": {"type": "string", "description": "The search query to perform."}
    }
    output_type = "string"

    def __init__(self, max_results=10, **kwargs):
        super().__init__()
        self.max_results = max_results
        self.ddgs = DDGS(**kwargs)

    def forward(self, query: str) -> str:
        results = self.ddgs.text(query, max_results=self.max_results)
        if len(results) == 0:
            raise Exception("No results found! Try a less restrictive/shorter query.")
        postprocessed_results = [
            f"[{result['title']}]({result['href']})\n{result['body']}"
            for result in results
        ]
        return "## Search Results\n\n" + "\n\n".join(postprocessed_results)