"""
Document chunking service for splitting documents into retrievable chunks.

Supports:
- Hierarchical chunking (chapter → section → paragraph)
- Separate streams for narrative, math, and code
- Metadata preservation for grounding
"""

import re
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple

from .ingestion import ParsedDocument, ContentBlock
from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


class ChunkType(Enum):
    """Types of content chunks."""
    NARRATIVE = "narrative"  # Regular text, explanations, definitions
    MATH = "math"            # Mathematical formulas and equations
    CODE = "code"            # Code blocks, algorithms, pseudocode
    HEADING = "heading"      # Section headings
    MIXED = "mixed"          # Contains multiple types


@dataclass
class Chunk:
    """Represents a document chunk for indexing and retrieval."""
    id: str
    content: str
    chunk_type: ChunkType
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    # Source information
    document_id: Optional[str] = None
    page_numbers: List[int] = field(default_factory=list)
    section_hierarchy: List[str] = field(default_factory=list)
    
    # Position information
    start_char: Optional[int] = None
    end_char: Optional[int] = None
    chunk_index: int = 0
    
    @classmethod
    def create(
        cls,
        content: str,
        chunk_type: ChunkType,
        **kwargs
    ) -> "Chunk":
        """Create a new chunk with auto-generated ID."""
        return cls(
            id=str(uuid.uuid4()),
            content=content,
            chunk_type=chunk_type,
            **kwargs
        )
    
    @property
    def token_estimate(self) -> int:
        """Rough estimate of token count (words * 1.3)."""
        return int(len(self.content.split()) * 1.3)
    
    @property
    def char_count(self) -> int:
        """Character count of content."""
        return len(self.content)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert chunk to dictionary for storage."""
        return {
            "id": self.id,
            "content": self.content,
            "chunk_type": self.chunk_type.value,
            "document_id": self.document_id,
            "page_numbers": self.page_numbers,
            "section_hierarchy": self.section_hierarchy,
            "chunk_index": self.chunk_index,
            **self.metadata
        }


class TextSplitter:
    """Base text splitter with configurable separators."""
    
    def __init__(
        self,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        separators: Optional[List[str]] = None,
        length_function: Optional[callable] = None
    ):
        """
        Initialize the text splitter.
        
        Args:
            chunk_size: Target size for chunks (in characters)
            chunk_overlap: Overlap between consecutive chunks
            separators: List of separators to try, in order of preference
            length_function: Function to measure text length
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.separators = separators or ["\n\n", "\n", ". ", " "]
        self.length_function = length_function or len
    
    def split_text(self, text: str) -> List[str]:
        """
        Split text into chunks.
        
        Args:
            text: Text to split
            
        Returns:
            List of text chunks
        """
        return self._split_text_recursive(text, self.separators)
    
    def _split_text_recursive(self, text: str, separators: List[str]) -> List[str]:
        """Recursively split text using available separators."""
        if not separators:
            # No separators left, split by chunk size
            return self._split_by_size(text)
        
        separator = separators[0]
        remaining_separators = separators[1:]
        
        # Split by current separator
        splits = text.split(separator)
        
        chunks = []
        current_chunk = []
        current_length = 0
        
        for split in splits:
            split_length = self.length_function(split)
            
            # If single split is too large, recursively split it
            if split_length > self.chunk_size:
                # First, add current accumulated chunk if any
                if current_chunk:
                    chunks.append(separator.join(current_chunk))
                    current_chunk = []
                    current_length = 0
                
                # Recursively split the large segment
                sub_chunks = self._split_text_recursive(split, remaining_separators)
                chunks.extend(sub_chunks)
                continue
            
            # Check if adding this split exceeds chunk size
            new_length = current_length + split_length + (len(separator) if current_chunk else 0)
            
            if new_length > self.chunk_size and current_chunk:
                # Save current chunk and start new one
                chunks.append(separator.join(current_chunk))
                
                # Apply overlap
                overlap_splits = self._get_overlap_splits(current_chunk, separator)
                current_chunk = overlap_splits + [split]
                current_length = self.length_function(separator.join(current_chunk))
            else:
                current_chunk.append(split)
                current_length = new_length
        
        # Add final chunk
        if current_chunk:
            chunks.append(separator.join(current_chunk))
        
        return [c for c in chunks if c.strip()]
    
    def _split_by_size(self, text: str) -> List[str]:
        """Split text by size when no separators work."""
        chunks = []
        start = 0
        
        while start < len(text):
            end = start + self.chunk_size
            
            # Try to break at a word boundary
            if end < len(text):
                space_idx = text.rfind(' ', start, end)
                if space_idx > start:
                    end = space_idx
            
            chunks.append(text[start:end].strip())
            start = end - self.chunk_overlap
        
        return chunks
    
    def _get_overlap_splits(self, splits: List[str], separator: str) -> List[str]:
        """Get splits for overlap from the end of current chunk."""
        if not splits or self.chunk_overlap <= 0:
            return []
        
        overlap_splits = []
        current_length = 0
        
        for split in reversed(splits):
            split_length = self.length_function(split) + len(separator)
            if current_length + split_length > self.chunk_overlap:
                break
            overlap_splits.insert(0, split)
            current_length += split_length
        
        return overlap_splits


