from typing import Dict, List, Any, Optional
import os
import requests
import time
import random
import logging
from requests.exceptions import RequestException
from urllib.parse import quote_plus
from langchain_core.language_models import BaseLLM

from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class GooglePSESearchEngine(BaseSearchEngine):
    """Google Programmable Search Engine implementation"""

    def __init__(self,
                max_results: int = 10,
                region: str = "us",
                safe_search: bool = True,
                search_language: str = "English",
                api_key: Optional[str] = None,
                search_engine_id: Optional[str] = None,
                llm: Optional[BaseLLM] = None,
                include_full_content: bool = False,
                max_filtered_results: Optional[int] = None,
                max_retries: int = 3,
                retry_delay: float = 2.0,
                **kwargs):
        """
        Initialize the Google Programmable Search Engine.
        
        Args:
            max_results: Maximum number of search results
            region: Region code for search results
            safe_search: Whether to enable safe search
            search_language: Language for search results
            api_key: Google API key (can also be set in GOOGLE_PSE_API_KEY env)
            search_engine_id: Google CSE ID (can also be set in GOOGLE_PSE_ENGINE_ID env)
            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
            max_retries: Maximum number of retry attempts for API requests
            retry_delay: Base delay in seconds between retry attempts
            **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
        
        # Retry configuration
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        
        # Rate limiting - keep track of last request time
        self.last_request_time = 0
        self.min_request_interval = 0.5  # Minimum time between requests in seconds
        
        # Language code mapping
        language_code_mapping = {
            "english": "en",
            "spanish": "es",
            "french": "fr",
            "german": "de",
            "italian": "it",
            "japanese": "ja",
            "korean": "ko",
            "portuguese": "pt",
            "russian": "ru",
            "chinese": "zh-CN"
        }
        
        # Get language code
        search_language = search_language.lower()
        self.language = language_code_mapping.get(search_language, "en")
        
        # Safe search setting
        self.safe = "active" if safe_search else "off"
        
        # Region/Country setting
        self.region = region
        
        # API key and Search Engine ID
        self.api_key = api_key or os.getenv("GOOGLE_PSE_API_KEY")
        self.search_engine_id = search_engine_id or os.getenv("GOOGLE_PSE_ENGINE_ID")
        
        if not self.api_key:
            raise ValueError("Google API key is required. Set it in the GOOGLE_PSE_API_KEY environment variable.")
        if not self.search_engine_id:
            raise ValueError("Google Search Engine ID is required. Set it in the GOOGLE_PSE_ENGINE_ID environment variable.")
            
        # Validate connection and credentials
        self._validate_connection()
    
    def _validate_connection(self):
        """Test the connection to ensure API key and Search Engine ID are valid"""
        try:
            # Make a minimal test query
            response = self._make_request("test")
            
            # Check if we got a valid response
            if response.get("error"):
                error_msg = response["error"].get("message", "Unknown error")
                raise ValueError(f"Google PSE API error: {error_msg}")
                
            # If we get here, the connection is valid
            logger.info("Google PSE connection validated successfully")
            return True
            
        except Exception as e:
            # Log the error and re-raise
            logger.error(f"Error validating Google PSE connection: {str(e)}")
            raise
    
    def _respect_rate_limit(self):
        """Ensure we don't exceed rate limits by adding appropriate delay between requests"""
        current_time = time.time()
        elapsed = current_time - self.last_request_time
        
        # If we've made a request recently, wait until the minimum interval has passed
        if elapsed < self.min_request_interval:
            sleep_time = self.min_request_interval - elapsed
            logger.debug(f"Rate limiting: sleeping for {sleep_time:.2f}s")
            time.sleep(sleep_time)
        
        # Update the last request time
        self.last_request_time = time.time()

    def _make_request(self, query: str, start_index: int = 1) -> Dict:
        """
        Make a request to the Google PSE API with retry logic and rate limiting
        
        Args:
            query: Search query string
            start_index: Starting index for pagination
            
        Returns:
            JSON response from the API
            
        Raises:
            RequestException: If all retry attempts fail
        """
        # Base URL for the API
        url = "https://www.googleapis.com/customsearch/v1"
        
        # Parameters for the request
        params = {
            "key": self.api_key,
            "cx": self.search_engine_id,
            "q": query,
            "num": min(10, self.max_results),  # Max 10 per request
            "start": start_index,
            "safe": self.safe,
            "lr": f"lang_{self.language}",
            "gl": self.region
        }
        
        # Implement retry logic with exponential backoff
        attempt = 0
        last_exception = None
        
        while attempt < self.max_retries:
            try:
                # Respect rate limits
                self._respect_rate_limit()
                
                # Add jitter to retries after the first attempt
                if attempt > 0:
                    jitter = random.uniform(0.5, 1.5)
                    sleep_time = self.retry_delay * (2 ** (attempt - 1)) * jitter
                    logger.info(f"Retry attempt {attempt+1}/{self.max_retries} for query '{query}'. Waiting {sleep_time:.2f}s")
                    time.sleep(sleep_time)
                
                # Make the request
                logger.debug(f"Making request to Google PSE API: {query} (start_index={start_index})")
                response = requests.get(url, params=params, timeout=10)
                
                # Check for HTTP errors
                response.raise_for_status()
                
                # Return the JSON response
                return response.json()
                
            except RequestException as e:
                logger.warning(f"Request error on attempt {attempt+1}/{self.max_retries}: {str(e)}")
                last_exception = e
            except Exception as e:
                logger.warning(f"Error on attempt {attempt+1}/{self.max_retries}: {str(e)}")
                last_exception = e
            
            attempt += 1
        
        # If we get here, all retries failed
        error_msg = f"Failed to get response from Google PSE API after {self.max_retries} attempts"
        logger.error(error_msg)
        
        if last_exception:
            raise RequestException(f"{error_msg}: {str(last_exception)}")
        else:
            raise RequestException(error_msg)
        
    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """Get search result previews/snippets"""
        results = []
        
        # Google PSE API returns a maximum of 10 results per request
        # We may need to make multiple requests to get the desired number
        start_index = 1
        total_results = 0
        
        while total_results < self.max_results:
            try:
                response = self._make_request(query, start_index)
                
                # Break if no items
                if "items" not in response:
                    break
                    
                items = response.get("items", [])
                
                # Process each result
                for item in items:
                    title = item.get("title", "")
                    snippet = item.get("snippet", "")
                    url = item.get("link", "")
                    
                    # Skip results without URL
                    if not url:
                        continue
                        
                    results.append({
                        "title": title,
                        "snippet": snippet,
                        "url": url,
                        "source": "Google Programmable Search"
                    })
                    
                    total_results += 1
                    if total_results >= self.max_results:
                        break
                        
                # Check if there are more results
                if not items or total_results >= self.max_results:
                    break
                    
                # Update start index for next request
                start_index += len(items)
                
                # Add a small delay between multiple requests to be respectful of the API
                if total_results < self.max_results:
                    time.sleep(self.min_request_interval)
                
            except Exception as e:
                logger.error(f"Error getting search results: {str(e)}")
                break
                
        logger.info(f"Retrieved {len(results)} search results for query: '{query}'")
        return results
        
    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Get full content for search results"""
        # Use the BaseSearchEngine implementation
        return super()._get_full_content(relevant_items)
        
    def run(self, query: str) -> List[Dict[str, Any]]:
        """Run the search engine to get results for a query"""
        # Get search result previews/snippets
        search_results = self._get_previews(query)
        
        # Filter for relevance if we have an LLM and max_filtered_results
        if self.llm and self.max_filtered_results:
            search_results = self._filter_for_relevance(query, search_results)
            
        # Get full content if needed
        if self.include_full_content:
            search_results = self._get_full_content(search_results)
            
        return search_results 