import os
from typing import Dict, Any
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
from langchain_core.tools import tool
import logging

logger = logging.getLogger(__name__)

@tool
def wiki_search(query: str) -> Dict[str, str]:
    """Search Wikipedia for a query and return maximum 2 results.
    Args:
        query: The search query."""
    try:
        logger.info(f"Searching Wikipedia for: {query}")
        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
        if not search_docs:
            logger.warning("No Wikipedia results found")
            return {"wiki_results": "No results found"}
            
        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.metadata.get("source", "")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
                for doc in search_docs
            ])
        logger.info(f"Found {len(search_docs)} Wikipedia results")
        return {"wiki_results": formatted_search_docs}
    except Exception as e:
        logger.error(f"Error searching Wikipedia: {str(e)}")
        return {"wiki_results": f"Error searching Wikipedia: {str(e)}"}

@tool
def web_search(query: str) -> Dict[str, str]:
    """Search Tavily for a query and return maximum 3 results.
    Args:
        query: The search query."""
    try:
        logger.info(f"Searching web for: {query}")
        search = TavilySearchResults(max_results=3)
        search_docs = search.invoke({"query": query})
        
        if not search_docs:
            logger.warning("No web results found")
            return {"web_results": "No results found"}
            
        if isinstance(search_docs, list):
            formatted_search_docs = "\n\n---\n\n".join(
                [
                    f'<Document source="{doc.get("source", "")}" page="{doc.get("page", "")}"/>\n{doc.get("content", "")}\n</Document>'
                    for doc in search_docs
                ])
            logger.info(f"Found {len(search_docs)} web results")
            return {"web_results": formatted_search_docs}
        logger.warning(f"Unexpected response format from Tavily: {type(search_docs)}")
        return {"web_results": f"Error: Unexpected response format from Tavily"}
    except Exception as e:
        logger.error(f"Error searching web: {str(e)}")
        return {"web_results": f"Error searching web: {str(e)}"}

@tool
def arxiv_search(query: str) -> Dict[str, str]:
    """Search Arxiv for a query and return maximum 3 results.
    Args:
        query: The search query."""
    try:
        logger.info(f"Searching Arxiv for: {query}")
        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
        if not search_docs:
            logger.warning("No Arxiv results found")
            return {"arxiv_results": "No results found"}
            
        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.metadata.get("source", "")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
                for doc in search_docs
            ])
        logger.info(f"Found {len(search_docs)} Arxiv results")
        return {"arxiv_results": formatted_search_docs}
    except Exception as e:
        logger.error(f"Error searching Arxiv: {str(e)}")
        return {"arxiv_results": f"Error searching Arxiv: {str(e)}"}

@tool
def wiki_api_search(query: str) -> Dict[str, str]:
    """Search Wikipedia using API wrapper for better results.
    Args:
        query: The search query."""
    try:
        logger.info(f"Searching Wikipedia API for: {query}")
        wikipedia = WikipediaAPIWrapper(top_k_results=3, doc_content_chars_max=4000)
        results = wikipedia.run(query)
        
        if not results or results.strip() == "No good Wikipedia Search Result was found":
            logger.warning("No Wikipedia API results found")
            return {"wiki_api_results": "No results found"}
            
        logger.info(f"Found Wikipedia API results")
        return {"wiki_api_results": results}
    except Exception as e:
        logger.error(f"Error searching Wikipedia API: {str(e)}")
        return {"wiki_api_results": f"Error searching Wikipedia API: {str(e)}"}

# List of all search tools
SEARCH_TOOLS = [wiki_search, web_search, arxiv_search, wiki_api_search]

class SearchTools:
    """Wrapper class for search tools to provide a unified interface"""
    
    def __init__(self):
        """Initialize search tools"""
        pass
    
    def search_wikipedia(self, query: str) -> str:
        """Search Wikipedia and return formatted results"""
        result = wiki_search(query)
        return result.get("wiki_results", "")
    
    def search_wikipedia_api(self, query: str) -> str:
        """Search Wikipedia using API wrapper and return formatted results"""
        result = wiki_api_search(query)
        return result.get("wiki_api_results", "")
    
    def search_web(self, query: str) -> str:
        """Search web and return formatted results"""
        result = web_search(query)
        return result.get("web_results", "")
    
    def search_arxiv(self, query: str) -> str:
        """Search Arxiv and return formatted results"""
        result = arxiv_search(query)
        return result.get("arxiv_results", "")