class MathAwareTextSplitter(TextSplitter):
    """Text splitter that preserves mathematical expressions."""
    
    # Patterns for math content
    MATH_PATTERNS = [
        (r'\$\$[\s\S]*?\$\$', 'display'),           # Display math
        (r'\$[^\$]+\$', 'inline'),                   # Inline math
        (r'\\begin\{equation\}[\s\S]*?\\end\{equation\}', 'equation'),
        (r'\\begin\{align\*?\}[\s\S]*?\\end\{align\*?\}', 'align'),
        (r'\\begin\{eqnarray\*?\}[\s\S]*?\\end\{eqnarray\*?\}', 'eqnarray'),
        (r'\\begin\{gather\*?\}[\s\S]*?\\end\{gather\*?\}', 'gather'),
        (r'\\\[[\s\S]*?\\\]', 'display_bracket'),     # \[ \] math
        (r'\\\([\s\S]*?\\\)', 'inline_bracket'),     # \( \) math
    ]
    
    def __init__(self, *args, preserve_math: bool = True, **kwargs):
        super().__init__(*args, **kwargs)
        self.preserve_math = preserve_math
    
    def split_text(self, text: str) -> List[str]:
        """Split text while preserving math blocks."""
        if not self.preserve_math:
            return super().split_text(text)
        
        # Replace math blocks with placeholders
        math_blocks = {}
        protected_text = text
        
        for pattern, math_type in self.MATH_PATTERNS:
            for i, match in enumerate(re.finditer(pattern, protected_text)):
                placeholder = f"__MATH_{len(math_blocks)}__"
                math_blocks[placeholder] = match.group(0)
                protected_text = protected_text.replace(match.group(0), placeholder, 1)
        
        # Split the protected text
        chunks = super().split_text(protected_text)
        
        # Restore math blocks
        restored_chunks = []
        for chunk in chunks:
            for placeholder, math_content in math_blocks.items():
                chunk = chunk.replace(placeholder, math_content)
            restored_chunks.append(chunk)
        
        return restored_chunks


class CodeAwareTextSplitter(TextSplitter):
    """Text splitter that preserves code blocks."""
    
    CODE_PATTERNS = [
        (r'```[\s\S]*?```', 'fenced'),              # Fenced code blocks
        (r'\\begin\{lstlisting\}[\s\S]*?\\end\{lstlisting\}', 'lstlisting'),
        (r'\\begin\{verbatim\}[\s\S]*?\\end\{verbatim\}', 'verbatim'),
        (r'\\begin\{algorithm\}[\s\S]*?\\end\{algorithm\}', 'algorithm'),
        (r'\\begin\{minted\}[\s\S]*?\\end\{minted\}', 'minted'),
    ]
    
    def __init__(self, *args, preserve_code: bool = True, **kwargs):
        super().__init__(*args, **kwargs)
        self.preserve_code = preserve_code
    
    def split_text(self, text: str) -> List[str]:
        """Split text while preserving code blocks."""
        if not self.preserve_code:
            return super().split_text(text)
        
        # Replace code blocks with placeholders
        code_blocks = {}
        protected_text = text
        
        for pattern, code_type in self.CODE_PATTERNS:
            for match in re.finditer(pattern, protected_text):
                placeholder = f"__CODE_{len(code_blocks)}__"
                code_blocks[placeholder] = match.group(0)
                protected_text = protected_text.replace(match.group(0), placeholder, 1)
        
        # Split the protected text
        chunks = super().split_text(protected_text)
        
        # Restore code blocks
        restored_chunks = []
        for chunk in chunks:
            for placeholder, code_content in code_blocks.items():
                chunk = chunk.replace(placeholder, code_content)
            restored_chunks.append(chunk)
        
        return restored_chunks


