"""Section -> chunk logic (pure; no filesystem, no embedding client).

ADR 0004: chunk unit is `(drug_id, section_key)`. A section under the token
ceiling becomes one chunk verbatim. Only the long-tail sections above it are
sub-chunked, with a sentence-boundary-aware sliding window.
"""
from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Dict, Iterable, Iterator, List, Mapping, Optional, Sequence

from ..segment.models import Monograph, SectionSpan, TableBlock
from ..tables.classify import SHAPE_FORMULA_2D
from .models import (
    CHUNK_KIND_BLOCK_DESCRIPTOR,
    CHUNK_KIND_PROSE,
    Chunk,
    ChunkAttachment,
)
from .sentences import split_sentences
from .tokens import TokenCounter, count_tokens

CEILING_TOKENS = 800
TARGET_TOKENS = 650
OVERLAP_TOKENS = 65

KIND_TABLE = "table"
KIND_FORMULA = "formula"

# A header row is only safe to embed when it is genuinely a row of labels.
# Measured on the corpus: 42 of 124 simple-table headers (34%) contain a
# digit, and AMIODARON's (physical page 183) is
# "Thời gian liệu pháp tĩnh mạch Liều 720 mg/ngày (0,5 mg/phút)" — a dose,
# inside what pdfplumber called a header, from an extraction never verified by
# eye. A label carrying no digit cannot be mistaken for a dose; a long cell is
# content rather than a label.
_DIGIT = re.compile(r"\d")
_DOSE_VALUE = re.compile(
    r"\d+(?:[.,]\d+)?\s*(?:mg|g|ml|microgam|mcg|µg|%|đơn vị|iu)\b",
    re.IGNORECASE,
)
_SUBGROUP_LABEL = re.compile(
    r"\b(?:trẻ|người lớn|người cao tuổi|sơ sinh|thiếu tháng|bệnh nhân|"
    r"suy thận|suy gan|clcr|cân nặng)\b",
    re.IGNORECASE,
)
HEADER_CELL_MAX_CHARS = 40


def _is_label_row(cells: Sequence[str]) -> bool:
    kept = [c for c in cells if c and c.strip()]
    if not kept:
        return False
    return all(
        not _DIGIT.search(cell) and len(cell.strip()) <= HEADER_CELL_MAX_CHARS
        for cell in kept
    )


def _split_trailing_label(atom: str) -> List[str]:
    """Separate a label suffix from clinical text that precedes it.

    The sentence splitter deliberately ends atoms at ``:``. In list-like dose
    prose that can yield ``"7,5 mg ... .\nBước 5:"`` as one atom. Treating the
    whole atom as the new label makes the dose at its beginning lose ``Bước 4``
    when copied across a seam. Split only at an explicit newline, sentence, or
    semicolon boundary and preserve every character exactly.
    """
    stripped = atom.rstrip()
    if not stripped.endswith(":"):
        return [atom]

    body = stripped[:-1]
    cuts = []
    newline = body.rfind("\n")
    if newline >= 0:
        cuts.append(newline + 1)
    for marker in (". ", "; "):
        position = body.rfind(marker)
        if position >= 0:
            cuts.append(position + len(marker))
    if not cuts:
        return [atom]

    cut = max(cuts)
    if not atom[cut:].strip():
        return [atom]
    return [atom[:cut], atom[cut:]]


def _label_has_embedded_content(atom: str) -> bool:
    """Whether a colon-ending atom contains content before its final label."""
    prefix = atom.rstrip()[:-1]
    return (
        "\n" in prefix
        or ". " in prefix
        or "; " in prefix
        or bool(_DOSE_VALUE.search(prefix))
    )


def _is_subgroup_label(atom: str) -> bool:
    """Population/organ-function labels that can sit under a route/indication."""
    return bool(_SUBGROUP_LABEL.search(atom))


