"""
Follow-up Context Manager

Manages and processes past research context for follow-up questions.
This is a standalone class that doesn't inherit from BaseKnowledgeGenerator
to avoid implementing many abstract methods.
"""

from typing import Dict, List, Any, Optional
from loguru import logger

from langchain_core.language_models.chat_models import BaseChatModel
from ...utilities.search_utilities import remove_think_tags


class FollowUpContextHandler:
    """
    Manages past research context for follow-up research.

    This class handles:
    1. Loading and structuring past research data
    2. Summarizing findings for follow-up context
    3. Extracting relevant information for new searches
    4. Building comprehensive context for strategies
    """

    def __init__(
        self, model: BaseChatModel, settings_snapshot: Optional[Dict] = None
    ):
        """
        Initialize the context manager.

        Args:
            model: Language model for processing context
            settings_snapshot: Optional settings snapshot
        """
        self.model = model
        self.settings_snapshot = settings_snapshot or {}
        self.past_research_cache = {}

    def build_context(
        self, research_data: Dict[str, Any], follow_up_query: str
    ) -> Dict[str, Any]:
        """
        Build comprehensive context from past research.

        Args:
            research_data: Past research data including findings, sources, etc.
            follow_up_query: The follow-up question being asked

        Returns:
            Structured context dictionary for follow-up research
        """
        logger.info(f"Building context for follow-up: {follow_up_query}")

        # Extract all components
        return {
            "parent_research_id": research_data.get("research_id", ""),
            "original_query": research_data.get("query", ""),
            "follow_up_query": follow_up_query,
            "past_findings": self._extract_findings(research_data),
            "past_sources": self._extract_sources(research_data),
            "key_entities": self._extract_entities(research_data),
            "summary": self._create_summary(research_data, follow_up_query),
            "report_content": research_data.get("report_content", ""),
            "formatted_findings": research_data.get("formatted_findings", ""),
            "all_links_of_system": research_data.get("all_links_of_system", []),
            "metadata": self._extract_metadata(research_data),
        }

    def _extract_findings(self, research_data: Dict) -> str:
        """
        Extract and format findings from past research.

        Args:
            research_data: Past research data

        Returns:
            Formatted findings string
        """
        findings_parts = []

        # Check various possible locations for findings
        if formatted := research_data.get("formatted_findings"):
            findings_parts.append(formatted)

        if report := research_data.get("report_content"):
            # Take first part of report if no formatted findings
            if not findings_parts:
                findings_parts.append(report[:2000])

        if not findings_parts:
            return "No previous findings available"

        return "\n\n".join(findings_parts)

    def _extract_sources(self, research_data: Dict) -> List[Dict]:
        """
        Extract and structure sources from past research.

        Args:
            research_data: Past research data

        Returns:
            List of source dictionaries
        """
        sources = []
        seen_urls = set()

        # Check all possible source fields
        for field in ["resources", "all_links_of_system", "past_links"]:
            if field_sources := research_data.get(field, []):
                for source in field_sources:
                    url = source.get("url", "")
                    # Avoid duplicates by URL
                    if url and url not in seen_urls:
                        sources.append(source)
                        seen_urls.add(url)
                    elif not url:
                        # Include sources without URLs (shouldn't happen but be safe)
                        sources.append(source)

        return sources

    def _extract_entities(self, research_data: Dict) -> List[str]:
        """
        Extract key entities from past research.

        Args:
            research_data: Past research data

        Returns:
            List of key entities
        """
        findings = self._extract_findings(research_data)

        if not findings or not self.model:
            return []

        prompt = f"""
Extract key entities (names, places, organizations, concepts) from these research findings:

{findings[:2000]}

Return up to 10 most important entities, one per line.
"""

        try:
            response = self.model.invoke(prompt)
            entities = [
                line.strip()
                for line in remove_think_tags(response.content)
                .strip()
                .split("\n")
                if line.strip()
            ]
            return entities[:10]
        except Exception:
            logger.warning("Failed to extract entities")
            return []

    def _create_summary(self, research_data: Dict, follow_up_query: str) -> str:
        """
        Create a targeted summary of past research relevant to the follow-up question.
        This is used internally for building context.

        Args:
            research_data: Past research data
            follow_up_query: The follow-up question

        Returns:
            Targeted summary for context building
        """
        findings = self._extract_findings(research_data)
        original_query = research_data.get("query", "")

        # For internal context, create a brief targeted summary
        return self._generate_summary(
            findings=findings,
            query=follow_up_query,
            original_query=original_query,
            max_sentences=5,
            purpose="context",
        )

    def _extract_metadata(self, research_data: Dict) -> Dict:
        """
        Extract metadata from past research.

        Args:
            research_data: Past research data

        Returns:
            Metadata dictionary
        """
        return {
            "strategy": research_data.get("strategy", ""),
            "mode": research_data.get("mode", ""),
            "created_at": research_data.get("created_at", ""),
            "research_meta": research_data.get("research_meta", {}),
        }

    def summarize_for_followup(
        self, findings: str, query: str, max_length: int = 1000
    ) -> str:
        """
        Create a concise summary of findings for external use (e.g., in prompts).
        This creates a length-constrained summary suitable for inclusion in LLM prompts.

        Args:
            findings: Past research findings
            query: Follow-up query
            max_length: Maximum length of summary in characters

        Returns:
            Concise summary constrained to max_length
        """
        # Use the shared summary generation with specific parameters for external use
        return self._generate_summary(
            findings=findings,
            query=query,
            original_query=None,
            max_sentences=max_length
            // 100,  # Approximate sentences based on length
            purpose="prompt",
            max_length=max_length,
        )

    def _generate_summary(
        self,
        findings: str,
        query: str,
        original_query: Optional[str] = None,
        max_sentences: int = 5,
        purpose: str = "context",
        max_length: Optional[int] = None,
    ) -> str:
        """
        Shared summary generation logic.

        Args:
            findings: Research findings to summarize
            query: Follow-up query
            original_query: Original research query (optional)
            max_sentences: Maximum number of sentences
            purpose: Purpose of summary ("context" or "prompt")
            max_length: Maximum character length (optional)

        Returns:
            Generated summary
        """
        if not findings:
            return ""

        # If findings are already short enough, return as-is
        if max_length and len(findings) <= max_length:
            return findings

        if not self.model:
            # Fallback without model
            if max_length:
                return findings[:max_length] + "..."
            return findings[:500] + "..."

        # Build prompt based on purpose
        if purpose == "context" and original_query:
            prompt = f"""
Create a brief summary of previous research findings that are relevant to this follow-up question:

Original research question: "{original_query}"
Follow-up question: "{query}"

Previous findings:
{findings[:3000]}

Provide a {max_sentences}-sentence summary focusing on aspects relevant to the follow-up question.
"""
        else:
            prompt = f"""
Summarize these research findings in relation to the follow-up question:

Follow-up question: "{query}"

Findings:
{findings[:4000]}

Create a summary of {max_sentences} sentences that captures the most relevant information.
"""

        try:
            response = self.model.invoke(prompt)
            summary = remove_think_tags(response.content).strip()

            # Apply length constraint if specified
            if max_length and len(summary) > max_length:
                summary = summary[:max_length] + "..."

            return summary
        except Exception:
            logger.warning("Summary generation failed")
            # Fallback to truncation
            if max_length:
                return findings[:max_length] + "..."
            return findings[:500] + "..."

    def identify_gaps(
        self, research_data: Dict, follow_up_query: str
    ) -> List[str]:
        """
        Identify information gaps that the follow-up should address.

        Args:
            research_data: Past research data
            follow_up_query: Follow-up question

        Returns:
            List of identified gaps
        """
        findings = self._extract_findings(research_data)

        if not findings or not self.model:
            return []

        prompt = f"""
Based on the previous research and the follow-up question, identify information gaps:

Previous research findings:
{findings[:2000]}

Follow-up question: "{follow_up_query}"

What specific information is missing or needs clarification? List up to 5 gaps, one per line.
"""

        try:
            response = self.model.invoke(prompt)
            gaps = [
                line.strip()
                for line in remove_think_tags(response.content)
                .strip()
                .split("\n")
                if line.strip()
            ]
            return gaps[:5]
        except Exception:
            logger.warning("Failed to identify gaps")
            return []

    def format_for_settings_snapshot(
        self, context: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Format context for inclusion in settings snapshot.
        Only includes essential metadata, not actual content.

        Args:
            context: Full context dictionary

        Returns:
            Minimal metadata for settings snapshot
        """
        # Only include minimal metadata in settings snapshot
        # Settings snapshot should be for settings, not data
        return {
            "followup_metadata": {
                "parent_research_id": context.get("parent_research_id"),
                "is_followup": True,
                "has_context": bool(context.get("past_findings")),
            }
        }

    def get_relevant_context_for_llm(
        self, context: Dict[str, Any], max_tokens: int = 2000
    ) -> str:
        """
        Get a concise version of context for LLM prompts.

        Args:
            context: Full context dictionary
            max_tokens: Approximate maximum tokens

        Returns:
            Concise context string
        """
        parts = []

        # Add original and follow-up queries
        parts.append(f"Original research: {context.get('original_query', '')}")
        parts.append(
            f"Follow-up question: {context.get('follow_up_query', '')}"
        )

        # Add summary
        if summary := context.get("summary"):
            parts.append(f"\nPrevious findings summary:\n{summary}")

        # Add key entities
        if entities := context.get("key_entities"):
            parts.append(f"\nKey entities: {', '.join(entities[:5])}")

        # Add source count
        if sources := context.get("past_sources"):
            parts.append(f"\nAvailable sources: {len(sources)}")

        result = "\n".join(parts)

        # Truncate if needed (rough approximation: 4 chars per token)
        max_chars = max_tokens * 4
        if len(result) > max_chars:
            result = result[:max_chars] + "..."

        return result
