"""
Section-aware chunker.

Rules:
- Chunks follow semantic heading boundaries, not pages.
- Atomic elements (tables, code blocks, display equations, figures, images)
  are NEVER split across chunks.
- Overlap (default 0) does NOT cross section boundaries.
- Zero-token elements are handled gracefully.
- Chunks with token_count < 1 are not emitted.
"""

from __future__ import annotations

import re
import zlib as _zlib
from collections.abc import Iterable, Iterator
from typing import TYPE_CHECKING

import tiktoken

from .lang_detect import token_budget_scale as _lang_scale
from .schema import (
    ATOMIC_TYPES,
    SCHEMA_VERSION,
    BaseElement,
    ChunkMetadata,
    ChunkRecord,
    EntityAnnotation,
    chunk_id,
)

if TYPE_CHECKING:
    from .schema import ChunkQuality


def _compute_quality(
    text: str,
    section_path: list[str],
    entity_count: int,
    token_count: int,
) -> ChunkQuality:
    import re as _re

    from .schema import ChunkQuality

    words = _re.findall(r"\b[a-zA-Z]{2,}\b", text.lower())
    lexical_density = len(set(words)) / max(len(words), 1)
    entity_density = (entity_count / max(token_count, 1)) * 100
    raw = text.encode("utf-8")
    compressed = _zlib.compress(raw, level=6)
    compression_ratio = 1.0 - len(compressed) / max(len(raw), 1)
    heading_coverage = min(len(section_path) / 4.0, 1.0)
    return ChunkQuality(
        lexical_density=round(lexical_density, 4),
        entity_density=round(entity_density, 4),
        compression_ratio=round(compression_ratio, 4),
        heading_coverage=round(heading_coverage, 4),
        text_length=len(text),
    )


_SENTENCE_END = re.compile(r"(?<=[.!?])\s+(?=[A-Z])")


def _majority_lang(buffer: list[BaseElement]) -> str | None:
    """Return the most common non-None language tag in the buffer."""
    langs = [el.metadata.language for el in buffer if el.metadata.language]
    if not langs:
        return None
    return max(set(langs), key=langs.count)


def _count_tokens(text: str, enc: tiktoken.Encoding) -> int:
    return len(enc.encode(text))


def _sentence_split_tail(text: str, max_tokens: int, enc: tiktoken.Encoding) -> str:
    """Split text into sentences and return the tail that fits within max_tokens."""
    sentences = _SENTENCE_END.split(text)
    tail_sentences: list[str] = []
    token_count = 0
    for sentence in reversed(sentences):
        t = _count_tokens(sentence, enc)
        if token_count + t > max_tokens:
            break
        tail_sentences.insert(0, sentence)
        token_count += t
    return " ".join(tail_sentences)


def _build_chunk_entities(
    elements: list[BaseElement],
) -> list[EntityAnnotation] | None:
    """
    Aggregate element-level entity annotations into a chunk-level rollup.

    Offsets are rebased to chunk-relative positions: each element's start
    offset within the joined chunk text (separator="\n\n") is added to
    ann.start / ann.end. If either offset is None, chunk-level offsets are None.

    Returns:
        None  — if ALL elements have entities=None (NER not run)
        []    — if NER ran but no entities were found across all elements
        [...]  — deduplicated entity list (by entity_type + text.lower(), highest score wins)
    """
    _SEP_LEN = 2  # len("\n\n")

    any_ner_ran = any(el.entities is not None for el in elements)
    if not any_ner_ran:
        return None

    # Build offset map: element id -> char offset within joined chunk text
    offset_map: dict[str, int] = {}
    cursor = 0
    for el in elements:
        if el.text:
            offset_map[el.id] = cursor
            cursor += len(el.text) + _SEP_LEN

    # Deduplicate by (entity_type, text.lower()), keep highest score
    best: dict[tuple[str, str], EntityAnnotation] = {}
    for el in elements:
        el_offset = offset_map.get(el.id)
        for ann in el.entities or []:
            key = (ann.entity_type, ann.text.lower())
            existing = best.get(key)
            if el_offset is not None and ann.start is not None and ann.end is not None:
                chunk_start: int | None = el_offset + ann.start
                chunk_end: int | None = el_offset + ann.end
            else:
                chunk_start = None
                chunk_end = None
            if existing is None or ann.score > existing.score:
                best[key] = EntityAnnotation(
                    entity_type=ann.entity_type,
                    text=ann.text,
                    start=chunk_start,
                    end=chunk_end,
                    score=ann.score,
                    extraction_method=ann.extraction_method,
                    source_element_id=el.id,
                )

    return list(best.values())


def _build_chunk_entity_types(
    entities: list[EntityAnnotation] | None,
) -> list[str] | None:
    """
    Return sorted deduplicated list of entity_type values from entities rollup.
    Returns None if entities is None (NER not run).
    """
    if entities is None:
        return None
    return sorted({ann.entity_type for ann in entities})


