from langchain_community.tools import BraveSearch
from typing import Dict, List, Any, Optional
import os
from langchain_core.language_models import BaseLLM

from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine
from local_deep_research import config
import logging
logger = logging.getLogger(__name__)

class BraveSearchEngine(BaseSearchEngine):
    """Brave search engine implementation with two-phase approach"""

    def __init__(self,
                max_results: int = 10,
                region: str = "US",
                time_period: str = "y",
                safe_search: bool = True,
                search_language: str = "English",
                api_key: Optional[str] = None,
                language_code_mapping: Optional[Dict[str, str]] = None,
                llm: Optional[BaseLLM] = None,
                include_full_content: bool = True,
                max_filtered_results: Optional[int] = None,
                **kwargs):
        """
        Initialize the Brave search engine.
        
        Args:
            max_results: Maximum number of search results
            region: Region code for search results
            time_period: Time period for search results
            safe_search: Whether to enable safe search
            search_language: Language for search results
            api_key: Brave Search API key (can also be set in BRAVE_API_KEY env)
            language_code_mapping: Mapping from language names to codes
            llm: Language model for relevance filtering
            include_full_content: Whether to include full webpage content in results
            max_filtered_results: Maximum number of results to keep after filtering
            **kwargs: Additional parameters (ignored but accepted for compatibility)
        """
        # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
        super().__init__(llm=llm, max_filtered_results=max_filtered_results, max_results=max_results)
        self.include_full_content = include_full_content
        
        # Set up language code mapping
        if language_code_mapping is None:
            language_code_mapping = {
                "english": "en",
                "spanish": "es",
                "chinese": "zh",
                "hindi": "hi",
                "french": "fr",
                "arabic": "ar",
                "bengali": "bn",
                "portuguese": "pt",
                "russian": "ru",
            }
        
        # Get API key
        brave_api_key = api_key or os.getenv("BRAVE_API_KEY")
        if not brave_api_key:
            raise ValueError("BRAVE_API_KEY not found. Please provide api_key or set the BRAVE_API_KEY environment variable.")
        
        # Get language code
        language_code = language_code_mapping.get(search_language.lower(), "en")
        
        # Convert time period format to Brave's format
        brave_time_period = f"p{time_period}"
        
        # Convert safe search to Brave's format
        brave_safe_search = "moderate" if safe_search else "off"
        
        # Initialize Brave Search
        self.engine = BraveSearch.from_api_key(
            api_key=brave_api_key,
            search_kwargs={
                "count": min(20, max_results),
                "country": region.upper(),
                "search_lang": language_code,
                "safesearch": brave_safe_search,
                "freshness": brave_time_period,
            }
        )
        
        # Set user agent for Brave Search
        os.environ["USER_AGENT"] = "Local Deep Research/1.0"
        
        # If full content is requested, initialize FullSearchResults
        if include_full_content:
            # Import FullSearchResults only if needed
            try:
                from local_deep_research.web_search_engines.engines.full_search import FullSearchResults
                self.full_search = FullSearchResults(
                    llm=llm,
                    web_search=self.engine,
                    language=search_language,
                    max_results=max_results,
                    region=region,
                    time=time_period,
                    safesearch=brave_safe_search
                )
            except ImportError:
                logger.warning("Warning: FullSearchResults not available. Full content retrieval disabled.")
                self.include_full_content = False
    
    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information from Brave Search.
        
        Args:
            query: The search query
            
        Returns:
            List of preview dictionaries
        """
        logger.info("Getting search results from Brave Search")
        
        try:
            # Get search results from Brave Search
            raw_results = self.engine.run(query[:400])
            
            # Parse results if they're in string format
            if isinstance(raw_results, str):
                try:
                    import json
                    raw_results = json.loads(raw_results)
                except json.JSONDecodeError:
                    logger.error("Error: Unable to parse BraveSearch response as JSON.")
                    return []
            
            # Format results as previews
            previews = []
            for i, result in enumerate(raw_results):
                preview = {
                    "id": i,  # Use index as ID
                    "title": result.get("title", ""),
                    "link": result.get("link", ""),
                    "snippet": result.get("snippet", ""),
                    "displayed_link": result.get("link", ""),
                    "position": i
                }
                
                # Store full Brave result for later
                preview["_full_result"] = result
                
                previews.append(preview)
            
            # Store the previews for potential full content retrieval
            self._search_results = previews
            
            return previews
            
        except Exception as e:
            logger.error(f"Error getting Brave Search results: {e}")
            return []
    
    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant search results.
        If include_full_content is True and FullSearchResults is available,
        retrieves full webpage content for the results.
        
        Args:
            relevant_items: List of relevant preview dictionaries
            
        Returns:
            List of result dictionaries with full content if requested
        """
        # Check if we should get full content
        if hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY:
            logger.info("Snippet-only mode, skipping full content retrieval")
            
            # Return the relevant items with their full Brave information
            results = []
            for item in relevant_items:
                # Use the full result if available, otherwise use the preview
                if "_full_result" in item:
                    result = item["_full_result"]
                    # Remove temporary field
                    if "_full_result" in result:
                        del result["_full_result"]
                else:
                    result = item
                    
                results.append(result)
                
            return results
            
        # If full content retrieval is enabled
        if self.include_full_content and hasattr(self, 'full_search'):
            logger.info("Retrieving full webpage content")
            
            try:
                # Extract only the links from relevant items
                links = [item.get("link") for item in relevant_items if item.get("link")]
                
                # Use FullSearchResults to get full content
                results_with_content = self.full_search._get_full_content(relevant_items)
                
                return results_with_content
                
            except Exception as e:
                logger.error(f"Error retrieving full content: {e}")
                # Fall back to returning the items without full content
        
        # Return items with their full Brave information
        results = []
        for item in relevant_items:
            # Use the full result if available, otherwise use the preview
            if "_full_result" in item:
                result = item["_full_result"].copy()
                # Remove temporary field
                if "_full_result" in result:
                    del result["_full_result"]
            else:
                result = item.copy()
                if "_full_result" in result:
                    del result["_full_result"]
                
            results.append(result)
            
        return results
    
    def run(self, query: str) -> List[Dict[str, Any]]:
        """
        Execute a search using Brave Search with the two-phase approach.
        
        Args:
            query: The search query
            
        Returns:
            List of search results
        """
        logger.info("---Execute a search using Brave Search---")
        
        # Use the implementation from the parent class which handles all phases
        results = super().run(query)
        
        # Clean up
        if hasattr(self, '_search_results'):
            del self._search_results
            
        return results