def _atoms(text: str, measure: TokenCounter) -> List[str]:
    """Smallest units the packer may not split.

    Normally a sentence. A drug-interaction list, though, is one "sentence"
    hundreds of names long: VORICONAZOL's `tương tác thuốc` produced two parts
    of 981 and 888 tokens even after sentence packing. Left that size the
    embedding truncates them, and a truncated interaction list reads as "this
    drug is not listed" — a false negative in exactly the direction that
    matters. Such a run is comma-separated by construction, so a comma is a
    lossless place to break it.
    """
    atoms: List[str] = []
    sentences = [
        atom
        for sentence in split_sentences(text)
        for atom in _split_trailing_label(sentence)
    ]
    for sentence in sentences:
        # Split against the packing target, not the ceiling: an atom sized
        # right up to the ceiling leaves no room for the overlap prepended to
        # it, which is how a 710-token atom became a 981-token chunk.
        if measure(sentence) <= TARGET_TOKENS or "," not in sentence:
            atoms.append(sentence)
            continue
        piece = ""
        for fragment in sentence.split(","):
            candidate = f"{piece},{fragment}" if piece else fragment
            if piece and measure(candidate) > TARGET_TOKENS:
                atoms.append(piece + ",")
                piece = fragment
            else:
                piece = candidate
        if piece:
            atoms.append(piece)
    return atoms


@dataclass(frozen=True)
class _PackedPart:
    atoms: List[str]
    text: str
    source_text: str
    context_labels: List[str]