def _build_chunk(
    buffer: list[BaseElement],
    source_filename: str,
    source_sha256: str,
    source_media_type: str,
    max_tokens: int,
    overlap_tokens: int,
    overlap_source: str,
    enc: tiktoken.Encoding,
    chunk_sequence: int,
    inject_context_prefix: bool = False,
    emit_quality: bool = False,
) -> ChunkRecord | None:
    """Build a ChunkRecord from a buffer of elements. Returns None if empty or zero tokens."""
    if not buffer:
        return None

    combined_text = "\n\n".join(e.text for e in buffer if e.text)
    token_count = _count_tokens(combined_text, enc)

    if token_count < 1:
        return None

    element_ids = [e.id for e in buffer]
    cid = chunk_id(element_ids, source_sha256, chunk_sequence)

    section_path = buffer[0].section_path if buffer else []
    section_path_tier = buffer[0].section_path_tier if buffer else 1

    # Context prefix injection
    context_prefix: str | None = None
    raw_text: str | None = None
    prefix_tokens = 0
    if inject_context_prefix:
        section_path_str = " > ".join(section_path) if section_path else "\u2014"
        context_prefix = f"[Source: {source_filename} | Section: {section_path_str}]\n\n"
        raw_text = combined_text
        combined_text = context_prefix + combined_text
        token_count = _count_tokens(combined_text, enc)
        prefix_tokens = len(enc.encode(context_prefix))

    chunk_entities = _build_chunk_entities(buffer)
    entity_types = _build_chunk_entity_types(chunk_entities)

    quality = (
        _compute_quality(
            combined_text,
            section_path,
            len(chunk_entities) if chunk_entities else 0,
            token_count,
        )
        if emit_quality
        else None
    )

    meta = ChunkMetadata(
        element_ids=element_ids,
        chunking_strategy="section-aware-v1",
        max_tokens=max_tokens,
        overlap_tokens=overlap_tokens,
        overlap_source=overlap_source,
        has_table=any(e.element_type.startswith("table:") for e in buffer),
        has_code=any(e.element_type == "code:block" for e in buffer),
        has_equation=any(e.element_type == "scientific:equation_display" for e in buffer),
        prefix_tokens=prefix_tokens,
        entity_types=entity_types,
    )

    return ChunkRecord(
        id=cid,
        element_type="composite:chunk",
        text=combined_text,
        section_path=section_path,
        section_path_tier=section_path_tier,
        sequence_index=chunk_sequence,
        extraction_method=buffer[0].extraction_method if buffer else "unknown",
        confidence=None,
        schema_version=SCHEMA_VERSION,
        page=buffer[0].page if buffer else None,
        source_filename=source_filename,
        source_sha256=source_sha256,
        source_media_type=source_media_type,
        token_count=token_count,
        chunk_metadata=meta,
        context_prefix=context_prefix,
        raw_text=raw_text,
        entities=chunk_entities,
        quality=quality,
    )


