"""
Concept extraction chain for identifying and summarizing key concepts.
"""

import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from ..core.retrieval import RetrievalService, RetrievalResult
from ..core.models import LLMClient, GenerationConfig
from ..core.prompts import get_prompt
from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class Concept:
    """Represents an extracted concept."""
    name: str
    definition: str = ""
    intuition: str = ""
    related: List[str] = field(default_factory=list)
    source_sections: List[str] = field(default_factory=list)
    importance_score: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "definition": self.definition,
            "intuition": self.intuition,
            "related": self.related,
            "source_sections": self.source_sections,
            "importance_score": self.importance_score
        }


class ConceptExtractor:
    """
    Chain for extracting concepts from documents.
    
    Uses RAG to identify and define key concepts with
    definitions and intuitive explanations.
    """
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        max_concepts: int = 50,
        include_definitions: bool = True,
        include_intuitions: bool = True
    ):
        """
        Initialize concept extractor.
        
        Args:
            retrieval_service: Service for retrieving context
            llm_client: LLM for generation
            max_concepts: Maximum number of concepts to extract
            include_definitions: Include formal definitions
            include_intuitions: Include intuitive explanations
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.max_concepts = max_concepts
        self.include_definitions = include_definitions
        self.include_intuitions = include_intuitions
    
    @trace_function(name="extract_concepts")
    def extract(
        self,
        queries: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> List[Concept]:
        """
        Extract concepts from the indexed document.
        
        Args:
            queries: Optional specific queries for concept retrieval
            config: Generation configuration
            
        Returns:
            List of extracted concepts
        """
        config = config or GenerationConfig()
        
        # Default queries for concept extraction
        if queries is None:
            queries = [
                "key concepts definitions terminology",
                "fundamental ideas principles",
                "important terms notation",
                "core concepts main ideas",
            ]
        
        all_concepts = []
        seen_names = set()
        
        for query in queries:
            # Retrieve relevant content
            results = self.retrieval_service.retrieve_narrative(query, top_k=5)
            
            if not results:
                continue
            
            # Extract concepts from context
            concepts = self._extract_from_context(results, config)
            
            # Deduplicate
            for concept in concepts:
                if concept.name.lower() not in seen_names:
                    all_concepts.append(concept)
                    seen_names.add(concept.name.lower())
        
        # Sort by importance and limit
        all_concepts.sort(key=lambda c: c.importance_score, reverse=True)
        
        logger.info(f"Extracted {len(all_concepts)} concepts")
        return all_concepts[:self.max_concepts]
    
    def _extract_from_context(
        self,
        results: List[RetrievalResult],
        config: GenerationConfig
    ) -> List[Concept]:
        """Extract concepts from retrieval results."""
        context = "\n\n".join(r.content for r in results)
        
        # Build prompt
        prompt = get_prompt("concept_extraction")(context=context)
        
        # Generate
        result = self.llm_client.generate(prompt, config)
        
        # Parse response
        concepts = self._parse_response(result.text, results)
        
        return concepts
    
    def _parse_response(
        self,
        text: str,
        results: List[RetrievalResult]
    ) -> List[Concept]:
        """Parse concept extraction response."""
        concepts = []
        current: Dict[str, Any] = {}
        
        # Track source sections
        source_sections = list(set(
            section
            for r in results
            for section in r.section_hierarchy
            if section
        ))
        
        for line in text.split("\n"):
            line = line.strip()
            
            if line.startswith("CONCEPT:"):
                if current and "name" in current:
                    concepts.append(self._create_concept(current, source_sections))
                current = {"name": line[8:].strip()}
            elif line.startswith("DEFINITION:") and self.include_definitions:
                current["definition"] = line[11:].strip()
            elif line.startswith("INTUITION:") and self.include_intuitions:
                current["intuition"] = line[10:].strip()
            elif line.startswith("RELATED:"):
                related_str = line[8:].strip()
                current["related"] = [r.strip() for r in related_str.split(",") if r.strip()]
        
        # Don't forget the last concept
        if current and "name" in current:
            concepts.append(self._create_concept(current, source_sections))
        
        return concepts
    
    def _create_concept(
        self,
        data: Dict[str, Any],
        source_sections: List[str]
    ) -> Concept:
        """Create a Concept from parsed data."""
        # Calculate importance score based on completeness
        score = 0.0
        if data.get("definition"):
            score += 0.4
        if data.get("intuition"):
            score += 0.3
        if data.get("related"):
            score += 0.2
        if len(data.get("name", "")) > 3:
            score += 0.1
        
        return Concept(
            name=data.get("name", ""),
            definition=data.get("definition", ""),
            intuition=data.get("intuition", ""),
            related=data.get("related", []),
            source_sections=source_sections,
            importance_score=score
        )
    
    def summarize_concepts(
        self,
        concepts: List[Concept],
        max_words: int = 500,
        config: Optional[GenerationConfig] = None
    ) -> str:
        """
        Summarize concepts into a cohesive text.
        
        Args:
            concepts: Concepts to summarize
            max_words: Maximum words
            config: Generation configuration
            
        Returns:
            Summarized concept text
        """
        config = config or GenerationConfig()
        
        # Format concepts for prompt
        concept_text = "\n\n".join(
            f"- **{c.name}**: {c.definition}"
            for c in concepts
            if c.definition
        )
        
        prompt = get_prompt("concept_summary")(
            concepts=concept_text,
            max_words=max_words
        )
        
        result = self.llm_client.generate(prompt, config)
        return result.text
    
    def group_by_topic(self, concepts: List[Concept]) -> Dict[str, List[Concept]]:
        """
        Group concepts by related topics.
        
        Args:
            concepts: Concepts to group
            
        Returns:
            Dictionary mapping topic to concepts
        """
        # Simple grouping based on related concepts
        groups: Dict[str, List[Concept]] = {"General": []}
        
        for concept in concepts:
            if concept.related:
                # Use first related concept as group
                group_name = concept.related[0]
                if group_name not in groups:
                    groups[group_name] = []
                groups[group_name].append(concept)
            else:
                groups["General"].append(concept)
        
        # Remove empty groups
        return {k: v for k, v in groups.items() if v}
