from loguru import logger
from datetime import datetime, UTC
from typing import Dict, List

import justext
from langchain_community.document_loaders import AsyncChromiumLoader
from langchain_community.document_transformers import BeautifulSoupTransformer
from langchain_core.language_models import BaseLLM

from ...config.search_config import QUALITY_CHECK_DDG_URLS
from ...utilities.json_utils import extract_json, get_llm_response_text


class FullSearchResults:
    def __init__(
        self,
        llm: BaseLLM,  # Add LLM parameter
        web_search: list,
        output_format: str = "list",
        language: str = "English",
        max_results: int = 10,
        region: str = "wt-wt",
        time: str = "y",
        safesearch: str | int = "Moderate",
    ):
        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

        self.bs_transformer = BeautifulSoupTransformer()
        self.tags_to_extract = ["p", "div", "span"]

    def check_urls(self, results: List[Dict], query: str) -> List[Dict]:
        if not results:
            return results

        now = datetime.now(UTC)
        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:
            {results}

            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)
            response_text = get_llm_response_text(response)
            good_indices = extract_json(response_text, expected_type=list)

            if good_indices is None:
                good_indices = []

            # Return only the results with good URLs
            return [r for i, r in enumerate(results) if i in good_indices]
        except Exception:
            logger.exception("URL filtering error")
            return []

    def remove_boilerplate(self, html: str) -> str:
        if not html or not html.strip():
            return ""
        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

    def run(self, query: str):
        nr_full_text = 0
        # Step 1: Get search results
        search_results = self.web_search.invoke(query)
        if not isinstance(search_results, list):
            raise ValueError("Expected the search results in list format.")

        # Step 2: Filter URLs using LLM
        if QUALITY_CHECK_DDG_URLS:
            filtered_results = self.check_urls(search_results, query)
        else:
            filtered_results = search_results

        # Extract URLs from filtered results
        urls = [
            result.get("link")
            for result in filtered_results
            if result.get("link")
        ]

        if not urls:
            logger.error("\n === NO VALID LINKS ===\n")
            return []

        # Step 3: Download the full HTML pages for filtered URLs
        loader = AsyncChromiumLoader(urls)
        html_docs = loader.load()

        # Step 4: Process the HTML using BeautifulSoupTransformer
        full_docs = self.bs_transformer.transform_documents(
            html_docs, tags_to_extract=self.tags_to_extract
        )

        # Step 5: Remove boilerplate from each document
        url_to_content = {}
        for doc in full_docs:
            nr_full_text = 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 filtered result
        for result in filtered_results:
            link = result.get("link")
            result["full_content"] = url_to_content.get(link)

        logger.info("FULL SEARCH WITH FILTERED URLS")
        logger.info("Full text retrieved: ", nr_full_text)
        return filtered_results

    def invoke(self, query: str):
        return self.run(query)

    def __call__(self, query: str):
        return self.invoke(query)
