"""Document-type-specific chunk strategies.

Each strategy defines target/max sizes, overlap, and boundary-preservation
rules for a particular document type.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field


@dataclass(frozen=True)
class ChunkStrategy:
    """Configuration for how a document type should be chunked.

    Sizes are in characters (approximate 4 chars/token for English text).
    """

    target_chars: int
    max_chars: int
    overlap_chars: int
    preserve_boundaries: list[str] = field(default_factory=list)
    boundary_patterns: list[re.Pattern[str]] = field(default_factory=list)

    def __post_init__(self) -> None:
        if self.target_chars <= 0:
            raise ValueError("target_chars must be positive")
        if self.max_chars < self.target_chars:
            raise ValueError("max_chars must be >= target_chars")
        if self.overlap_chars < 0:
            raise ValueError("overlap_chars must be non-negative")


# ~4 chars per token approximation
# News / press release: 700-1000 tokens target, 100 tokens overlap
ArticleStrategy = ChunkStrategy(
    target_chars=3200,  # ~800 tokens
    max_chars=4000,  # ~1000 tokens
    overlap_chars=400,  # ~100 tokens
    preserve_boundaries=["paragraph", "heading"],
    boundary_patterns=[
        re.compile(r"\n\n+"),  # Paragraph breaks
        re.compile(r"\n#{1,6}\s"),  # Markdown headings
        re.compile(r"\n[A-Z][A-Z\s]{3,}(?:\n|$)"),  # ALL-CAPS headings
    ],
)

# Filing: 900-1300 tokens target, 150 tokens overlap
FilingStrategy = ChunkStrategy(
    target_chars=4400,  # ~1100 tokens
    max_chars=5200,  # ~1300 tokens
    overlap_chars=600,  # ~150 tokens
    preserve_boundaries=["section", "item", "heading"],
    boundary_patterns=[
        re.compile(r"\n(?:Item\s+\d+[A-Z]?[\.\:])", re.IGNORECASE),  # SEC item headers
        re.compile(r"\n(?:PART\s+[IVX]+)", re.IGNORECASE),  # Part headers
        re.compile(r"\n#{1,6}\s"),  # Markdown headings
        re.compile(r"\n[A-Z][A-Z\s]{3,}(?:\n|$)"),  # ALL-CAPS headings
        re.compile(r"\n\n+"),  # Paragraph breaks
    ],
)

# Transcript: 700-1000 tokens target, 100 tokens overlap
TranscriptStrategy = ChunkStrategy(
    target_chars=3200,  # ~800 tokens
    max_chars=4000,  # ~1000 tokens
    overlap_chars=400,  # ~100 tokens
    preserve_boundaries=["speaker", "paragraph"],
    boundary_patterns=[
        # Speaker turn patterns: "John Smith:" or "OPERATOR:" or "John Smith - CEO:"
        re.compile(r"\n(?:[A-Z][a-zA-Z\s\-\.]+(?:\s*[-–—]\s*[A-Za-z\s,]+)?)\s*:\s*"),
        re.compile(r"\n[A-Z][A-Z\s]{2,}:\s*"),  # ALL-CAPS speaker
        re.compile(r"\n\n+"),  # Paragraph breaks
    ],
)

# Macro event: 500-800 tokens target, 80 tokens overlap
MacroEventStrategy = ChunkStrategy(
    target_chars=2400,  # ~600 tokens
    max_chars=3200,  # ~800 tokens
    overlap_chars=320,  # ~80 tokens
    preserve_boundaries=["paragraph"],
    boundary_patterns=[
        re.compile(r"\n\n+"),  # Paragraph breaks
        re.compile(r"\n#{1,6}\s"),  # Markdown headings
    ],
)

# Strategy lookup by document type
STRATEGY_MAP: dict[str, ChunkStrategy] = {
    "article": ArticleStrategy,
    "news": ArticleStrategy,
    "press_release": ArticleStrategy,
    "filing": FilingStrategy,
    "transcript": TranscriptStrategy,
    "macro_event": MacroEventStrategy,
    "macro": MacroEventStrategy,
}

DEFAULT_STRATEGY = ArticleStrategy


def get_strategy(document_type: str) -> ChunkStrategy:
    """Return the appropriate ChunkStrategy for a document type."""
    return STRATEGY_MAP.get(document_type.lower(), DEFAULT_STRATEGY)