"""
Global outline extraction pass.

Extracts high-level document structure and themes.
"""

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

from ..schema import PassResult
from ...core.retrieval import RetrievalService
from ...core.models import LLMClient, GenerationConfig
from ...infra.logging import get_logger
from ...infra.tracing import trace_function, Span

logger = get_logger(__name__)


class GlobalOutlinePass:
    """
    Pass for extracting global document outline.
    
    Provides:
    - Chapter/section structure
    - Main themes
    - Content prioritization hints
    """
    
    PASS_NAME = "global_outline"
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        top_k_retrieval: int = 5
    ):
        """
        Initialize global outline pass.
        
        Args:
            retrieval_service: RAG retrieval service
            llm_client: LLM for generation
            top_k_retrieval: Documents to retrieve per query
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.top_k = top_k_retrieval
    
    @trace_function(name="global_outline_pass")
    def run(
        self,
        document_title: str,
        table_of_contents: Optional[List[Dict[str, Any]]] = None,
        config: Optional[GenerationConfig] = None
    ) -> PassResult:
        """
        Run the global outline extraction pass.
        
        Args:
            document_title: Title of source document
            table_of_contents: Pre-extracted TOC if available
            config: Generation configuration
            
        Returns:
            PassResult with outline data
        """
        start_time = time.perf_counter()
        config = config or GenerationConfig(max_tokens=2000, temperature=0.3)
        
        errors = []
        total_tokens = 0
        
        try:
            with Span("extract_outline"):
                # If TOC provided, use it
                if table_of_contents:
                    chapters = self._process_toc(table_of_contents)
                else:
                    # Extract structure from content
                    chapters = self._extract_structure(document_title, config)
                
                # Generate themes
                themes, tokens = self._extract_themes(document_title, config)
                total_tokens += tokens
                
                # Generate overview
                overview, tokens = self._generate_overview(document_title, chapters, themes, config)
                total_tokens += tokens
                
                output = {
                    "title": document_title,
                    "chapters": chapters,
                    "themes": themes,
                    "overview": overview,
                    "priority_sections": self._identify_priorities(chapters, themes)
                }
            
            duration_ms = (time.perf_counter() - start_time) * 1000
            
            return PassResult(
                success=True,
                pass_name=self.PASS_NAME,
                output=output,
                errors=errors,
                metrics={
                    "chapters_found": len(chapters),
                    "themes_identified": len(themes),
                },
                duration_ms=duration_ms,
                tokens_used=total_tokens
            )
            
        except Exception as e:
            logger.error(f"Global outline pass failed: {e}")
            return PassResult(
                success=False,
                pass_name=self.PASS_NAME,
                errors=[str(e)],
                duration_ms=(time.perf_counter() - start_time) * 1000
            )
    
    def _process_toc(self, toc: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Process table of contents into chapter list."""
        chapters = []
        for item in toc:
            chapters.append({
                "title": item.get("title", ""),
                "level": item.get("level", 0),
            })
        return chapters
    
    def _extract_structure(
        self,
        title: str,
        config: GenerationConfig
    ) -> List[Dict[str, Any]]:
        """Extract document structure from content."""
        results = self.retrieval_service.retrieve(
            f"table of contents chapters sections {title}",
            top_k=self.top_k
        )
        
        if not results:
            return []
        
        context = "\n\n".join(r.content for r in results)
        
        prompt = f"""Identify the main chapters or sections from this text.

TEXT:
{context}

List the main sections (one per line):"""
        
        result = self.llm_client.generate(prompt, config)
        
        chapters = []
        for line in result.text.split("\n"):
            line = line.strip().lstrip("-•*0123456789. ")
            if line and len(line) > 3:
                chapters.append({"title": line, "level": 0})
        
        return chapters[:20]  # Limit
    
    def _extract_themes(
        self,
        title: str,
        config: GenerationConfig
    ) -> tuple[List[str], int]:
        """Extract main themes from document."""
        results = self.retrieval_service.retrieve(
            f"main themes topics introduction overview {title}",
            top_k=self.top_k
        )
        
        if not results:
            return [], 0
        
        context = "\n\n".join(r.content for r in results)
        
        prompt = f"""What are the 5-7 main themes or topics covered in this text about "{title}"?

TEXT:
{context}

Main themes (one per line):"""
        
        result = self.llm_client.generate(prompt, config)
        
        themes = []
        for line in result.text.split("\n"):
            line = line.strip().lstrip("-•*0123456789. ")
            if line and len(line) > 3:
                themes.append(line)
        
        tokens = result.tokens_used if hasattr(result, 'tokens_used') else 0
        return themes[:10], tokens
    
    def _generate_overview(
        self,
        title: str,
        chapters: List[Dict[str, Any]],
        themes: List[str],
        config: GenerationConfig
    ) -> tuple[str, int]:
        """Generate document overview."""
        chapter_list = ", ".join(c["title"] for c in chapters[:10])
        theme_list = ", ".join(themes[:7])
        
        prompt = f"""Write a concise overview (3-4 sentences) for a cheat sheet about "{title}".

Chapters: {chapter_list}
Themes: {theme_list}

Overview:"""
        
        result = self.llm_client.generate(prompt, config)
        tokens = result.tokens_used if hasattr(result, 'tokens_used') else 0
        
        return result.text.strip(), tokens
    
    def _identify_priorities(
        self,
        chapters: List[Dict[str, Any]],
        themes: List[str]
    ) -> List[str]:
        """Identify priority sections for the cheat sheet."""
        # Heuristic: prioritize earlier chapters and common theme keywords
        priority_keywords = [
            "introduction", "fundamental", "basic", "core",
            "key", "main", "important", "essential"
        ]
        
        priorities = []
        for chapter in chapters[:10]:
            title_lower = chapter["title"].lower()
            for keyword in priority_keywords:
                if keyword in title_lower:
                    priorities.append(chapter["title"])
                    break
        
        # Add first few themes
        priorities.extend(themes[:3])
        
        return priorities[:10]
