import requests
from typing import Dict, List, Any, Optional
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


class GuardianSearchEngine(BaseSearchEngine):
    """The Guardian API search engine implementation"""
    
    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):
        """
        Initialize The Guardian search engine.
        
        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
        """
        # Initialize the BaseSearchEngine with the LLM
        super().__init__(llm=llm)
        
        self.max_results = max_results
        self.api_key = api_key or os.getenv("GUARDIAN_API_KEY")
        
        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
        
        # API base URL
        self.api_url = "https://content.guardianapis.com/search"
    
    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:
            # Always request all fields for simplicity
            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": min(self.max_results, 50),  # 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
            
            # 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", [])
            
            # 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:
            print(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.
        Actually gets all data but returns only preview fields.
        
        Args:
            query: The search query
            
        Returns:
            List of preview dictionaries
        """
        print("Getting articles from The Guardian API")
        
        # Get all article data
        articles = self._get_all_data(query)
        
        # Store full articles for later use (implementation detail)
        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
        """
        print("Adding full content to 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 two-phase approach.
        
        Args:
            query: The search query
            
        Returns:
            List of search results
        """
        print("---Execute a search using The Guardian---")
        
        # Use the implementation from the parent class which handles all phases
        results = super().run(query)
        
        # Clean up the cache after use
        if hasattr(self, '_full_articles'):
            del self._full_articles
            
        return results
    
    def get_article_by_id(self, article_id: str) -> Dict[str, Any]:
        """
        Get a specific article by its ID.
        
        Args:
            article_id: The Guardian article ID
            
        Returns:
            Dictionary with article information
        """
        try:
            # Guardian article API URL
            url = f"https://content.guardianapis.com/{article_id}"
            
            # Always request all fields
            response = requests.get(
                url, 
                params={
                    "api-key": self.api_key,
                    "show-fields": "headline,trailText,body,byline,publication",
                    "show-tags": "keyword"
                }
            )
            response.raise_for_status()
            
            data = response.json()
            article = data.get("response", {}).get("content", {})
            
            if not article:
                return {}
                
            fields = article.get("fields", {})
            
            # Format the article with all fields
            result = {
                "id": article_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", "")
            }
            
            # Only include full content if not in snippet-only mode
            if not hasattr(config, 'SEARCH_SNIPPETS_ONLY') or not config.SEARCH_SNIPPETS_ONLY:
                result["content"] = fields.get("body", "")
                result["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"]
            
            return result
            
        except Exception as e:
            print(f"Error getting article details: {e}")
            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 search 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