"""
Follow-up Relevance Filter

Filters and ranks past research sources based on their relevance
to follow-up questions.
"""

from typing import Dict, List
from loguru import logger

from .base_filter import BaseFilter
from ...utilities.json_utils import extract_json, get_llm_response_text


class FollowUpRelevanceFilter(BaseFilter):
    """
    Filters past research sources by relevance to follow-up questions.

    This filter analyzes sources from previous research and determines
    which ones are most relevant to the new follow-up question.
    """

    def filter_results(
        self, results: List[Dict], query: str, max_results: int = 10, **kwargs
    ) -> List[Dict]:
        """
        Filter search results by relevance to the follow-up query.

        Args:
            results: List of source dictionaries from past research
            query: The follow-up query
            max_results: Maximum number of results to return (default: 10)
            **kwargs: Additional parameters:
                - past_findings: Summary of past findings for context
                - original_query: The original research query

        Returns:
            Filtered list of relevant sources
        """
        if not results:
            return []

        past_findings = kwargs.get("past_findings", "")
        original_query = kwargs.get("original_query", "")

        # Use LLM to select relevant sources
        relevant_indices = self._select_relevant_sources(
            results, query, past_findings, max_results, original_query
        )

        # Return selected sources
        filtered = [results[i] for i in relevant_indices if i < len(results)]

        logger.info(
            f"Filtered {len(results)} sources to {len(filtered)} relevant ones "
            f"for follow-up query. Kept indices: {relevant_indices}"
        )

        return filtered

    def _select_relevant_sources(
        self,
        sources: List[Dict],
        query: str,
        context: str,
        max_results: int,
        original_query: str = "",
    ) -> List[int]:
        """
        Select relevant sources using LLM.

        Args:
            sources: List of source dictionaries
            query: The follow-up query
            context: Past findings context
            max_results: Maximum number of sources to select
            original_query: The original research query

        Returns:
            List of indices of relevant sources
        """
        if not self.model:
            # If no model available, return first max_results
            return list(range(min(max_results, len(sources))))

        # Build source list for LLM
        source_list = []
        for i, source in enumerate(sources):
            title = source.get("title") or "Unknown"
            url = source.get("url") or ""
            snippet = (
                source.get("snippet") or source.get("content_preview") or ""
            )[:150]
            source_list.append(
                f"{i}. {title}\n   URL: {url}\n   Content: {snippet}"
            )

        sources_text = "\n\n".join(source_list)

        # Include context if available for better selection
        context_section = ""
        if context or original_query:
            parts = []
            if original_query:
                parts.append(f"Original research question: {original_query}")
            if context:
                parts.append(f"Previous research findings:\n{context}")

            context_section = f"""
Previous Research Context:
{chr(10).join(parts)}

---
"""

        prompt = f"""
Select the most relevant sources for answering this follow-up question based on the previous research context.
{context_section}
Follow-up question: "{query}"

Available sources from previous research:
{sources_text}

Instructions:
- Select sources that are most relevant to the follow-up question given the context
- Consider which sources directly address the question or provide essential information
- Think about what the user is asking for in relation to the previous findings
- Return ONLY a JSON array of source numbers (e.g., [0, 2, 5, 7])
- Do not include any explanation or other text

Return the indices of relevant sources as a JSON array:"""

        try:
            response = self.model.invoke(prompt)
            content = get_llm_response_text(response)

            # Parse JSON response
            indices = extract_json(content, expected_type=list)

            if indices is not None:
                # Validate it's a list of integers
                indices = [
                    int(i)
                    for i in indices
                    if isinstance(i, (int, float)) and int(i) < len(sources)
                ]
            else:
                logger.debug("Failed to parse JSON, attempting regex fallback")
                # Fallback to regex extraction
                import re

                numbers = re.findall(r"\d+", content)
                indices = [int(n) for n in numbers if int(n) < len(sources)]

            return indices
        except Exception as e:
            logger.debug(f"LLM source selection failed: {e}")
            # Fallback to first max_results sources
            return list(range(min(max_results, len(sources))))
