"""
Technical Content Chunking

Intelligent chunking for technical documents that preserves:
- Complete tables (never split)
- Equations with surrounding context
- Section hierarchy
- Paragraph boundaries

Designed for nuclear criticality safety documents where
accurate context preservation is critical.
"""

import logging
import re
import uuid
from dataclasses import dataclass
from enum import Enum

logger = logging.getLogger(__name__)


class ChunkType(Enum):
    """Type of content in a chunk."""

    TEXT = "text"
    TABLE = "table"
    EQUATION = "equation"
    FIGURE = "figure"
    CODE = "code"
    MIXED = "mixed"


@dataclass
class Chunk:
    """
    Represents a chunk of document content.

    Chunks are the basic units for vector storage and retrieval.
    Each chunk maintains full provenance to its source.
    """

    id: str
    content: str
    chunk_type: ChunkType = ChunkType.TEXT
    chunk_index: int = 0  # Position in document

    # Source tracking
    document_id: str | None = None
    doc_id: str | None = None  # Human-readable ID (e.g., "ANSI/ANS-8.1-2014")
    section_path: str = ""  # e.g., "Chapter 2 > Section 2.1 > Table 1"
    page_number: int | None = None
    page_range: tuple[int, int] | None = None  # For multi-page chunks

    # Content metadata
    has_table: bool = False
    has_equation: bool = False
    has_figure: bool = False
    token_count: int = 0

    # Overlap tracking
    overlap_start: int = 0  # Chars from previous chunk
    overlap_end: int = 0  # Chars repeated in next chunk

    def __post_init__(self):
        if not self.id:
            self.id = str(uuid.uuid4())

    @property
    def is_complete_unit(self) -> bool:
        """Check if chunk is a complete semantic unit (table, figure, etc.)."""
        return self.chunk_type in [ChunkType.TABLE, ChunkType.FIGURE, ChunkType.EQUATION]


@dataclass
class ChunkingConfig:
    """Configuration for the chunking process."""

    # Size parameters (in tokens, approximated)
    target_chunk_size: int = 512
    max_chunk_size: int = 1024
    min_chunk_size: int = 100

    # Overlap parameters
    overlap_ratio: float = 0.1  # 10% overlap
    overlap_sentences: int = 2  # Minimum sentences to overlap

    # Content handling
    preserve_tables: bool = True  # Never split tables
    preserve_equations: bool = True  # Keep equations with context
    equation_context_sentences: int = 3  # Sentences before/after equation

    # Structure handling
    respect_section_boundaries: bool = True
    respect_paragraph_boundaries: bool = True
    include_section_path: bool = True

    # Token estimation (chars per token, approximate)
    chars_per_token: float = 4.0


