"""
Prompt Chunker — Split large prompts across model context windows.

When a prompt exceeds a model's context limit, splits at natural boundaries
(file separators, section breaks, paragraphs) with configurable overlap for
context continuity. Supports a map-reduce pattern for multi-chunk processing.
"""

import re
from typing import Dict, List, Optional


# Context window limits (input tokens, leaving room for output)
MODEL_LIMITS: Dict[str, int] = {
    "gpt-4o": 100_000,        # 128K total, reserve 28K for output
    "gpt-4o-mini": 100_000,
    "gemini": 900_000,         # 1M total, reserve 100K for output
    "gemini-2.5-flash": 900_000,
    "grok": 100_000,           # 131K total, reserve 31K for output
    "grok-3-fast": 100_000,
    "qwen": 28_000,            # 32K total, reserve 4K for output
    "ollama": 6_000,            # 8K typical, reserve 2K for output
    "llama": 100_000,
    "claude": 150_000,          # 200K total, reserve 50K for output
}

# Patterns for splitting at natural boundaries (ordered by preference)
_SPLIT_PATTERNS = [
    r'\n---+\n',               # Markdown horizontal rules
    r'\n#{1,3} ',               # Markdown headers
    r'\n```\n',                 # Code block boundaries
    r'\n\[FILE: [^\]]+\]\n',    # El Gringo file markers
    r'\n\n',                    # Double newline (paragraphs)
    r'\n',                      # Single newline (last resort)
]


def estimate_tokens(text: str) -> int:
    """Estimate token count. 1 token ≈ 4 chars for English text/code."""
    return len(text) // 4 if text else 0


def get_model_limit(model: str) -> int:
    """Look up context limit for a model name (fuzzy match)."""
    model_lower = model.lower()
    # Direct match
    if model_lower in MODEL_LIMITS:
        return MODEL_LIMITS[model_lower]
    # Partial match (e.g., "gemini-coder" → "gemini")
    for key, limit in MODEL_LIMITS.items():
        if key in model_lower or model_lower in key:
            return limit
    return 16_000  # Safe default


def needs_chunking(text: str, model: str) -> bool:
    """Check if text exceeds the model's input context limit."""
    return estimate_tokens(text) > get_model_limit(model)


def chunk(
    text: str,
    model: str,
    overlap_tokens: int = 200,
    max_chunk_tokens: Optional[int] = None,
) -> List[str]:
    """Split text into chunks that fit the model's context window.

    Splits at natural boundaries with overlap for context continuity.
    Returns [text] unchanged if no chunking is needed.
    """
    limit = max_chunk_tokens or get_model_limit(model)
    if estimate_tokens(text) <= limit:
        return [text]

    # Reserve space for overlap and chunk header
    target_chars = (limit - overlap_tokens - 50) * 4
    overlap_chars = overlap_tokens * 4

    # Split into sections at natural boundaries
    sections = _split_at_boundaries(text)

    chunks: List[str] = []
    current = ""

    for section in sections:
        if len(current) + len(section) > target_chars and current:
            chunks.append(current)
            # Carry overlap from end of previous chunk
            current = current[-overlap_chars:] + section if overlap_chars else section
        else:
            current += section

    if current:
        chunks.append(current)

    # Add chunk headers for context
    if len(chunks) > 1:
        total = len(chunks)
        chunks = [
            f"[Chunk {i + 1}/{total}]\n\n{c}" for i, c in enumerate(chunks)
        ]

    return chunks


def _split_at_boundaries(text: str) -> List[str]:
    """Split text at the best available natural boundary."""
    for pattern in _SPLIT_PATTERNS:
        parts = re.split(f'({pattern})', text)
        if len(parts) > 1:
            # Rejoin split delimiters with their following section
            sections = []
            for i, part in enumerate(parts):
                if i % 2 == 0:
                    sections.append(part)
                else:
                    # Attach delimiter to next section if possible
                    if sections:
                        sections[-1] += part
            return [s for s in sections if s.strip()]
    # No boundaries found — split by character count
    chunk_size = 50_000  # ~12.5K tokens
    return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