def _pack_parts(sentences: List[str], measure: TokenCounter) -> List[_PackedPart]:
    """Greedily pack sentences up to TARGET_TOKENS, overlapping by OVERLAP_TOKENS.

    A single sentence longer than the target becomes its own part rather than
    being cut mid-sentence — the caller flags it instead of splitting it.
    """
    IndexedAtom = tuple[int, str]
    PackedAtom = tuple[int, str, bool]  # index, text, retrieval-only context
    parts: List[List[PackedAtom]] = []
    current: List[PackedAtom] = []
    current_tokens = 0

    def is_label(atom: str) -> bool:
        return atom.rstrip().endswith(":")

    def contains(items: Sequence[PackedAtom], atom: IndexedAtom) -> bool:
        return any((index, text) == atom for index, text, _ in items)

    # Record the active label for every atom before packing. Looking only inside
    # the current window is insufficient: a population may span several parts,
    # so its label can have fallen outside both the 650-token buffer and the
    # 65-token overlap by the time another dose reaches a seam.
    contexts: List[IndexedAtom | None] = []
    scope_contexts: List[IndexedAtom | None] = []
    active_label: IndexedAtom | None = None
    scope_label: IndexedAtom | None = None
    for index, atom in enumerate(sentences):
        # Context at the *start* of an atom comes from the preceding label. An
        # atom may contain a dose and only end with the next label (Bisoprolol:
        # "7,5 mg ... Bước 5:"); assigning that atom to its own final label
        # makes the dose at its beginning context-free.
        contexts.append(active_label)
        next_scope = scope_label
        if is_label(atom):
            if _is_subgroup_label(atom):
                if active_label is not None and not _is_subgroup_label(active_label[1]):
                    next_scope = active_label
            else:
                next_scope = None
        # A pure subgroup label needs its parent scope when it itself lands at
        # the start of a continuation, not only when the following dose arrives.
        scope_contexts.append(next_scope)
        if is_label(atom):
            scope_label = next_scope
            active_label = (index, atom)

    def context_chain(index: int) -> List[IndexedAtom]:
        """Labels needed to make atom ``index`` independently interpretable.

        Some malformed sentence atoms contain a dose and only end with the next
        label. In that case the atom itself is context for what follows, but it
        still needs its own preceding label. Walk only until a plain label;
        this keeps the chain clinically complete without dragging the entire
        section into every overlap.
        """
        chain: List[IndexedAtom] = []

        def add(atom: IndexedAtom | None) -> None:
            if atom is not None and atom not in chain:
                chain.append(atom)

        add(scope_contexts[index])

        target = sentences[index]
        # A plain label introduces new sibling context, so it needs only its
        # retained parent scope. A compound colon-ending atom still contains
        # clinical material before that new label and therefore needs the
        # preceding active label as well.
        if is_label(target) and not _label_has_embedded_content(target):
            return chain

        nested: List[IndexedAtom] = []
        context = contexts[index]
        while context is not None and context not in nested:
            nested.append(context)
            add(scope_contexts[context[0]])
            if not _label_has_embedded_content(context[1]):
                break
            context = contexts[context[0]]
        for atom in reversed(nested):
            add(atom)
        return chain

    for index, sentence in enumerate(sentences):
        tokens = measure(sentence)
        if current and current_tokens + tokens > TARGET_TOKENS:
            # Never end a part on a label. `split_sentences` treats ':' as a
            # boundary, so "Người lớn: 500 mg mỗi 8 giờ." splits after the
            # colon; flushing there leaves a chunk ending "Người lớn:" with
            # the dose in the next one. Measured before this rule: 38 chunks,
            # AMOXICILIN's ending on a Lyme-disease indication followed by a
            # bare "Người lớn:". Outlier item 17 counted population markers on
            # 1,121 of ~1,400 monograph pages, so this is the common shape,
            # and a dose separated from the population it applies to is a
            # patient-safety defect rather than a cosmetic one.
            carried: List[PackedAtom] = []
            while current and is_label(current[-1][1]):
                carried.insert(0, current.pop())
            if not current:
                # A label is safety context, not a useful standalone retrieval
                # unit. Keep it with the following atom even when that makes a
                # synthetic pathological atom oversized; the caller will flag
                # the oversize instead of publishing an unqualified dose.
                current = carried
                current_tokens = sum(measure(atom) for _, atom, _ in current)
            else:
                parts.append(current)

                overlap: List[PackedAtom] = []
                acc = 0
                for prev_index, prev, is_context in reversed(current):
                    # Stop *before* exceeding the budget. A governing label
                    # longer than the budget is retained by itself because
                    # clinical context wins over the overlap target.
                    size = measure(prev)
                    if acc + size > OVERLAP_TOKENS:
                        break
                    overlap.insert(0, (prev_index, prev, is_context))
                    acc += size

                # If dose/detail atoms are copied, their own governing label
                # must precede them even when a newer trailing label is being
                # carried for the incoming sentence.
                if overlap:
                    prefix = [
                        (context_index, context_text, True)
                        for context_index, context_text in context_chain(overlap[0][0])
                        if not contains(overlap, (context_index, context_text))
                    ]
                    overlap = prefix + overlap
                current = overlap + carried
                current_tokens = sum(measure(atom) for _, atom, _ in current)

            # The incoming atom may belong to a label that was flushed several
            # chunks ago. Repeat that label before the new material. Do not add
            # it twice when the incoming atom is itself the label or it was
            # already carried/overlapped.
            incoming_scope = scope_contexts[index]
            for incoming_context in context_chain(index):
                if not contains(current, incoming_context):
                    packed_context = (*incoming_context, True)
                    if incoming_context == incoming_scope:
                        current.insert(0, packed_context)
                    else:
                        current.append(packed_context)
                    current_tokens += measure(incoming_context[1])

        current.append((index, sentence, False))
        current_tokens += tokens

    if current:
        parts.append(current)
    return [
        _PackedPart(
            atoms=[atom for _, atom, _ in part],
            text="".join(atom for _, atom, _ in part).strip(),
            source_text="".join(
                atom for _, atom, is_context in part if not is_context
            ).strip(),
            context_labels=[
                atom.strip() for _, atom, is_context in part if is_context
            ],
        )
        for part in parts
    ]


def _pack(sentences: List[str], measure: TokenCounter) -> List[List[str]]:
    """Compatibility surface used by focused packer tests."""
    return [part.atoms for part in _pack_parts(sentences, measure)]


def _block_kind(block: TableBlock) -> str:
    return KIND_FORMULA if block.shape == SHAPE_FORMULA_2D else KIND_TABLE


