"""
Core concepts extraction pass.

Extracts key concepts, definitions, and relationships from the document.
"""

import time
from typing import Any, Dict, List, Optional

from ..schema import ConceptItem, CoreConceptsSection, PassResult, ImportanceLevel
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, Span

logger = get_logger(__name__)


class CoreConceptsPass:
    """
    Pass for extracting core concepts from documents.
    
    Uses RAG to identify and define key concepts with:
    - Formal definitions
    - Intuitive explanations
    - Assumptions and prerequisites
    - Relationships to other concepts
    """
    
    PASS_NAME = "core_concepts"
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        max_concepts: int = 30,
        top_k_retrieval: int = 8
    ):
        """
        Initialize core concepts pass.
        
        Args:
            retrieval_service: RAG retrieval service
            llm_client: LLM for generation
            max_concepts: Maximum concepts to extract
            top_k_retrieval: Documents to retrieve per query
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.max_concepts = max_concepts
        self.top_k = top_k_retrieval
    
    @trace_function(name="core_concepts_pass")
    def run(
        self,
        document_title: str,
        chapter_hints: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> PassResult:
        """
        Run the core concepts extraction pass.
        
        Args:
            document_title: Title of the source document
            chapter_hints: Optional chapter/section titles to focus on
            config: Generation configuration
            
        Returns:
            PassResult with CoreConceptsSection output
        """
        start_time = time.perf_counter()
        config = config or GenerationConfig(max_tokens=2048, temperature=0.3)
        
        errors = []
        all_concepts = []
        total_tokens = 0
        
        # Build queries based on document structure
        queries = self._build_queries(document_title, chapter_hints)
        
        try:
            with Span("extract_concepts", {"query_count": len(queries)}):
                for query in queries:
                    try:
                        concepts, tokens = self._extract_for_query(query, config)
                        all_concepts.extend(concepts)
                        total_tokens += tokens
                    except Exception as e:
                        logger.warning(f"Concept extraction failed for query '{query}': {e}")
                        errors.append(f"Query '{query[:50]}...': {str(e)}")
                
                # Deduplicate and rank
                unique_concepts = self._deduplicate_concepts(all_concepts)
                ranked_concepts = self._rank_concepts(unique_concepts)
                final_concepts = ranked_concepts[:self.max_concepts]
                
                # Build summary
                summary = self._generate_summary(final_concepts, document_title, config)
                
                section = CoreConceptsSection(
                    concepts=final_concepts,
                    summary=summary
                )
            
            duration_ms = (time.perf_counter() - start_time) * 1000
            
            return PassResult(
                success=True,
                pass_name=self.PASS_NAME,
                output=section,
                errors=errors,
                metrics={
                    "concepts_extracted": len(final_concepts),
                    "queries_processed": len(queries),
                    "deduplication_ratio": len(unique_concepts) / max(len(all_concepts), 1)
                },
                duration_ms=duration_ms,
                tokens_used=total_tokens
            )
            
        except Exception as e:
            logger.error(f"Core concepts pass failed: {e}")
            return PassResult(
                success=False,
                pass_name=self.PASS_NAME,
                errors=[str(e)],
                duration_ms=(time.perf_counter() - start_time) * 1000
            )
    
    def _build_queries(
        self,
        title: str,
        chapter_hints: Optional[List[str]]
    ) -> List[str]:
        """Build retrieval queries."""
        base_queries = [
            f"key concepts definitions in {title}",
            "fundamental principles and terminology",
            "important definitions and notation",
            "core ideas and main concepts",
        ]
        
        if chapter_hints:
            for chapter in chapter_hints[:5]:  # Limit chapter queries
                base_queries.append(f"concepts in {chapter}")
        
        return base_queries
    
    def _extract_for_query(
        self,
        query: str,
        config: GenerationConfig
    ) -> tuple[List[ConceptItem], int]:
        """Extract concepts for a single query."""
        # Retrieve context
        results = self.retrieval_service.retrieve_narrative(query, top_k=self.top_k)
        
        if not results:
            return [], 0
        
        context = "\n\n".join(r.content for r in results)
        
        # Build and run prompt
        prompt = get_prompt("concept_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse response
        concepts = self._parse_concepts(result.text)
        tokens = result.tokens_used if hasattr(result, 'tokens_used') else 0
        
        return concepts, tokens
    
    def _parse_concepts(self, text: str) -> List[ConceptItem]:
        """Parse LLM response into ConceptItem objects."""
        concepts = []
        current: Dict[str, Any] = {}
        
        for line in text.split("\n"):
            line = line.strip()
            
            if line.startswith("CONCEPT:"):
                if current and "name" in current:
                    concepts.append(self._build_concept(current))
                current = {"name": line[8:].strip()}
            elif line.startswith("DEFINITION:"):
                current["definition"] = line[11:].strip()
            elif line.startswith("INTUITION:"):
                current["intuition"] = line[10:].strip()
            elif line.startswith("ASSUMPTIONS:"):
                current["assumptions"] = [a.strip() for a in line[12:].split(",")]
            elif line.startswith("RELATED:"):
                current["related"] = [r.strip() for r in line[8:].split(",")]
        
        if current and "name" in current:
            concepts.append(self._build_concept(current))
        
        return concepts
    
    def _build_concept(self, data: Dict[str, Any]) -> ConceptItem:
        """Build ConceptItem from parsed data."""
        return ConceptItem(
            name=data.get("name", ""),
            definition=data.get("definition", ""),
            intuition=data.get("intuition", ""),
            assumptions=data.get("assumptions", []),
            related_concepts=data.get("related", []),
            importance=ImportanceLevel.MEDIUM
        )
    
    def _deduplicate_concepts(self, concepts: List[ConceptItem]) -> List[ConceptItem]:
        """Remove duplicate concepts."""
        seen = set()
        unique = []
        
        for concept in concepts:
            key = concept.name.lower().strip()
            if key not in seen and key:
                seen.add(key)
                unique.append(concept)
        
        return unique
    
    def _rank_concepts(self, concepts: List[ConceptItem]) -> List[ConceptItem]:
        """Rank concepts by importance."""
        def score(c: ConceptItem) -> float:
            s = 0.0
            if c.definition:
                s += 2.0
            if c.intuition:
                s += 1.5
            if c.related_concepts:
                s += 0.5 * min(len(c.related_concepts), 5)
            return s
        
        return sorted(concepts, key=score, reverse=True)
    
    def _generate_summary(
        self,
        concepts: List[ConceptItem],
        title: str,
        config: GenerationConfig
    ) -> str:
        """Generate a summary of extracted concepts."""
        if not concepts:
            return ""
        
        concept_names = ", ".join(c.name for c in concepts[:15])
        
        prompt = f"""Write a brief (2-3 sentence) summary connecting these key concepts from "{title}":

Concepts: {concept_names}

Summary:"""
        
        try:
            result = self.llm_client.generate(prompt, config)
            return result.text.strip()
        except Exception:
            return f"Key concepts from {title} including: {concept_names}"
