"""Subsection-boundary chunking with overlap."""

import re

from ars25.models.statute import StatuteRecord, StatuteChunk

CHUNK_TARGET_TOKENS = 500
CHUNK_OVERLAP_CHARS = 200

# Matches subsection boundaries: "A. ", "B. ", "1. ", "2. " at start of line
SUBSECTION_PATTERN = re.compile(r"(?=\n\s*(?:[A-Z]\.|[0-9]+\.)\s)")


def estimate_tokens(text: str) -> int:
    return int(len(text.split()) * 1.3 + 0.5)


def chunk_statute(record: StatuteRecord, full_text: str) -> list[StatuteChunk]:
    """Split a statute's full text into chunks at subsection boundaries.

    Returns chunks with metadata but without vectors (embedding happens later).
    """
    chunks: list[StatuteChunk] = []

    if not full_text or not full_text.strip():
        return chunks

    # Try splitting on subsection boundaries
    parts = SUBSECTION_PATTERN.split(full_text)
    parts = [p for p in parts if p.strip()]

    if len(parts) <= 1 or estimate_tokens(full_text) <= CHUNK_TARGET_TOKENS:
        # Small enough for a single chunk
        chunks.append(StatuteChunk(
            section=record.section,
            title=record.title,
            chapter_number=record.chapter_number,
            chapter_title=record.chapter_title,
            article_number=record.article_number,
            article_title=record.article_title,
            url=record.url,
            chunk_index=0,
            chunk_text=full_text.strip(),
        ))
        return chunks

    # Group subsections into chunks targeting ~500 tokens
    current_text = ""
    chunk_index = 0

    for part in parts:
        part_tokens = estimate_tokens(part)
        current_tokens = estimate_tokens(current_text)

        if current_tokens + part_tokens > CHUNK_TARGET_TOKENS and current_text:
            # Save current chunk
            chunks.append(StatuteChunk(
                section=record.section,
                title=record.title,
                chapter_number=record.chapter_number,
                chapter_title=record.chapter_title,
                article_number=record.article_number,
                article_title=record.article_title,
                url=record.url,
                chunk_index=chunk_index,
                chunk_text=current_text.strip(),
            ))
            chunk_index += 1

            # Start new chunk with overlap from end of previous
            overlap_text = current_text[-CHUNK_OVERLAP_CHARS:]
            current_text = overlap_text + part
        else:
            current_text += part

    # Save final chunk
    if current_text.strip():
        chunks.append(StatuteChunk(
            section=record.section,
            title=record.title,
            chapter_number=record.chapter_number,
            chapter_title=record.chapter_title,
            article_number=record.article_number,
            article_title=record.article_title,
            url=record.url,
            chunk_index=chunk_index,
            chunk_text=current_text.strip(),
        ))

    return chunks
