"""
Dzielenie dokumentu OWU na fragmenty (chunki) z zachowaniem numeracji paragrafów.

Dlaczego nie dzielimy "co N znaków": w dokumentach OWU najważniejsza jest
możliwość zacytowania dokładnego paragrafu/ustępu. Jeśli chunk nie wie, z
jakiego "§" pochodzi, AI nie może tego podać w odpowiedzi -- a to jest
absolutny wymóg jakościowy (patrz plan-techniczny-mvp.md, sekcja 5).
"""

import re
from dataclasses import dataclass, field


@dataclass
class Chunk:
    text: str
    source_file: str
    paragraph_label: str  # np. "§ 12 ust. 1" albo "(brak numeracji)"
    chunk_index: int
    metadata: dict = field(default_factory=dict)


# Rozpoznaje nagłówki typu "§ 12", "§12 ust. 3", "Art. 5" na początku linii.
PARAGRAPH_HEADER_RE = re.compile(
    r"^(§\s*\d+[a-z]?(?:\s*ust\.?\s*\d+)?|Art\.\s*\d+[a-z]?)",
    re.IGNORECASE | re.MULTILINE,
)


def split_into_paragraphs(raw_text: str) -> list[tuple[str, str]]:
    """
    Zwraca listę (etykieta_paragrafu, tekst) na podstawie nagłówków § / Art.
    Tekst przed pierwszym nagłówkiem trafia do etykiety "(wstęp)".
    """
    matches = list(PARAGRAPH_HEADER_RE.finditer(raw_text))
    if not matches:
        return [("(brak numeracji)", raw_text.strip())]

    parts: list[tuple[str, str]] = []
    if matches[0].start() > 0:
        intro = raw_text[: matches[0].start()].strip()
        if intro:
            parts.append(("(wstęp)", intro))

    for i, m in enumerate(matches):
        start = m.start()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(raw_text)
        segment = raw_text[start:end].strip()
        label = m.group(1).strip()
        parts.append((label, segment))

    return parts


def chunk_document(
    raw_text: str,
    source_file: str,
    max_chars: int = 1200,
) -> list[Chunk]:
    """
    Dzieli dokument na chunki po granicach paragrafów. Jeśli pojedynczy
    paragraf jest bardzo długi (> max_chars), dzieli go dalej po akapitach,
    zachowując tę samą etykietę paragrafu (żeby cytowanie było wciąż poprawne).
    """
    chunks: list[Chunk] = []
    idx = 0

    for label, segment in split_into_paragraphs(raw_text):
        if len(segment) <= max_chars:
            chunks.append(Chunk(text=segment, source_file=source_file,
                                 paragraph_label=label, chunk_index=idx))
            idx += 1
            continue

        # Paragraf zbyt długi — dziel po akapitach, ta sama etykieta.
        buffer = ""
        for para in segment.split("\n\n"):
            if len(buffer) + len(para) > max_chars and buffer:
                chunks.append(Chunk(text=buffer.strip(), source_file=source_file,
                                     paragraph_label=label, chunk_index=idx))
                idx += 1
                buffer = ""
            buffer += para + "\n\n"
        if buffer.strip():
            chunks.append(Chunk(text=buffer.strip(), source_file=source_file,
                                 paragraph_label=label, chunk_index=idx))
            idx += 1

    return chunks