def chunk_elements(
    elements: Iterator[BaseElement],
    source_filename: str,
    source_sha256: str,
    source_media_type: str,
    max_tokens: int = 512,
    overlap_tokens: int = 0,
    tokenizer: str = "cl100k_base",
    inject_context_prefix: bool = False,
    emit_quality: bool = False,
) -> Iterator[ChunkRecord]:
    """
    Convert a stream of BaseElement records into ChunkRecord records.

    Chunking strategy: section-aware-v1
    - Flush buffer at heading boundaries (no overlap carried across sections)
    - Flush buffer when adding an element would exceed max_tokens
    - Atomic elements are emitted as solo chunks if they exceed max_tokens
    - Overlap only applies to token-overflow flushes, not section-boundary flushes
    """
    enc = tiktoken.get_encoding(tokenizer)

    HEADING_TYPES = {
        "structural:section_header",
        "structural:title",
    }

    buffer: list[BaseElement] = []
    overlap_tail: list[BaseElement] = []
    chunk_sequence = 0

    def flush(carry_overlap: bool) -> Iterator[ChunkRecord]:
        nonlocal chunk_sequence, overlap_tail

        if not buffer:
            return

        used_overlap = overlap_tail if carry_overlap else []
        full_buffer = used_overlap + buffer

        chunk = _build_chunk(
            full_buffer,
            source_filename,
            source_sha256,
            source_media_type,
            max_tokens,
            overlap_tokens,
            "element" if used_overlap else "none",
            enc,
            chunk_sequence,
            inject_context_prefix=inject_context_prefix,
            emit_quality=emit_quality,
        )
        if chunk is not None:
            yield chunk
            chunk_sequence += 1

        # Compute new overlap tail (only when carrying overlap is requested)
        if carry_overlap and overlap_tokens > 0:
            new_tail: list[BaseElement] = []
            tail_tokens = 0
            for elem in reversed(buffer):
                elem_tokens = _count_tokens(elem.text, enc)
                if tail_tokens + elem_tokens <= overlap_tokens:
                    new_tail.insert(0, elem)
                    tail_tokens += elem_tokens
                else:
                    # Try sentence-level split on the last text element
                    if not new_tail and elem.text:
                        tail_text = _sentence_split_tail(
                            elem.text, overlap_tokens - tail_tokens, enc
                        )
                        if tail_text:
                            # Create a synthetic overlap element (shallow copy with new text)
                            synthetic = elem.model_copy(update={"text": tail_text})
                            new_tail.insert(0, synthetic)
                    break
            overlap_tail = new_tail
        else:
            overlap_tail = []

        buffer.clear()

    for element in elements:
        elem_tokens = _count_tokens(element.text, enc) if element.text else 0
        is_atomic = element.element_type in ATOMIC_TYPES
        is_heading = element.element_type in HEADING_TYPES

        # Skip zero-token non-atomic elements by appending without token impact
        if elem_tokens == 0:
            if is_atomic:
                # Zero-token atomic: flush current buffer, skip the atomic (don't emit solo)
                yield from flush(carry_overlap=False)
            else:
                buffer.append(element)
            continue

        # Atomic element larger than max_tokens: solo chunk
        if is_atomic and elem_tokens > max_tokens:
            yield from flush(carry_overlap=False)
            # Emit atomic solo chunk
            solo_buffer = [element]
            solo = _build_chunk(
                solo_buffer,
                source_filename,
                source_sha256,
                source_media_type,
                max_tokens,
                0,
                "none",
                enc,
                chunk_sequence,
                inject_context_prefix=inject_context_prefix,
                emit_quality=emit_quality,
            )
            if solo is not None:
                yield solo
                chunk_sequence += 1
            overlap_tail = []
            continue

        # Heading: flush current buffer WITHOUT carrying overlap (section boundary)
        if is_heading and buffer:
            yield from flush(carry_overlap=False)

        # Adding this element would exceed max_tokens: flush WITH overlap, then add
        # Include overlap_tail tokens so the guard reflects the actual resulting chunk size
        current_tokens = sum(_count_tokens(e.text, enc) for e in overlap_tail if e.text) + sum(
            _count_tokens(e.text, enc) for e in buffer if e.text
        )
        effective_max = int(max_tokens * _lang_scale(_majority_lang(buffer)))
        if buffer and current_tokens + elem_tokens > effective_max:
            yield from flush(carry_overlap=overlap_tokens > 0)

        buffer.append(element)

    # Flush remaining buffer
    yield from flush(carry_overlap=False)


def chunk_elements_dual(
    elements: Iterable[BaseElement],
    source_filename: str,
    source_sha256: str,
    source_media_type: str,
    child_size: int = 128,
    parent_size: int = 512,
    overlap: int = 0,
    tokenizer: str = "cl100k_base",
    inject_context_prefix: bool = False,
    emit_quality: bool = False,
) -> Iterator[ChunkRecord]:
    """
    Dual-granularity chunking: emits paired coarse (parent) + fine (child) chunks.

    Emit order: for each parent chunk, emit the parent then its children.
    This allows streaming consumption by ParentDocumentRetriever patterns.

    If all elements fit in a single chunk at both sizes, emits granularity='single'
    for the parent and skips children.
    """
    # First pass: collect all elements into a list
    all_elements = list(elements)

    # Coarse pass
    coarse_chunks = list(
        chunk_elements(
            iter(all_elements),
            source_filename=source_filename,
            source_sha256=source_sha256,
            source_media_type=source_media_type,
            max_tokens=parent_size,
            overlap_tokens=overlap,
            tokenizer=tokenizer,
            inject_context_prefix=inject_context_prefix,
            emit_quality=emit_quality,
        )
    )

    # Check for degenerate case: only one chunk at parent size
    degenerate = len(coarse_chunks) <= 1

    for parent in coarse_chunks:
        # Upgrade parent chunk type and granularity
        parent.element_type = "composite:parent_chunk"
        if parent.chunk_metadata:
            parent.chunk_metadata.granularity = "single" if degenerate else "coarse"
            parent.chunk_metadata.parent_size = parent_size
        yield parent

        if degenerate:
            # No children needed — parent IS the only chunk
            continue

        # Find the source elements for this parent by matching element IDs
        parent_element_ids = set(parent.chunk_metadata.element_ids if parent.chunk_metadata else [])
        child_source = [el for el in all_elements if el.id in parent_element_ids]

        if not child_source:
            continue

        # Fine pass over just this parent's elements
        for child in chunk_elements(
            iter(child_source),
            source_filename=source_filename,
            source_sha256=source_sha256,
            source_media_type=source_media_type,
            max_tokens=child_size,
            overlap_tokens=overlap,
            tokenizer=tokenizer,
            inject_context_prefix=inject_context_prefix,
            emit_quality=emit_quality,
        ):
            child.parent_chunk_id = parent.id
            if child.chunk_metadata:
                child.chunk_metadata.granularity = "fine"
                child.chunk_metadata.parent_size = parent_size
            yield child
