import justext
from langchain_community.document_loaders import AsyncChromiumLoader
from langchain_community.document_transformers import BeautifulSoupTransformer
from langchain_core.language_models import BaseLLM
from typing import List, Dict, Any, Optional, Union
import json
import os
from .utilties.search_utilities import remove_think_tags
from datetime import datetime
from local_deep_research import config

class FullSearchResults:
    """
    Enhanced web content retrieval class that works with the BaseSearchEngine architecture.
    Can be used as a wrapper around web-based search engines like DuckDuckGo and SerpAPI.
    """
    
    def __init__(
        self,
        llm: BaseLLM,
        web_search,        
        output_format: str = "list",
        language: str = "English",
        max_results: int = 10,
        region: str = "wt-wt",
        time: str = "y",
        safesearch: str = "Moderate"
    ):
        """
        Initialize the full search results processor.
        
        Args:
            llm: Language model instance for relevance filtering
            web_search: Web search engine instance that provides initial results
            output_format: Format of output ('list' or other formats)
            language: Language for content processing
            max_results: Maximum number of search results
            region: Search region
            time: Time period for search results
            safesearch: Safe search setting
        """
        self.llm = llm
        self.output_format = output_format
        self.language = language
        self.max_results = max_results
        self.region = region
        self.time = time
        self.safesearch = safesearch
        self.web_search = web_search
        os.environ["USER_AGENT"] = "Local Deep Research/1.0"

        self.bs_transformer = BeautifulSoupTransformer()
        self.tags_to_extract = ["p", "div", "span"]
    
    def run(self, query: str) -> List[Dict[str, Any]]:
        """
        Legacy method that performs a full search in one step.
        Respects config parameters:
        - SEARCH_SNIPPETS_ONLY: If True, only returns snippets without full content
        - SKIP_RELEVANCE_FILTER: If True, returns all results without filtering
        
        Args:
            query: The search query
            
        Returns:
            List of search results with full content (unless SEARCH_SNIPPETS_ONLY is True)
        """
        # Phase 1: Get search results from the web search engine
        previews = self._get_previews(query)
        if not previews:
            return []
            
        # Phase 2: Filter URLs using LLM (unless SKIP_RELEVANCE_FILTER is True)
        if hasattr(config, 'SKIP_RELEVANCE_FILTER') and config.SKIP_RELEVANCE_FILTER:
            relevant_items = previews
            print("Skipping relevance filtering as per config")
        else:
            relevant_items = self._filter_relevant_items(previews, query)
            if not relevant_items:
                return []
            
        # Phase 3: Get full content for relevant items (unless SEARCH_SNIPPETS_ONLY is True)
        if hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY:
            print("Returning snippet-only results as per config")
            return relevant_items
        else:
            results = self._get_full_content(relevant_items)
            return results
    
    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information from the web search engine.
        
        Args:
            query: The search query
            
        Returns:
            List of preview dictionaries
        """
        try:
            # Get search results from the web search engine
            search_results = self.web_search.invoke(query)
            
            if not isinstance(search_results, list):
                print("Error: Expected search results in list format")
                return []
            
            # Return the results as previews
            return search_results
            
        except Exception as e:
            print(f"Error getting previews: {e}")
            return []
    
    def _filter_relevant_items(self, previews: List[Dict[str, Any]], query: str) -> List[Dict[str, Any]]:
        """
        Filter previews for relevance using LLM.
        
        Args:
            previews: List of preview dictionaries
            query: The original search query
            
        Returns:
            List of relevant preview dictionaries
        """
        # Skip filtering if disabled in config or no previews
        if not config.QUALITY_CHECK_DDG_URLS or not previews:
            return previews
        
        # Format for LLM evaluation
        now = datetime.now()
        current_time = now.strftime("%Y-%m-%d")
        prompt = f"""ONLY Return a JSON array. The response contains no letters. Evaluate these URLs for:
            1. Timeliness (today: {current_time})
            2. Factual accuracy (cross-reference major claims)
            3. Source reliability (prefer official company websites, established news outlets)
            4. Direct relevance to query: {query}

            URLs to evaluate:
            {json.dumps(previews, indent=2)}

            Return a JSON array of indices (0-based) for sources that meet ALL criteria.
            ONLY Return a JSON array of indices (0-based) and nothing else. No letters. 
            Example response: \n[0, 2, 4]\n\n"""
        
        try:
            # Get LLM's evaluation
            response = self.llm.invoke(prompt)
            
            # Extract JSON array from response
            response_text = remove_think_tags(response.content)
            # Clean up response to handle potential formatting issues
            response_text = response_text.strip()
            
            # Find the first occurrence of '[' and the last occurrence of ']'
            start_idx = response_text.find('[')
            end_idx = response_text.rfind(']')
            
            if start_idx >= 0 and end_idx > start_idx:
                array_text = response_text[start_idx:end_idx+1]
                good_indices = json.loads(array_text)
                
                # Return only the results with good indices
                return [r for i, r in enumerate(previews) if i in good_indices]
            else:
                print("Could not find JSON array in response, returning all previews")
                return previews
                
        except Exception as e:
            print(f"URL filtering error: {e}")
            # Fall back to returning all previews on error
            return previews
    
    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant items by retrieving and processing web pages.
        
        Args:
            relevant_items: List of relevant preview dictionaries
            
        Returns:
            List of result dictionaries with full content
        """
        nr_full_text = 0
        
        # Extract URLs from relevant items
        urls = [item.get("link") for item in relevant_items if item.get("link")]
        
        if not urls:
            print("\n === NO VALID LINKS ===\n")
            return relevant_items
        
        try:
            # Download the full HTML pages for filtered URLs
            loader = AsyncChromiumLoader(urls)
            html_docs = loader.load()
            
            # Process the HTML using BeautifulSoupTransformer
            full_docs = self.bs_transformer.transform_documents(
                html_docs, tags_to_extract=self.tags_to_extract
            )
            
            # Remove boilerplate from each document
            url_to_content = {}
            for doc in full_docs:
                nr_full_text += 1
                source = doc.metadata.get("source")
                if source:
                    cleaned_text = self._remove_boilerplate(doc.page_content)
                    url_to_content[source] = cleaned_text
            
            # Attach the cleaned full content to each result
            results = []
            for item in relevant_items:
                new_item = item.copy()
                link = item.get("link")
                new_item["full_content"] = url_to_content.get(link, None)
                results.append(new_item)
            
            print(f"FULL SEARCH WITH FILTERED URLS - Full text retrieved: {nr_full_text}")
            return results
            
        except Exception as e:
            print(f"Error retrieving full content: {e}")
            # Return original items if full content retrieval fails
            return relevant_items
    
    def _remove_boilerplate(self, html: str) -> str:
        """
        Remove boilerplate content from HTML.
        
        Args:
            html: HTML content
            
        Returns:
            Cleaned text content
        """
        if not html or not html.strip():
            return ""
        try:
            paragraphs = justext.justext(html, justext.get_stoplist(self.language))
            cleaned = "\n".join([p.text for p in paragraphs if not p.is_boilerplate])
            return cleaned
        except Exception as e:
            print(f"Error removing boilerplate: {e}")
            return html
    
    def invoke(self, query: str) -> List[Dict[str, Any]]:
        """Compatibility method for LangChain tools"""
        return self.run(query)
    
    def __call__(self, query: str) -> List[Dict[str, Any]]:
        """Make the class callable like a function"""
        return self.invoke(query)