class TechnicalChunker:
    """
    Chunks technical documents while preserving structure.

    Designed specifically for nuclear criticality safety documents
    where tables, equations, and precise references are critical.

    Features:
    - Never splits tables
    - Keeps equations with surrounding context
    - Maintains section hierarchy
    - Respects paragraph boundaries
    - Configurable overlap for context continuity

    Example:
        chunker = TechnicalChunker()
        chunks = chunker.chunk(document)
        for chunk in chunks:
            print(f"[{chunk.section_path}] {chunk.content[:100]}...")
    """

    # Patterns for detecting special content
    TABLE_PATTERN = re.compile(r"^\|.*\|$", re.MULTILINE)
    EQUATION_PATTERN = re.compile(r"\$\$.*?\$\$|\$.*?\$|\\begin\{equation\}.*?\\end\{equation\}", re.DOTALL)
    FIGURE_PATTERN = re.compile(r"Figure\s+\d+[.:]\s*", re.IGNORECASE)
    SECTION_PATTERN = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)

    def __init__(self, config: ChunkingConfig | None = None):
        """
        Initialize the chunker.

        Args:
            config: Chunking configuration (uses defaults if not provided)
        """
        self.config = config or ChunkingConfig()

    def chunk(
        self,
        content: str,
        document_id: str | None = None,
        doc_id: str | None = None,
        metadata: dict | None = None,
    ) -> list[Chunk]:
        """
        Chunk document content.

        Args:
            content: Document text content
            document_id: Internal document UUID
            doc_id: Human-readable document ID (e.g., "NUREG-1520")
            metadata: Additional metadata to include

        Returns:
            List of Chunk objects
        """
        if not content or not content.strip():
            return []

        metadata = metadata or {}

        # Step 1: Parse document structure
        segments = self._parse_structure(content)

        # Step 2: Identify special content (tables, equations)
        segments = self._identify_special_content(segments)

        # Step 3: Create chunks respecting boundaries
        chunks = self._create_chunks(segments, document_id, doc_id)

        # Step 4: Add overlap between chunks
        if self.config.overlap_ratio > 0:
            chunks = self._add_overlap(chunks)

        logger.info(f"Created {len(chunks)} chunks from document")
        return chunks

    def _parse_structure(self, content: str) -> list[dict]:
        """Parse document into structural segments."""
        segments: list[dict] = []
        current_section_path: list[str] = []
        current_content: list[str] = []
        current_page = 1

        lines = content.split("\n")

        for line in lines:
            # Check for page markers
            page_match = re.match(r"---\s*Page\s+(\d+)\s*---", line)
            if page_match:
                # Save current content
                if current_content:
                    segments.append({
                        "content": "\n".join(current_content),
                        "section_path": " > ".join(current_section_path) or "Document",
                        "page_number": current_page,
                        "type": "text",
                    })
                    current_content = []

                current_page = int(page_match.group(1))
                continue

            # Check for section headers
            section_match = self.SECTION_PATTERN.match(line)
            if section_match:
                # Save current content
                if current_content:
                    segments.append({
                        "content": "\n".join(current_content),
                        "section_path": " > ".join(current_section_path) or "Document",
                        "page_number": current_page,
                        "type": "text",
                    })
                    current_content = []

                level = len(section_match.group(1))
                title = section_match.group(2).strip()

                # Update section path
                while len(current_section_path) >= level:
                    current_section_path.pop()
                current_section_path.append(title)

                continue

            current_content.append(line)

        # Don't forget the last segment
        if current_content:
            segments.append({
                "content": "\n".join(current_content),
                "section_path": " > ".join(current_section_path) or "Document",
                "page_number": current_page,
                "type": "text",
            })

        return segments

    def _identify_special_content(self, segments: list[dict]) -> list[dict]:
        """Identify and mark special content (tables, equations)."""
        result = []

        for segment in segments:
            content = segment["content"]

            # Check for tables (markdown format)
            if self.TABLE_PATTERN.search(content):
                # Extract tables as separate segments
                parts = self._split_tables(content, segment)
                result.extend(parts)
            elif self.EQUATION_PATTERN.search(content):
                # Mark as containing equations
                segment["has_equation"] = True
                result.append(segment)
            else:
                result.append(segment)

        return result

    def _split_tables(self, content: str, base_segment: dict) -> list[dict]:
        """Split content to extract tables as separate chunks."""
        parts: list[dict] = []
        lines = content.split("\n")

        current_text: list[str] = []
        current_table: list[str] = []
        in_table = False

        for line in lines:
            is_table_line = bool(self.TABLE_PATTERN.match(line))

            if is_table_line:
                # Starting or continuing a table
                if not in_table and current_text:
                    # Save preceding text
                    parts.append({
                        **base_segment,
                        "content": "\n".join(current_text),
                        "type": "text",
                    })
                    current_text = []

                in_table = True
                current_table.append(line)
            else:
                if in_table:
                    # End of table
                    parts.append({
                        **base_segment,
                        "content": "\n".join(current_table),
                        "type": "table",
                        "has_table": True,
                    })
                    current_table = []
                    in_table = False

                current_text.append(line)

        # Handle remaining content
        if current_table:
            parts.append({
                **base_segment,
                "content": "\n".join(current_table),
                "type": "table",
                "has_table": True,
            })
        if current_text:
            parts.append({
                **base_segment,
                "content": "\n".join(current_text),
                "type": "text",
            })

        return parts

    def _create_chunks(
        self,
        segments: list[dict],
        document_id: str | None,
        doc_id: str | None,
    ) -> list[Chunk]:
        """Create chunks from segments."""
        chunks = []
        chunk_index = 0

        for segment in segments:
            content = segment["content"].strip()
            if not content:
                continue

            segment_type = segment.get("type", "text")
            token_estimate = len(content) / self.config.chars_per_token

            # Tables are always complete chunks
            if segment_type == "table" and self.config.preserve_tables:
                chunks.append(Chunk(
                    id=str(uuid.uuid4()),
                    content=content,
                    chunk_type=ChunkType.TABLE,
                    chunk_index=chunk_index,
                    document_id=document_id,
                    doc_id=doc_id,
                    section_path=segment.get("section_path", ""),
                    page_number=segment.get("page_number"),
                    has_table=True,
                    token_count=int(token_estimate),
                ))
                chunk_index += 1
                continue

            # For text, split if too large
            if token_estimate <= self.config.max_chunk_size:
                chunk_type = ChunkType.TEXT
                if segment.get("has_equation"):
                    chunk_type = ChunkType.EQUATION

                chunks.append(Chunk(
                    id=str(uuid.uuid4()),
                    content=content,
                    chunk_type=chunk_type,
                    chunk_index=chunk_index,
                    document_id=document_id,
                    doc_id=doc_id,
                    section_path=segment.get("section_path", ""),
                    page_number=segment.get("page_number"),
                    has_equation=segment.get("has_equation", False),
                    token_count=int(token_estimate),
                ))
                chunk_index += 1
            else:
                # Split large text segments
                sub_chunks = self._split_large_segment(
                    content,
                    segment,
                    document_id,
                    doc_id,
                    chunk_index,
                )
                chunks.extend(sub_chunks)
                chunk_index += len(sub_chunks)

        return chunks

    def _split_large_segment(
        self,
        content: str,
        segment: dict,
        document_id: str | None,
        doc_id: str | None,
        start_index: int,
    ) -> list[Chunk]:
        """Split large text segment into smaller chunks."""
        chunks = []
        target_chars = int(self.config.target_chunk_size * self.config.chars_per_token)

        # Split by paragraphs first
        paragraphs = re.split(r"\n\n+", content)

        current_content: list[str] = []
        current_chars = 0
        chunk_index = start_index

        for para in paragraphs:
            para = para.strip()
            if not para:
                continue

            para_chars = len(para)

            # If single paragraph exceeds target, split by sentences
            if para_chars > target_chars:
                if current_content:
                    # Save current
                    chunk_text = "\n\n".join(current_content)
                    chunks.append(Chunk(
                        id=str(uuid.uuid4()),
                        content=chunk_text,
                        chunk_type=ChunkType.TEXT,
                        chunk_index=chunk_index,
                        document_id=document_id,
                        doc_id=doc_id,
                        section_path=segment.get("section_path", ""),
                        page_number=segment.get("page_number"),
                        token_count=int(len(chunk_text) / self.config.chars_per_token),
                    ))
                    chunk_index += 1
                    current_content = []
                    current_chars = 0

                # Split paragraph by sentences
                sentence_chunks = self._split_by_sentences(
                    para, segment, document_id, doc_id, chunk_index
                )
                chunks.extend(sentence_chunks)
                chunk_index += len(sentence_chunks)
                continue

            # Check if adding this paragraph exceeds target
            if current_chars + para_chars > target_chars and current_content:
                chunk_text = "\n\n".join(current_content)
                chunks.append(Chunk(
                    id=str(uuid.uuid4()),
                    content=chunk_text,
                    chunk_type=ChunkType.TEXT,
                    chunk_index=chunk_index,
                    document_id=document_id,
                    doc_id=doc_id,
                    section_path=segment.get("section_path", ""),
                    page_number=segment.get("page_number"),
                    token_count=int(len(chunk_text) / self.config.chars_per_token),
                ))
                chunk_index += 1
                current_content = []
                current_chars = 0

            current_content.append(para)
            current_chars += para_chars

        # Don't forget remaining content
        if current_content:
            chunk_text = "\n\n".join(current_content)
            chunks.append(Chunk(
                id=str(uuid.uuid4()),
                content=chunk_text,
                chunk_type=ChunkType.TEXT,
                chunk_index=chunk_index,
                document_id=document_id,
                doc_id=doc_id,
                section_path=segment.get("section_path", ""),
                page_number=segment.get("page_number"),
                token_count=int(len(chunk_text) / self.config.chars_per_token),
            ))

        return chunks

    def _split_by_sentences(
        self,
        text: str,
        segment: dict,
        document_id: str | None,
        doc_id: str | None,
        start_index: int,
    ) -> list[Chunk]:
        """Split text by sentences."""
        chunks = []
        target_chars = int(self.config.target_chunk_size * self.config.chars_per_token)

        # Simple sentence splitting (can be improved with NLP)
        sentences = re.split(r"(?<=[.!?])\s+", text)

        current_sentences: list[str] = []
        current_chars = 0
        chunk_index = start_index

        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue

            sent_chars = len(sentence)

            if current_chars + sent_chars > target_chars and current_sentences:
                chunk_text = " ".join(current_sentences)
                chunks.append(Chunk(
                    id=str(uuid.uuid4()),
                    content=chunk_text,
                    chunk_type=ChunkType.TEXT,
                    chunk_index=chunk_index,
                    document_id=document_id,
                    doc_id=doc_id,
                    section_path=segment.get("section_path", ""),
                    page_number=segment.get("page_number"),
                    token_count=int(len(chunk_text) / self.config.chars_per_token),
                ))
                chunk_index += 1
                current_sentences = []
                current_chars = 0

            current_sentences.append(sentence)
            current_chars += sent_chars

        if current_sentences:
            chunk_text = " ".join(current_sentences)
            chunks.append(Chunk(
                id=str(uuid.uuid4()),
                content=chunk_text,
                chunk_type=ChunkType.TEXT,
                chunk_index=chunk_index,
                document_id=document_id,
                doc_id=doc_id,
                section_path=segment.get("section_path", ""),
                page_number=segment.get("page_number"),
                token_count=int(len(chunk_text) / self.config.chars_per_token),
            ))

        return chunks

    def _add_overlap(self, chunks: list[Chunk]) -> list[Chunk]:
        """Add overlap between consecutive text chunks."""
        if len(chunks) < 2:
            return chunks

        overlap_chars = int(
            self.config.target_chunk_size
            * self.config.chars_per_token
            * self.config.overlap_ratio
        )

        for i in range(1, len(chunks)):
            prev_chunk = chunks[i - 1]
            curr_chunk = chunks[i]

            # Don't add overlap to/from complete units
            if prev_chunk.is_complete_unit or curr_chunk.is_complete_unit:
                continue

            # Get overlap text from end of previous chunk
            prev_content = prev_chunk.content
            overlap_text = prev_content[-overlap_chars:] if len(prev_content) > overlap_chars else ""

            # Add to start of current chunk
            if overlap_text:
                curr_chunk.content = overlap_text + "\n\n" + curr_chunk.content
                curr_chunk.overlap_start = len(overlap_text)
                prev_chunk.overlap_end = len(overlap_text)
                curr_chunk.token_count = int(len(curr_chunk.content) / self.config.chars_per_token)

        return chunks


def chunk_document(
    content: str,
    config: ChunkingConfig | None = None,
    **kwargs,
) -> list[Chunk]:
    """
    Convenience function to chunk a document.

    Args:
        content: Document text content
        config: Optional chunking configuration
        **kwargs: Additional arguments passed to TechnicalChunker.chunk()

    Returns:
        List of Chunk objects
    """
    chunker = TechnicalChunker(config)
    return chunker.chunk(content, **kwargs)
