from loguru import logger
from datetime import datetime, UTC
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable

from langchain_core.language_models import BaseLLM

from ...config.search_config import QUALITY_CHECK_DDG_URLS
from ...research_library.downloaders.extraction import (
    batch_fetch_and_extract,
)
from ...security.egress.policy import PolicyDeniedError
from ...security.ssrf_validator import redact_url_for_log, validate_url
from ...utilities.js_rendering import (
    read_js_rendering_setting as _read_js_rendering_setting,
)
from ...utilities.json_utils import extract_json, get_llm_response_text


@runtime_checkable
class _Invokable(Protocol):
    def invoke(self, query: str) -> Any: ...


class FullSearchResults:
    def __init__(
        self,
        llm: Optional[BaseLLM],
        web_search: _Invokable,
        output_format: str = "list",
        language: str = "English",
        max_results: int = 10,
        region: str = "wt-wt",
        time: Optional[str] = "y",
        safesearch: str | int = "Moderate",
        settings_snapshot: Optional[Dict] = None,
        egress_context: Optional[Any] = None,
    ):
        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.settings_snapshot = settings_snapshot
        # Set by the factory when the parent engine is gated against a
        # specific scope; used to evaluate per-URL fetches below.
        self.egress_context = egress_context

    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:
            if self.llm is None:
                return results
            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 [r for i, r in enumerate(results) if i in good_indices]
        except PolicyDeniedError:
            # The URL-quality LLM was denied by egress policy (e.g. a cloud
            # LLM under require_local / PRIVATE_ONLY). Do NOT fall through to
            # the unfiltered-results fallback below — that would let the run
            # proceed despite the user's policy refusing the LLM. Fail closed.
            raise
        except Exception:
            logger.exception("URL filtering error")
            logger.warning(
                "URL quality filter unavailable — returning {} unfiltered "
                "results as fallback",
                len(results),
            )
            return results  # Fall back to original results on LLM error

    def run(self, query: str):
        # 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 []

        # SSRF-validate + egress-scope-validate each URL. SSRF only checks
        # IP class; scope enforcement is the separate axis that blocks
        # public hosts under PRIVATE_ONLY / STRICT.
        evaluate_url_fn = None
        egress_ctx = self.egress_context
        if egress_ctx is not None:
            from ...security.egress.policy import evaluate_url as _ev_url

            evaluate_url_fn = _ev_url
        safe_urls: List[str] = []
        for url in urls:
            if url is None:
                continue
            if not validate_url(url):
                logger.warning(
                    "SSRF validation blocked URL from full content fetch: "
                    f"{redact_url_for_log(url)}."
                )
                continue
            if evaluate_url_fn is not None and egress_ctx is not None:
                url_decision = evaluate_url_fn(url, egress_ctx)
                if not url_decision.allowed:
                    logger.bind(policy_audit=True).warning(
                        "full_search URL denied by egress policy",
                        url=redact_url_for_log(url),
                        scope=egress_ctx.scope.value,
                        reason=url_decision.reason,
                    )
                    continue
            safe_urls.append(url)

        if not safe_urls:
            logger.warning(
                "All URLs were blocked by SSRF validation — returning results "
                "without full content. This can happen when search results "
                "point to internal/private network addresses."
            )
            for result in filtered_results:
                result["full_content"] = None
            return filtered_results

        # Fetch and extract all pages — specialized downloaders (arXiv,
        # PubMed, etc.) are tried first, with HTML crawling as fallback.
        url_to_content = batch_fetch_and_extract(
            safe_urls,
            language=self.language,
            enable_js_rendering=_read_js_rendering_setting(
                self.settings_snapshot
            ),
        )

        nr_full_text = sum(1 for v in url_to_content.values() if v)
        for result in filtered_results:
            link = result.get("link")
            result["full_content"] = url_to_content.get(link) if link else None

        logger.info(f"Full search: retrieved content from {nr_full_text} pages")
        return filtered_results

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Fetch and attach full content to an existing list of items."""
        evaluate_url_fn = None
        egress_ctx = self.egress_context
        if egress_ctx is not None:
            from ...security.egress.policy import evaluate_url as _ev_url

            evaluate_url_fn = _ev_url
        urls: List[str] = []
        for item in relevant_items:
            link = item.get("link")
            if link is None:
                continue
            if not validate_url(link):
                logger.warning(
                    "SSRF validation blocked URL from full content fetch: "
                    f"{redact_url_for_log(link)}."
                )
                continue
            if evaluate_url_fn is not None and egress_ctx is not None:
                url_decision = evaluate_url_fn(link, egress_ctx)
                if not url_decision.allowed:
                    logger.bind(policy_audit=True).warning(
                        "full_search _get_full_content URL denied by egress policy",
                        url=redact_url_for_log(link),
                        scope=egress_ctx.scope.value,
                        reason=url_decision.reason,
                    )
                    continue
            urls.append(link)

        if not urls:
            for item in relevant_items:
                item["full_content"] = None
            return relevant_items

        try:
            url_to_content = batch_fetch_and_extract(
                urls,
                language=self.language,
                enable_js_rendering=_read_js_rendering_setting(
                    self.settings_snapshot
                ),
            )
        except Exception:
            logger.exception("Error fetching full content")
            for item in relevant_items:
                item["full_content"] = None
            return relevant_items

        for item in relevant_items:
            link = item.get("link")
            item["full_content"] = url_to_content.get(link) if link else None

        return relevant_items

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

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