class HierarchicalChunker:
    """
    Hierarchical chunking that respects document structure.
    
    Creates chunks at multiple levels:
    - Document level summaries
    - Chapter/section chunks
    - Paragraph-level chunks
    """
    
    def __init__(
        self,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        min_chunk_size: int = 100,
        max_chunk_size: int = 2000,
        preserve_math: bool = True,
        preserve_code: bool = True
    ):
        """
        Initialize the hierarchical chunker.
        
        Args:
            chunk_size: Target chunk size
            chunk_overlap: Overlap between chunks
            min_chunk_size: Minimum chunk size
            max_chunk_size: Maximum chunk size
            preserve_math: Preserve math expressions
            preserve_code: Preserve code blocks
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.min_chunk_size = min_chunk_size
        self.max_chunk_size = max_chunk_size
        
        # Create specialized splitter
        self.splitter = MathAwareTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            preserve_math=preserve_math
        )
        
        self.preserve_math = preserve_math
        self.preserve_code = preserve_code
    
    def chunk_document(self, document: ParsedDocument) -> List[Chunk]:
        """
        Chunk a parsed document hierarchically.
        
        Args:
            document: ParsedDocument to chunk
            
        Returns:
            List of Chunks with metadata
        """
        chunks = []
        
        # Group content blocks by section
        sections = self._group_by_section(document.content_blocks)
        
        # Process each section
        for section_path, blocks in sections.items():
            section_chunks = self._chunk_section(
                blocks,
                section_hierarchy=list(section_path) if section_path else [],
                document_id=document.metadata.source_path
            )
            chunks.extend(section_chunks)
        
        # Assign chunk indices
        for i, chunk in enumerate(chunks):
            chunk.chunk_index = i
        
        logger.info(f"Created {len(chunks)} chunks from document")
        return chunks
    
    def _group_by_section(
        self, 
        blocks: List[ContentBlock]
    ) -> Dict[Tuple[str, ...], List[ContentBlock]]:
        """Group content blocks by their section hierarchy."""
        sections: Dict[Tuple[str, ...], List[ContentBlock]] = {}
        
        for block in blocks:
            section_key = tuple(block.section_hierarchy) if block.section_hierarchy else ()
            if section_key not in sections:
                sections[section_key] = []
            sections[section_key].append(block)
        
        return sections
    
    def _chunk_section(
        self,
        blocks: List[ContentBlock],
        section_hierarchy: List[str],
        document_id: Optional[str]
    ) -> List[Chunk]:
        """Chunk a single section's content."""
        chunks = []
        
        # Separate blocks by type
        narrative_blocks = []
        math_blocks = []
        code_blocks = []
        
        for block in blocks:
            if block.is_math:
                math_blocks.append(block)
            elif block.is_code:
                code_blocks.append(block)
            else:
                narrative_blocks.append(block)
        
        # Chunk narrative content
        if narrative_blocks:
            narrative_text = "\n\n".join(b.content for b in narrative_blocks)
            page_numbers = list(set(
                b.page_number for b in narrative_blocks if b.page_number
            ))
            
            text_chunks = self.splitter.split_text(narrative_text)
            
            for text in text_chunks:
                if len(text) >= self.min_chunk_size:
                    chunks.append(Chunk.create(
                        content=text,
                        chunk_type=ChunkType.NARRATIVE,
                        document_id=document_id,
                        section_hierarchy=section_hierarchy,
                        page_numbers=page_numbers
                    ))
        
        # Create separate chunks for math content
        for block in math_blocks:
            chunks.append(Chunk.create(
                content=block.content,
                chunk_type=ChunkType.MATH,
                document_id=document_id,
                section_hierarchy=section_hierarchy,
                page_numbers=[block.page_number] if block.page_number else [],
                metadata=block.metadata
            ))
        
        # Create separate chunks for code content
        for block in code_blocks:
            chunks.append(Chunk.create(
                content=block.content,
                chunk_type=ChunkType.CODE,
                document_id=document_id,
                section_hierarchy=section_hierarchy,
                page_numbers=[block.page_number] if block.page_number else [],
                metadata=block.metadata
            ))
        
        return chunks


