from __future__ import annotations

from dataclasses import dataclass

from app.rag.extractors import ExtractedBlock


@dataclass(frozen=True)
class ChunkMetadata:
    content: str
    page_start: int | None = None
    page_end: int | None = None
    section_title: str = ""
    section_path: str = ""
    paragraph_index: int | None = None
    source_ref: str = ""


def chunk_text(text: str, *, chunk_size: int = 800, overlap: int = 120) -> list[str]:
    cleaned = (text or "").strip()
    if not cleaned:
        return []
    if len(cleaned) <= chunk_size:
        return [cleaned]
    chunks: list[str] = []
    start = 0
    while start < len(cleaned):
        end = min(len(cleaned), start + chunk_size)
        piece = cleaned[start:end].strip()
        if piece:
            chunks.append(piece)
        if end >= len(cleaned):
            break
        start = max(0, end - overlap)
    return chunks


def chunk_blocks(
    blocks: list[ExtractedBlock],
    *,
    chunk_size: int = 1400,
    overlap_blocks: int = 1,
) -> list[ChunkMetadata]:
    normalized = [block for block in blocks if block.text.strip()]
    if not normalized:
        return []

    chunks: list[ChunkMetadata] = []
    current: list[ExtractedBlock] = []
    current_len = 0

    for block in normalized:
        block_len = len(block.text)
        if current and current_len + block_len + 2 > chunk_size:
            chunks.append(_build_chunk(current))
            current = current[-overlap_blocks:] if overlap_blocks > 0 else []
            current_len = sum(len(item.text) + 2 for item in current)

        if block_len > chunk_size:
            if current:
                chunks.append(_build_chunk(current))
                current = []
                current_len = 0
            chunks.extend(_split_large_block(block, chunk_size=chunk_size))
            continue

        current.append(block)
        current_len += block_len + 2

    if current:
        chunks.append(_build_chunk(current))

    return chunks


def _build_chunk(blocks: list[ExtractedBlock]) -> ChunkMetadata:
    page_values = [block.page_start for block in blocks if block.page_start is not None]
    page_end_values = [block.page_end or block.page_start for block in blocks if block.page_start is not None]
    section_path = next((block.section_path for block in reversed(blocks) if block.section_path), "")
    section_title = next((block.section_title for block in reversed(blocks) if block.section_title), "")
    paragraph_values = [block.paragraph_index for block in blocks if block.paragraph_index is not None]
    page_start = min(page_values) if page_values else None
    page_end = max(page_end_values) if page_end_values else None
    source_ref = _format_source_ref(page_start, page_end, section_path or section_title)

    return ChunkMetadata(
        content="\n\n".join(block.text.strip() for block in blocks if block.text.strip()),
        page_start=page_start,
        page_end=page_end,
        section_title=section_title,
        section_path=section_path,
        paragraph_index=min(paragraph_values) if paragraph_values else None,
        source_ref=source_ref,
    )


def _split_large_block(block: ExtractedBlock, *, chunk_size: int) -> list[ChunkMetadata]:
    pieces = chunk_text(block.text, chunk_size=chunk_size, overlap=120)
    out: list[ChunkMetadata] = []
    for piece in pieces:
        out.append(
            ChunkMetadata(
                content=piece,
                page_start=block.page_start,
                page_end=block.page_end,
                section_title=block.section_title,
                section_path=block.section_path,
                paragraph_index=block.paragraph_index,
                source_ref=_format_source_ref(
                    block.page_start,
                    block.page_end,
                    block.section_path or block.section_title,
                ),
            )
        )
    return out


def _format_source_ref(page_start: int | None, page_end: int | None, section_path: str) -> str:
    refs: list[str] = []
    if page_start is not None:
        if page_end and page_end != page_start:
            refs.append(f"стр. {page_start}-{page_end}")
        else:
            refs.append(f"стр. {page_start}")
    if section_path:
        refs.append(section_path)
    return ", ".join(refs)