def _attachment(
    block: TableBlock,
    _header_row: List[str],
    printed_page: int | None = None,
) -> ChunkAttachment:
    return ChunkAttachment(
        block_id=block.table_id,
        kind=_block_kind(block),
        shape=block.shape,
        physical_page=block.physical_page,
        bbox=list(block.bbox),
        printed_page=printed_page,
        quarantined=block.quarantined,
        # Header extraction is not human-verified and continuation tables can
        # begin with a body row. Two ARSENIC TRIOXYD ADR rows were previously
        # shipped as "Cột:" metadata. Keep every cell value out of both the
        # retrieval text and vector payload until a reviewed logical-table
        # artifact can prove which row is a header.
        header_row=[],
    )


def _blocks_by_section(monograph: Monograph) -> Dict[str, List[TableBlock]]:
    grouped: Dict[str, List[TableBlock]] = {}
    for block in monograph.tables:
        if block.section_key:
            grouped.setdefault(block.section_key, []).append(block)
    return grouped


def describe_block(monograph: Monograph, section: SectionSpan,
                   attachment: ChunkAttachment,
                   printed_page: int | None = None) -> str:
    """Retrieval text for a block, built only from metadata.

    No cell value ever appears here. A header row is a row of labels;
    linearising it cannot invent a numeric relationship, which is exactly what
    linearising a body row does.
    """
    noun = "công thức" if attachment.kind == KIND_FORMULA else "bảng"
    page_label = (f"trang {printed_page}" if printed_page is not None
                  else "chưa xác định trang in")
    text = (f"{monograph.drug_name} — {section.display_name} — {noun}, "
            f"{page_label}.")
    text += (" Nội dung chỉ tra cứu được trên ảnh trang gốc, "
             "không trích dẫn được dưới dạng văn bản.")
    return text


def _supporting_pages(section: SectionSpan, source_text: str) -> List[int]:
    """Physical pages whose prose parts intersect one contiguous chunk span."""
    section_text = section.text.strip()
    if not source_text or section_text.count(source_text) != 1:
        raise ValueError(
            f"cannot map {section.key!r} chunk source text uniquely to its section"
        )
    chunk_start = section_text.index(source_text)
    chunk_end = chunk_start + len(source_text)

    pages: List[int] = []
    cursor = 0
    for part in section.parts:
        if part.kind != "prose" or part.quarantined or not part.text:
            continue
        part_start = section_text.find(part.text, cursor)
        if part_start < 0:
            raise ValueError(
                f"cannot map {section.key!r} part on physical page "
                f"{part.physical_page} back to section text"
            )
        part_end = part_start + len(part.text)
        cursor = part_end
        if part_start < chunk_end and part_end > chunk_start:
            pages.append(part.physical_page)

    if not pages:
        raise ValueError(f"no page provenance supports section {section.key!r} chunk")
    return sorted(set(pages))


def _page_ranges(
    physical_pages: Sequence[int],
    printed_page_map: Mapping[int, Optional[int]] | None,
    label: str,
) -> tuple[List[int], List[int]]:
    physical_range = [min(physical_pages), max(physical_pages)]
    if printed_page_map is None:
        return physical_range, []
    printed_pages = [printed_page_map.get(page) for page in physical_pages]
    if any(page is None for page in printed_pages):
        missing = [physical for physical, printed in zip(physical_pages, printed_pages, strict=True)
                   if printed is None]
        raise ValueError(
            f"cannot cite {label}: printed folio missing for physical pages {missing}"
        )
    verified = [page for page in printed_pages if page is not None]
    return physical_range, [min(verified), max(verified)]


