import requests
import logging
from typing import Dict, List, Any, Optional, Tuple
import os
from datetime import datetime, timedelta
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
from local_deep_research.utilties.search_utilities import remove_think_tags

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

class GuardianSearchEngine(BaseSearchEngine):
    """Enhanced Guardian API search engine implementation with LLM query optimization"""
    
    def __init__(self, 
                max_results: int = 10, 
                api_key: Optional[str] = None,
                from_date: Optional[str] = None,
                to_date: Optional[str] = None,
                section: Optional[str] = None,
                order_by: str = "relevance",
                llm: Optional[BaseLLM] = None,
                max_filtered_results: Optional[int] = None,
                optimize_queries: bool = True,
                adaptive_search: bool = True):
        """
        Initialize The Guardian search engine with enhanced features.
        
        Args:
            max_results: Maximum number of search results
            api_key: The Guardian API key (can also be set in GUARDIAN_API_KEY env)
            from_date: Start date for search (YYYY-MM-DD format, default 1 month ago)
            to_date: End date for search (YYYY-MM-DD format, default today)
            section: Filter by section (e.g., "politics", "technology", "sport")
            order_by: Sort order ("relevance", "newest", "oldest")
            llm: Language model for relevance filtering and query optimization
            max_filtered_results: Maximum number of results to keep after filtering
            optimize_queries: Whether to optimize queries using LLM
            adaptive_search: Whether to use adaptive search (adjusting date ranges)
        """
        # 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.api_key = api_key or os.getenv("GUARDIAN_API_KEY")
        self.optimize_queries = optimize_queries
        self.adaptive_search = adaptive_search
        
        if not self.api_key:
            raise ValueError("Guardian API key not found. Please provide api_key or set the GUARDIAN_API_KEY environment variable.")
        
        # Set date ranges if not provided
        if not from_date:
            # Default to one month ago
            one_month_ago = datetime.now() - timedelta(days=30)
            self.from_date = one_month_ago.strftime("%Y-%m-%d")
        else:
            self.from_date = from_date
            
        if not to_date:
            # Default to today
            self.to_date = datetime.now().strftime("%Y-%m-%d")
        else:
            self.to_date = to_date
            
        self.section = section
        self.order_by = order_by
        self._original_date_params = {
            "from_date": self.from_date,
            "to_date": self.to_date
        }
        
        # API base URL
        self.api_url = "https://content.guardianapis.com/search"
    
    def _optimize_query_for_guardian(self, query: str) -> str:
        """
        Optimize a natural language query for Guardian search.
        Uses LLM to transform questions into effective news search queries.
        
        Args:
            query: Natural language query
            
        Returns:
            Optimized query string for Guardian
        """
        # Handle extremely long queries by truncating first
        if len(query) > 150:
            simple_query = " ".join(query.split()[:10])
            logger.info(f"Query too long ({len(query)} chars), truncating to: {simple_query}")
            query = simple_query
        
        if not self.llm or not self.optimize_queries:
            # Return original query if no LLM available or optimization disabled
            return query
            
        try:
            # Prompt for query optimization
            prompt = f"""Transform this natural language question into a very short Guardian news search query.

Original query: "{query}"

CRITICAL RULES:
1. ONLY RETURN THE EXACT SEARCH QUERY - NO EXPLANATIONS, NO COMMENTS
2. Keep it EXTREMELY BRIEF - MAXIMUM 3-4 words total
3. Focus only on the main topic/person/event
4. Include proper names when relevant
5. Remove ALL unnecessary words
6. DO NOT use Boolean operators (no AND/OR)
7. DO NOT use quotes

EXAMPLE CONVERSIONS:
✓ "What's the impact of rising interest rates on UK housing market?" → "UK housing rates"
✓ "Latest developments in the Ukraine-Russia peace negotiations" → "Ukraine Russia negotiations"
✓ "How are tech companies responding to AI regulation?" → "tech AI regulation"
✓ "What is Donald Trump's current political activity?" → "Trump political activity"

Return ONLY the extremely brief search query.
"""
            
            # Get response from LLM
            response = self.llm.invoke(prompt)
            optimized_query = remove_think_tags(response.content).strip()
            
            # Clean up the query - remove any explanations
            lines = optimized_query.split('\n')
            for line in lines:
                line = line.strip()
                if line and not line.lower().startswith(('here', 'i would', 'the best', 'this query')):
                    optimized_query = line
                    break
            
            # Remove any quotes that wrap the entire query
            if optimized_query.startswith('"') and optimized_query.endswith('"') and optimized_query.count('"') == 2:
                optimized_query = optimized_query[1:-1]
            
            logger.info(f"Original query: '{query}'")
            logger.info(f"Optimized for Guardian: '{optimized_query}'")
            
            return optimized_query
            
        except Exception as e:
            logger.error(f"Error optimizing query: {e}")
            return query  # Fall back to original query on error
    
    def _adapt_dates_for_query_type(self, query: str) -> None:
        """
        Adapt date range based on query type (historical vs current).
        
        Args:
            query: The search query
        """
        # Fast path - for very short queries, default to recent news
        if len(query.split()) <= 4:
            logger.info("Short query detected, defaulting to recent news")
            # Default to 60 days for short queries
            recent = (datetime.now() - timedelta(days=60)).strftime("%Y-%m-%d")
            self.from_date = recent
            self.order_by = "newest"
            return
        
        if not self.llm or not self.adaptive_search:
            return
            
        try:
            prompt = f"""Is this query asking about HISTORICAL events or CURRENT events?

Query: "{query}"

ONE WORD ANSWER ONLY:
- "HISTORICAL" if about past events (older than 1 year)
- "CURRENT" if about recent events (within past year)
- "UNCLEAR" if can't determine

ONE WORD ONLY:"""
            
            response = self.llm.invoke(prompt)
            answer = remove_think_tags(response.content).strip().upper()
            
            # Reset to original parameters first
            self.from_date = self._original_date_params["from_date"]
            self.to_date = self._original_date_params["to_date"]
            
            if "HISTORICAL" in answer:
                # For historical queries, go back 10 years
                logger.info("Query classified as HISTORICAL - extending search timeframe")
                ten_years_ago = (datetime.now() - timedelta(days=3650)).strftime("%Y-%m-%d")
                self.from_date = ten_years_ago
                
            elif "CURRENT" in answer:
                # For current events, focus on recent content
                logger.info("Query classified as CURRENT - focusing on recent content")
                recent = (datetime.now() - timedelta(days=60)).strftime("%Y-%m-%d")
                self.from_date = recent
                self.order_by = "newest"  # Prioritize newest for current events
                
        except Exception as e:
            logger.error(f"Error adapting dates for query type: {e}")
            # Keep original date parameters on error
    
    def _adaptive_search(self, query: str) -> Tuple[List[Dict[str, Any]], str]:
        """
        Perform adaptive search that progressively adjusts parameters based on results.
        
        Args:
            query: The search query
            
        Returns:
            Tuple of (list of articles, search strategy used)
        """
        # Try with current parameters
        articles = self._get_all_data(query)
        strategy = "initial"
        
        # If no results or too few, try different strategies
        if len(articles) < 3 and self.adaptive_search:
            logger.info(f"Initial search found only {len(articles)} results, trying alternative strategies")
            
            # Try with expanded date range
            original_from_date = self.from_date
            original_order_by = self.order_by
            
            # Strategy 1: Expand to 6 months
            logger.info("Strategy 1: Expanding time range to 6 months")
            six_months_ago = (datetime.now() - timedelta(days=180)).strftime("%Y-%m-%d")
            self.from_date = six_months_ago
            
            articles1 = self._get_all_data(query)
            if len(articles1) > len(articles):
                articles = articles1
                strategy = "expanded_6mo"
                
            # Strategy 2: Expand to all time and try relevance order
            if len(articles) < 3:
                logger.info("Strategy 2: Expanding to all time with relevance ordering")
                self.from_date = "2000-01-01"  # Effectively "all time"
                self.order_by = "relevance"
                
                articles2 = self._get_all_data(query)
                if len(articles2) > len(articles):
                    articles = articles2
                    strategy = "all_time_relevance"
            
            # Strategy 3: Try removing section constraints
            if len(articles) < 3 and self.section:
                logger.info("Strategy 3: Removing section constraint")
                original_section = self.section
                self.section = None
                
                articles3 = self._get_all_data(query)
                if len(articles3) > len(articles):
                    articles = articles3
                    strategy = "no_section"
                    
                # Restore section setting
                self.section = original_section
            
            # Restore original settings
            self.from_date = original_from_date
            self.order_by = original_order_by
            
        logger.info(f"Adaptive search using strategy '{strategy}' found {len(articles)} results")
        return articles, strategy
    
    def _get_all_data(self, query: str) -> List[Dict[str, Any]]:
        """
        Get all article data from The Guardian API in a single call.
        Always requests all fields for simplicity.
        
        Args:
            query: The search query
            
        Returns:
            List of articles with all data
        """
        try:
            # Ensure query is not empty
            if not query or query.strip() == "":
                query = "news"
                logger.warning("Empty query provided, using 'news' as default")
            
            # Ensure query is not too long for API
            if len(query) > 100:
                logger.warning(f"Query too long for Guardian API ({len(query)} chars), truncating")
                query = query[:100]
            
            # Always request all fields for simplicity
            # Ensure max_results is an integer to avoid comparison errors
            page_size = min(int(self.max_results) if self.max_results is not None else 10, 50)
            
            # Log full parameters for debugging
            logger.info(f"Guardian API search query: '{query}'")
            logger.info(f"Guardian API date range: {self.from_date} to {self.to_date}")
            
            params = {
                "q": query,
                "api-key": self.api_key,
                "from-date": self.from_date,
                "to-date": self.to_date,
                "order-by": self.order_by,
                "page-size": page_size,  # API maximum is 50
                "show-fields": "headline,trailText,byline,body,publication",
                "show-tags": "keyword"
            }
            
            # Add section filter if specified
            if self.section:
                params["section"] = self.section
            
            # Log the complete request parameters (except API key)
            log_params = params.copy()
            log_params["api-key"] = "REDACTED"
            logger.info(f"Guardian API request parameters: {log_params}")
            
            # Execute the API request
            response = requests.get(self.api_url, params=params)
            response.raise_for_status()
            
            data = response.json()
            
            # Extract results from the response
            articles = data.get("response", {}).get("results", [])
            logger.info(f"Guardian API returned {len(articles)} articles")
            
            # Format results to include all data
            formatted_articles = []
            for i, article in enumerate(articles):
                if i >= self.max_results:
                    break
                    
                fields = article.get("fields", {})
                
                # Format the article with all fields
                result = {
                    "id": article.get("id", ""),
                    "title": fields.get("headline", article.get("webTitle", "")),
                    "link": article.get("webUrl", ""),
                    "snippet": fields.get("trailText", ""),
                    "publication_date": article.get("webPublicationDate", ""),
                    "section": article.get("sectionName", ""),
                    "author": fields.get("byline", ""),
                    "content": fields.get("body", ""),
                    "full_content": fields.get("body", "")
                }
                
                # Extract tags/keywords
                tags = article.get("tags", [])
                result["keywords"] = [tag.get("webTitle", "") for tag in tags if tag.get("type") == "keyword"]
                
                formatted_articles.append(result)
            
            return formatted_articles
            
        except Exception as e:
            logger.error(f"Error getting data from The Guardian API: {e}")
            return []
    
    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information for Guardian articles with enhanced optimization.
        
        Args:
            query: The search query
            
        Returns:
            List of preview dictionaries
        """
        logger.info(f"Getting articles from The Guardian API for query: {query}")
        
        # Step 1: Optimize the query using LLM
        optimized_query = self._optimize_query_for_guardian(query)
        
        # Step 2: Adapt date parameters based on query type
        self._adapt_dates_for_query_type(optimized_query)
        
        # Step 3: Perform adaptive search
        articles, strategy = self._adaptive_search(optimized_query)
        
        # Store search metadata for debugging
        self._search_metadata = {
            "original_query": query,
            "optimized_query": optimized_query,
            "strategy": strategy,
            "from_date": self.from_date,
            "to_date": self.to_date,
            "section": self.section,
            "order_by": self.order_by
        }
        
        # Store full articles for later use
        self._full_articles = {a["id"]: a for a in articles}
        
        # Return only preview fields for each article
        previews = []
        for article in articles:
            preview = {
                "id": article["id"],
                "title": article["title"],
                "link": article["link"],
                "snippet": article["snippet"],
                "publication_date": article["publication_date"],
                "section": article["section"],
                "author": article["author"],
                "keywords": article.get("keywords", [])
            }
            previews.append(preview)
        
        return previews
    
    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant Guardian articles.
        Restores full content from the cached data.
        
        Args:
            relevant_items: List of relevant preview dictionaries
            
        Returns:
            List of result dictionaries with full content
        """
        logger.info(f"Adding full content to {len(relevant_items)} relevant Guardian articles")
        
        # Check if we should add full content
        if hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY:
            return relevant_items
            
        # Get full articles for relevant items
        results = []
        for item in relevant_items:
            article_id = item.get("id", "")
            
            # Get the full article from our cache
            if hasattr(self, '_full_articles') and article_id in self._full_articles:
                results.append(self._full_articles[article_id])
            else:
                # If not found (shouldn't happen), just use the preview
                results.append(item)
        
        return results
    
    def run(self, query: str) -> List[Dict[str, Any]]:
        """
        Execute a search using The Guardian API with the enhanced approach.
        
        Args:
            query: The search query
            
        Returns:
            List of search results
        """
        logger.info(f"---Execute a search using The Guardian (enhanced)---")
        
        # Additional safety check for None query
        if query is None:
            logger.error("None query passed to Guardian search engine")
            query = "news"
        
        try:
            # Get previews with our enhanced method
            previews = self._get_previews(query)
            
            # If no results, try one more time with a simplified query
            if not previews:
                simple_query = " ".join([w for w in query.split() if len(w) > 3][:3])
                logger.warning(f"No Guardian articles found, trying simplified query: {simple_query}")
                previews = self._get_previews(simple_query)
                
                # If still no results, try with a very generic query as last resort
                if not previews and "trump" in query.lower():
                    logger.warning("Trying last resort query: 'Donald Trump'")
                    previews = self._get_previews("Donald Trump")
                elif not previews:
                    logger.warning("Trying last resort query: 'news'")
                    previews = self._get_previews("news")
            
            # If still no results after all attempts, return empty list
            if not previews:
                logger.warning(f"No Guardian articles found after multiple attempts")
                return []
            
            # Filter for relevance if we have an LLM
            if self.llm and hasattr(self, 'max_filtered_results') and self.max_filtered_results:
                filtered_items = self._filter_for_relevance(previews, query)
                if not filtered_items:
                    # Fall back to unfiltered results if everything was filtered out
                    logger.warning("All articles filtered out, using unfiltered results")
                    filtered_items = previews[:self.max_filtered_results]
            else:
                filtered_items = previews
            
            # Get full content for relevant items
            results = self._get_full_content(filtered_items)
            
            # Add source information to make it clear these are from The Guardian
            for result in results:
                if "source" not in result:
                    result["source"] = "The Guardian"
            
            # Clean up the cache after use
            if hasattr(self, '_full_articles'):
                del self._full_articles
                
            # Restore original date parameters
            self.from_date = self._original_date_params["from_date"]
            self.to_date = self._original_date_params["to_date"]
            
            # Log search metadata if available
            if hasattr(self, '_search_metadata'):
                logger.info(f"Search metadata: {self._search_metadata}")
                del self._search_metadata
                
            return results
            
        except Exception as e:
            logger.error(f"Error in Guardian search: {e}")
            
            # Restore original date parameters on error
            self.from_date = self._original_date_params["from_date"]
            self.to_date = self._original_date_params["to_date"]
            
            return []
    
    def search_by_section(self, section: str, max_results: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Search for articles in a specific section.
        
        Args:
            section: The Guardian section name (e.g., "politics", "technology")
            max_results: Maximum number of results (defaults to self.max_results)
            
        Returns:
            List of articles in the section
        """
        original_section = self.section
        original_max_results = self.max_results
        
        try:
            # Set section and max_results for this search
            self.section = section
            if max_results:
                self.max_results = max_results
                
            # Use empty query to get all articles in the section
            return self.run("")
            
        finally:
            # Restore original values
            self.section = original_section
            self.max_results = original_max_results
    
    def get_recent_articles(self, days: int = 7, max_results: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Get recent articles from The Guardian.
        
        Args:
            days: Number of days to look back
            max_results: Maximum number of results (defaults to self.max_results)
            
        Returns:
            List of recent articles
        """
        original_from_date = self.from_date
        original_order_by = self.order_by
        original_max_results = self.max_results
        
        try:
            # Set parameters for this search
            self.from_date = (datetime.now() - timedelta(days=days)).strftime("%Y-%m-%d")
            self.order_by = "newest"
            if max_results:
                self.max_results = max_results
                
            # Use empty query to get all recent articles
            return self.run("")
            
        finally:
            # Restore original values
            self.from_date = original_from_date
            self.order_by = original_order_by
            self.max_results = original_max_results