"""
Semantic chunking for large document processing.

Splits documents at semantic boundaries (sentences) with configurable
overlap for context continuity. Designed for processing 10M+ token
documents into manageable chunks.

Requirements: MEM-07 (semantic chunking)
"""

from __future__ import annotations

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


@dataclass
class Chunk:
    """A semantic chunk of a document.

    Represents a contiguous portion of a document that has been split
    at semantic boundaries. Chunks maintain token position information
    for reference and can hold embeddings for similarity search.

    Attributes:
        chunk_id: Unique identifier for this chunk
        content: The text content of the chunk
        start_token: Starting token position in original document
        end_token: Ending token position in original document
        embedding: Vector embedding for similarity search (populated later)
        summary: Compressed summary of chunk content (populated by compressor)
        importance_score: Relevance/importance score (0.0 to 1.0)
    """

    chunk_id: str
    content: str
    start_token: int
    end_token: int
    embedding: List[float] = field(default_factory=list)
    summary: str = ""
    importance_score: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        """Convert chunk to dictionary for serialization.

        Returns:
            Dictionary representation of the chunk
        """
        return {
            "chunk_id": self.chunk_id,
            "content": self.content,
            "start_token": self.start_token,
            "end_token": self.end_token,
            "embedding": self.embedding,
            "summary": self.summary,
            "importance_score": self.importance_score,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Chunk":
        """Create chunk from dictionary.

        Args:
            data: Dictionary with chunk data

        Returns:
            Chunk instance
        """
        return cls(
            chunk_id=data.get("chunk_id", ""),
            content=data.get("content", ""),
            start_token=data.get("start_token", 0),
            end_token=data.get("end_token", 0),
            embedding=data.get("embedding", []),
            summary=data.get("summary", ""),
            importance_score=data.get("importance_score", 0.0),
        )

    @property
    def token_count(self) -> int:
        """Number of tokens in this chunk.

        Returns:
            Token count (end - start position)
        """
        return self.end_token - self.start_token


class SemanticChunker:
    """Splits documents into semantic chunks with overlap.

    Chunks are created at sentence boundaries to maintain semantic coherence.
    Overlap between chunks preserves context for retrieval.

    Example:
        ```python
        chunker = SemanticChunker(chunk_size=512, chunk_overlap=50)
        chunks = chunker.chunk(large_document, "doc-001")
        print(f"Created {len(chunks)} chunks")
        ```
    """

    # Regex pattern for sentence boundaries
    # Matches: period, exclamation, question mark followed by space or end
    SENTENCE_PATTERN = re.compile(r"(?<=[.!?])\s+|(?<=[.!?])$")

    def __init__(self, chunk_size: int = 512, chunk_overlap: int = 50):
        """Initialize the semantic chunker.

        Args:
            chunk_size: Target number of tokens per chunk (default: 512)
            chunk_overlap: Number of tokens to overlap between chunks (default: 50)
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def _split_sentences(self, text: str) -> List[str]:
        """Split text into sentences at semantic boundaries.

        Uses regex to split at sentence-ending punctuation while
        preserving the punctuation with the sentence.

        Args:
            text: Text to split into sentences

        Returns:
            List of sentences (non-empty)
        """
        # Split at sentence boundaries
        parts = self.SENTENCE_PATTERN.split(text)

        # Filter empty strings and strip whitespace
        sentences = [s.strip() for s in parts if s.strip()]

        return sentences

    def _count_tokens(self, text: str) -> int:
        """Approximate token count for text.

        Uses a simple heuristic of ~4 characters per token, which
        works reasonably well for English text without requiring
        external tokenizers.

        Args:
            text: Text to count tokens for

        Returns:
            Approximate token count
        """
        # Simple approximation: ~4 chars per token
        # This avoids dependency on tiktoken or similar
        return max(1, len(text) // 4)

    def _merge_chunks(self, initial_chunks: List[Chunk], overlap: int) -> List[Chunk]:
        """Merge small chunks and add overlap for context continuity.

        Takes initial sentence-based chunks and merges adjacent ones
        that are too small, while ensuring overlap between final chunks.

        Args:
            initial_chunks: Chunks created from sentences
            overlap: Number of tokens to overlap

        Returns:
            Merged chunks with proper overlap
        """
        if not initial_chunks:
            return []

        merged = []
        current_content = []
        current_start = 0
        current_tokens = 0

        for chunk in initial_chunks:
            chunk_tokens = chunk.token_count

            # Check if adding this chunk would exceed target size
            if current_tokens + chunk_tokens > self.chunk_size and current_content:
                # Finalize current chunk
                content = " ".join(current_content)
                merged.append(
                    Chunk(
                        chunk_id=str(uuid.uuid4()),
                        content=content,
                        start_token=current_start,
                        end_token=current_start + self._count_tokens(content),
                    )
                )

                # Start new chunk with overlap
                # Take last N tokens worth of content for overlap
                if overlap > 0 and current_content:
                    overlap_text = content[
                        -overlap * 4 :
                    ]  # Approximate char count for overlap
                    current_content = [overlap_text]
                    current_start = current_start + current_tokens - overlap
                    current_tokens = self._count_tokens(overlap_text)
                else:
                    current_content = []
                    current_start = chunk.start_token
                    current_tokens = 0

            current_content.append(chunk.content)
            if not current_content or len(current_content) == 1 and current_tokens == 0:
                current_start = chunk.start_token
            current_tokens += chunk_tokens

        # Don't forget the last chunk
        if current_content:
            content = " ".join(current_content)
            merged.append(
                Chunk(
                    chunk_id=str(uuid.uuid4()),
                    content=content,
                    start_token=current_start,
                    end_token=current_start + self._count_tokens(content),
                )
            )

        return merged

    def chunk(self, content: str, document_id: str) -> List[Chunk]:
        """Split document into semantic chunks with overlap.

        Main entry point for chunking. Splits at sentence boundaries,
        then merges to achieve target chunk size with overlap.

        Args:
            content: Document text to chunk
            document_id: Unique identifier for the source document

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

        # Split into sentences first (semantic boundaries)
        sentences = self._split_sentences(content)

        if not sentences:
            return []

        # Create initial chunks from sentences
        initial_chunks = []
        token_position = 0

        for sentence in sentences:
            sentence_tokens = self._count_tokens(sentence)
            chunk = Chunk(
                chunk_id=f"{document_id}-{len(initial_chunks)}",
                content=sentence,
                start_token=token_position,
                end_token=token_position + sentence_tokens,
            )
            initial_chunks.append(chunk)
            token_position += sentence_tokens

        # Merge and add overlap
        final_chunks = self._merge_chunks(initial_chunks, self.chunk_overlap)

        # Update chunk IDs to use document_id prefix
        for i, chunk in enumerate(final_chunks):
            chunk.chunk_id = f"{document_id}-chunk-{i}"

        return final_chunks