class ChunkingService:
    """
    Main chunking service for document processing.
    
    Provides different chunking strategies based on document type.
    """
    
    def __init__(self, settings=None):
        """
        Initialize the chunking service.
        
        Args:
            settings: Optional ChunkingSettings
        """
        self.settings = settings
        
        # Default settings
        chunk_size = settings.chunk_size if settings else 1000
        chunk_overlap = settings.chunk_overlap if settings else 200
        min_chunk_size = settings.min_chunk_size if settings else 100
        max_chunk_size = settings.max_chunk_size if settings else 2000
        preserve_math = settings.preserve_math if settings else True
        preserve_code = settings.preserve_code if settings else True
        
        self.hierarchical_chunker = HierarchicalChunker(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            min_chunk_size=min_chunk_size,
            max_chunk_size=max_chunk_size,
            preserve_math=preserve_math,
            preserve_code=preserve_code
        )
        
        self.text_splitter = MathAwareTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            preserve_math=preserve_math
        )
    
    @trace_function(name="chunk_document")
    def chunk(
        self,
        document: ParsedDocument,
        strategy: str = "hierarchical"
    ) -> List[Chunk]:
        """
        Chunk a document using specified strategy.
        
        Args:
            document: ParsedDocument to chunk
            strategy: Chunking strategy ("hierarchical", "simple", "by_type")
            
        Returns:
            List of Chunks
        """
        if strategy == "hierarchical":
            return self.hierarchical_chunker.chunk_document(document)
        elif strategy == "simple":
            return self._simple_chunk(document)
        elif strategy == "by_type":
            return self._chunk_by_type(document)
        else:
            raise ValueError(f"Unknown chunking strategy: {strategy}")
    
    def _simple_chunk(self, document: ParsedDocument) -> List[Chunk]:
        """Simple flat chunking of document text."""
        chunks = []
        text_chunks = self.text_splitter.split_text(document.raw_text)
        
        for i, text in enumerate(text_chunks):
            chunks.append(Chunk.create(
                content=text,
                chunk_type=ChunkType.MIXED,
                document_id=document.metadata.source_path,
                chunk_index=i
            ))
        
        return chunks
    
    def _chunk_by_type(self, document: ParsedDocument) -> List[Chunk]:
        """Chunk document with separate streams by content type."""
        chunks = []
        
        # Collect content by type
        narrative_content = []
        math_content = []
        code_content = []
        
        for block in document.content_blocks:
            if block.is_math:
                math_content.append(block)
            elif block.is_code:
                code_content.append(block)
            elif not block.is_heading:
                narrative_content.append(block)
        
        # Chunk narrative content
        narrative_text = "\n\n".join(b.content for b in narrative_content)
        for text in self.text_splitter.split_text(narrative_text):
            chunks.append(Chunk.create(
                content=text,
                chunk_type=ChunkType.NARRATIVE,
                document_id=document.metadata.source_path
            ))
        
        # Keep math blocks as individual chunks
        for block in math_content:
            chunks.append(Chunk.create(
                content=block.content,
                chunk_type=ChunkType.MATH,
                document_id=document.metadata.source_path,
                section_hierarchy=block.section_hierarchy
            ))
        
        # Keep code blocks as individual chunks
        for block in code_content:
            chunks.append(Chunk.create(
                content=block.content,
                chunk_type=ChunkType.CODE,
                document_id=document.metadata.source_path,
                section_hierarchy=block.section_hierarchy
            ))
        
        # Assign indices
        for i, chunk in enumerate(chunks):
            chunk.chunk_index = i
        
        return chunks
    
    def chunk_text(self, text: str, chunk_type: ChunkType = ChunkType.MIXED) -> List[Chunk]:
        """
        Chunk raw text.
        
        Args:
            text: Text to chunk
            chunk_type: Type to assign to chunks
            
        Returns:
            List of Chunks
        """
        chunks = []
        text_chunks = self.text_splitter.split_text(text)
        
        for i, content in enumerate(text_chunks):
            chunks.append(Chunk.create(
                content=content,
                chunk_type=chunk_type,
                chunk_index=i
            ))
        
        return chunks