def chunk_section(monograph: Monograph, section: SectionSpan,
                  blocks: Sequence[TableBlock] = (),
                  header_rows: Dict[str, List[str]] | None = None,
                  printed_page_map: Mapping[int, Optional[int]] | None = None,
                  measure: TokenCounter = count_tokens) -> List[Chunk]:
    header_rows = header_rows or {}
    attachments = []
    for block in blocks:
        _, attachment_printed_range = _page_ranges(
            [block.physical_page], printed_page_map,
            f"{monograph.drug_id}/{block.table_id}",
        )
        attachments.append(_attachment(
            block,
            header_rows.get(block.table_id, []),
            attachment_printed_range[0] if attachment_printed_range else None,
        ))
    quarantined = any(a.quarantined for a in attachments)

    def build(body: str, source_body: str, context_labels: List[str],
              part_index: int, part_count: int) -> Chunk:
        tokens = measure(body)
        source_page_range, printed_page_range = _page_ranges(
            _supporting_pages(section, source_body), printed_page_map,
            f"{monograph.drug_id}/{section.key}/{part_index}",
        )
        return Chunk(
            chunk_id=f"{monograph.drug_id}__{section.key}__{part_index}",
            drug_id=monograph.drug_id,
            drug_name=monograph.drug_name,
            section_key=section.key,
            section_display_name=section.display_name,
            text=body,
            source_text=source_body,
            context_labels=list(context_labels),
            heading_physical_page=section.heading.physical_page,
            source_page_range=source_page_range,
            printed_page_range=list(printed_page_range),
            atc_codes=list(monograph.atc_codes),
            part_index=part_index,
            part_count=part_count,
            est_tokens=tokens,
            oversized=tokens > CEILING_TOKENS,
            chunk_kind=CHUNK_KIND_PROSE,
            attachments=list(attachments),
            has_quarantined_content=quarantined,
        )

    text = section.text.strip()
    prose: List[Chunk] = []
    if text:
        if measure(text) <= CEILING_TOKENS:
            prose = [build(text, text, [], 0, 1)]
        else:
            parts = [part for part in _pack_parts(_atoms(text, measure), measure)
                     if part.text]
            prose = [
                build(part.text, part.source_text, part.context_labels,
                      index, len(parts))
                for index, part in enumerate(parts)
            ]

    descriptors = []
    for attachment in attachments:
        source_page_range, attachment_printed_range = _page_ranges(
            [attachment.physical_page], printed_page_map,
            f"{monograph.drug_id}/{attachment.block_id}",
        )
        printed_page = (attachment_printed_range[0]
                        if attachment_printed_range else None)
        body = describe_block(monograph, section, attachment, printed_page)
        descriptors.append(Chunk(
            chunk_id=f"{monograph.drug_id}__{section.key}__block__{attachment.block_id}",
            drug_id=monograph.drug_id,
            drug_name=monograph.drug_name,
            section_key=section.key,
            section_display_name=section.display_name,
            text=body,
            source_text=body,
            heading_physical_page=section.heading.physical_page,
            source_page_range=source_page_range,
            printed_page_range=attachment_printed_range,
            atc_codes=list(monograph.atc_codes),
            est_tokens=measure(body),
            chunk_kind=CHUNK_KIND_BLOCK_DESCRIPTOR,
            attachments=[attachment],
            has_quarantined_content=attachment.quarantined,
        ))
    return prose + descriptors


def chunk_monograph(monograph: Monograph,
                    header_rows: Dict[str, List[str]] | None = None,
                    printed_page_map: Mapping[int, Optional[int]] | None = None,
                    measure: TokenCounter = count_tokens) -> List[Chunk]:
    grouped = _blocks_by_section(monograph)
    chunks: List[Chunk] = []
    for section in monograph.sections.values():
        chunks.extend(chunk_section(monograph, section,
                                    grouped.get(section.key, ()), header_rows,
                                    printed_page_map, measure))
    return chunks


def chunk_all(monographs: Iterable[Monograph],
              header_rows: Dict[str, List[str]] | None = None,
              *,
              printed_page_map: Mapping[int, Optional[int]] | None,
              measure: TokenCounter = count_tokens) -> Iterator[Chunk]:
    if printed_page_map is None:
        raise ValueError(
            "chunk_all requires a verified printed_page_map; refusing to emit "
            "an embedding corpus without printed-page provenance"
        )
    for monograph in monographs:
        yield from chunk_monograph(monograph, header_rows, printed_page_map, measure)