"""
Main extraction pipeline. Orchestrates:
  File -> Quarantine -> Classifier -> Parser Registry -> HeadingStack -> Chunker -> Writer
"""

from __future__ import annotations

import hashlib
import logging
from collections.abc import Iterator
from pathlib import Path

from . import quarantine as quarantine_mod
from .classifier import classify
from .heading_stack import HeadingStack
from .lang_detect import detect as _detect_lang
from .parsers import register as _register  # noqa: F401
from .parsers.base import FileRecord, RawElement
from .registry import registry
from .schema import (
    ELEMENT_TYPES,
    SCHEMA_VERSION,
    BaseElement,
    ElementMetadata,
    element_id,
)
from .schema import (
    list_id as make_list_id,
)
from .table_serializers import serialize as _serialize_table
from .telemetry import trace_stage

logger = logging.getLogger(__name__)
_MAX_OPEN_TABLES = 32


def _sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def _normalize(
    raw: RawElement,
    sequence_index: int,
    source_sha256: str,
    source_filename: str,
    source_media_type: str,
    heading_stack: HeadingStack,
) -> BaseElement:
    """Convert a RawElement to a normalized BaseElement."""
    from .schema import (
        AdmonitionData,
        CitationData,
        EquationData,
        FigureData,
        ImageData,
        ReferenceData,
        TableData,
        TableStructured,
        TranscriptData,
    )

    # Update heading stack if this is a heading
    HEADING_TYPES = {
        "structural:section_header",
        "structural:title",
        "structural:page_header",
    }

    if raw.element_type in HEADING_TYPES and raw.heading_level is not None:
        if raw.element_type != "structural:page_header":
            heading_stack.push(
                level=raw.heading_level,
                text=raw.text,
                detection_tier=raw.detection_tier,
            )

    section_path = heading_stack.current_section_path()
    section_path_tier = heading_stack.current_section_path_tier()

    eid = element_id(
        text=raw.text,
        page=raw.page,
        sequence_index=sequence_index,
        source_sha256=source_sha256,
    )

    # Build sub-objects from raw.extra dicts
    table = None
    if raw.table_data:
        td = raw.table_data
        structured = td.get("structured", {})
        from .schema import MergedCell

        raw_merge_map = td.get("merge_map") or []
        merge_map = [MergedCell(**m) if isinstance(m, dict) else m for m in raw_merge_map]
        table = TableData(
            markdown=td.get("markdown", ""),
            structured=TableStructured(
                headers=structured.get("headers", []),
                rows=structured.get("rows", []),
            ),
            has_header_row=td.get("has_header_row", False),
            row_count=td.get("row_count", 0),
            header_row_count=td.get("header_row_count", 0),
            col_count=td.get("col_count", 0),
            caption=td.get("caption"),
            merge_map=merge_map,
            html=td.get("html"),
            table_extraction_method=td.get("table_extraction_method"),
        )

    equation = None
    if raw.equation_data:
        ed = raw.equation_data
        equation = EquationData(
            latex=ed.get("latex", ""),
            plain_text=ed.get("plain_text"),
            mathml=ed.get("mathml"),
        )

    citation = None
    if raw.citation_data:
        cd = raw.citation_data
        citation = CitationData(
            reference_entry_id=cd.get("reference_entry_id"),
            citation_text=cd.get("citation_text", raw.text),
        )

    reference = None
    if raw.reference_data:
        rd = raw.reference_data
        reference = ReferenceData(
            authors=rd.get("authors", []),
            year=rd.get("year"),
            title=rd.get("title"),
            venue=rd.get("venue"),
            doi=rd.get("doi"),
            url=rd.get("url"),
            raw=rd.get("raw", raw.text),
        )

    figure = None
    if raw.figure_data:
        fd = raw.figure_data
        figure = FigureData(
            caption=fd.get("caption"),
            image_ref=fd.get("image_ref"),
            image_sha256=fd.get("image_sha256"),
            ocr_text=fd.get("ocr_text"),
            width_px=fd.get("width_px"),
            height_px=fd.get("height_px"),
            asset_path=fd.get("asset_path"),
            asset_format=fd.get("asset_format"),
            context_window=fd.get("context_window"),
        )

    image = None
    if raw.image_data:
        imd = raw.image_data
        image = ImageData(
            image_ref=imd.get("image_ref"),
            image_sha256=imd.get("image_sha256"),
            format=imd.get("format"),
            width_px=imd.get("width_px"),
            height_px=imd.get("height_px"),
            alt_text=imd.get("alt_text"),
            ocr_text=imd.get("ocr_text"),
            storage=imd.get("storage", "ref"),
        )

    admonition = None
    if raw.admonition_data:
        ad = raw.admonition_data
        admonition = AdmonitionData(
            kind=ad.get("kind", "note"),
            title=ad.get("title"),
        )

    transcript = None
    if raw.transcript_data:
        trd = raw.transcript_data
        transcript = TranscriptData(
            speaker=trd.get("speaker"),
            start_time_s=trd.get("start_time_s", 0.0),
            end_time_s=trd.get("end_time_s", 0.0),
            word_timestamps=trd.get("word_timestamps"),
        )

    # List fields
    list_id_val = None
    list_index = None
    list_depth = None
    if raw.list_data:
        ld = raw.list_data
        list_index = ld.get("list_index")
        list_depth = ld.get("list_depth", 0)
        first_idx = ld.get("first_item_sequence_index", sequence_index)
        list_id_val = make_list_id(source_sha256, first_idx)

    metadata = ElementMetadata(
        page_number=raw.page,
        bbox=raw.extra.get("bbox"),
        style_name=raw.extra.get("style_name"),
        language=raw.extra.get("language"),
        is_speaker_note=raw.extra.get("is_speaker_note"),
        merged_cell=raw.extra.get("merged_cell"),
    )

    return BaseElement(
        id=eid,
        element_type=raw.element_type,
        text=raw.text,
        section_path=section_path,
        section_path_tier=section_path_tier,
        sequence_index=sequence_index,
        extraction_method=raw.extra.get("extraction_method", "unknown"),
        confidence=raw.extra.get("confidence"),
        schema_version=SCHEMA_VERSION,
        page=raw.page,
        heading_level=raw.heading_level,
        source_filename=source_filename,
        source_sha256=source_sha256,
        source_media_type=source_media_type,
        table=table,
        equation=equation,
        citation=citation,
        reference=reference,
        figure=figure,
        image=image,
        admonition=admonition,
        transcript=transcript,
        list_id=list_id_val,
        list_index=list_index,
        list_depth=list_depth,
        metadata=metadata,
    )


def _reassemble_table_continuations(
    elements: Iterator[BaseElement],
) -> Iterator[BaseElement]:
    """Merge table:continuation elements into their parent table (streaming).

    Maintains a bounded cache of recently-seen table elements. If a continuation's
    parent is within the cache, the rows are merged in-place and the continuation
    is suppressed. If the parent has already been yielded (outside the window),
    the continuation is emitted unmerged with a warning.
    """
    # Cache: table_id -> parent element (not yet yielded)
    pending: dict[str, BaseElement] = {}
    pending_order: list[str] = []  # insertion order for eviction

    for el in elements:
        if el.element_type == "table:continuation":
            parent_id = el.metadata.continuation_of if el.metadata else None
            if parent_id and parent_id in pending:
                # Parent is still in our cache — merge rows
                parent = pending[parent_id]
                if parent.table and el.table:
                    merged_rows = list(parent.table.structured.rows) + list(
                        el.table.structured.rows
                    )
                    headers = parent.table.structured.headers
                    md_lines = []
                    if headers:
                        md_lines.append("| " + " | ".join(headers) + " |")
                        md_lines.append("|" + "|".join(" --- " for _ in headers) + "|")
                    for row in merged_rows:
                        padded = list(row) + [""] * max(0, len(headers) - len(row))
                        md_lines.append("| " + " | ".join(str(c) for c in padded) + " |")
                    new_markdown = "\n".join(md_lines)

                    existing_merged = list(parent.table.merged_continuation_ids or [])
                    existing_merged.append(el.id)

                    from .schema import TableData, TableStructured

                    new_table = TableData(
                        markdown=new_markdown,
                        structured=TableStructured(
                            headers=list(headers),
                            rows=merged_rows,
                        ),
                        has_header_row=parent.table.has_header_row,
                        row_count=len(merged_rows),
                        header_row_count=parent.table.header_row_count,
                        col_count=parent.table.col_count,
                        caption=parent.table.caption,
                        merged_continuation_ids=existing_merged,
                    )
                    pending[parent_id] = parent.model_copy(
                        update={"table": new_table, "text": new_markdown or parent.text}
                    )
                    # Don't yield the continuation — it's merged
                    continue
                else:
                    # Tables lack structured data — emit unmerged
                    logger.warning(
                        "table:continuation %s could not be merged into parent %s (missing table data)",
                        el.id,
                        parent_id,
                    )
            else:
                if parent_id:
                    logger.warning(
                        "table:continuation %s: parent %s not in cache (already yielded or missing)",
                        el.id,
                        parent_id,
                    )
            # Emit the continuation unmerged
            yield el
            continue

        # For table:simple and table:complex, hold in pending cache
        if el.element_type in ("table:simple", "table:complex"):
            # Evict oldest if cache is full
            while len(pending_order) >= _MAX_OPEN_TABLES:
                oldest_id = pending_order.pop(0)
                yield pending.pop(oldest_id)
            pending[el.id] = el
            pending_order.append(el.id)
            continue

        # For all other elements, yield directly — tables stay pending until evicted or end-of-stream
        yield el

    # Flush remaining pending tables in insertion order
    for tid in pending_order:
        yield pending[tid]


def run(
    path: Path,
    mode: str = "elements",
    chunk_size: int = 512,
    overlap: int = 0,
    strategy: str = "fast",
    include_types: list[str] | None = None,
    exclude_types: list[str] | None = None,
    tokenizer: str = "cl100k_base",
    inject_context_prefix: bool = False,
    parent_size: int = 512,
    child_size: int = 128,
    entities: bool = False,
    entities_model: str = "fastino/gliner2-base-v1",
    entities_types: list[str] | None = None,
    entities_threshold: float = 0.50,
    ner_warning_callback: callable[[str, str], None] | None = None,  # type: ignore[valid-type]
    ner_step_out: list[object] | None = None,  # if provided, ner_step is appended after creation
    canonicalize: bool = False,
    registry_path: str | None = None,
    use_embeddings: bool = False,
    classify_labels: list[str] | None = None,
    relation_types: list[str] | None = None,
    extract_schema: dict[str, object] | None = None,
    extract_on_types: list[str] | None = None,
    table_text_mode: str = "markdown",
    emit_quality: bool = False,
    figures_dir: Path | None = None,
) -> Iterator[BaseElement]:
    """Run the full extraction pipeline and yield elements or chunks."""
    from .chunker import chunk_elements, chunk_elements_dual

    _VALID_MODES = ("elements", "chunks", "dual-chunks")
    if mode not in _VALID_MODES:
        raise ValueError(f"Invalid mode {mode!r}. Expected one of: {', '.join(_VALID_MODES)}")
    _VALID_STRATEGIES = ("fast", "accurate", "ocr")
    if strategy not in _VALID_STRATEGIES:
        raise ValueError(
            f"Invalid strategy {strategy!r}. Expected one of: {', '.join(_VALID_STRATEGIES)}"
        )
    _all_prefixes = {t.split(":")[0] for t in ELEMENT_TYPES}
    for filter_name, type_list in [
        ("include_types", include_types),
        ("exclude_types", exclude_types),
    ]:
        if type_list:
            for t in type_list:
                base = t.rstrip("*").rstrip(":")
                if t not in ELEMENT_TYPES and base not in _all_prefixes:
                    valid_sample = sorted(ELEMENT_TYPES)[:8]
                    raise ValueError(
                        f"{filter_name}: unknown element type {t!r}. "
                        f"Valid types include: {valid_sample} ..."
                    )

    source_sha256 = _sha256_file(path)
    source_filename = path.name

    detected_mime = quarantine_mod.run(path)
    mime_type, subtype = classify(path, detected_mime)
    source_media_type = mime_type

    file_record = FileRecord(
        path=path,
        sha256=source_sha256,
        size_bytes=path.stat().st_size,
        mime_type=mime_type,
        subtype=subtype,
    )

    parser = registry.get(mime_type, subtype)
    heading_stack = HeadingStack(detection_tier=1)

    # Prepare figures_dir if provided
    _figures_dir: Path | None = None
    if figures_dir is not None:
        _figures_dir = Path(figures_dir)
        _figures_dir.mkdir(parents=True, exist_ok=True)

    def _normalize_stream() -> Iterator[BaseElement]:
        seq = 0
        if _figures_dir is not None and hasattr(parser, "parse_with_figures"):
            _raw_iter = parser.parse_with_figures(file_record, figures_dir=_figures_dir)
        else:
            _raw_iter = parser.parse(file_record)
        with trace_stage("parse", format=mime_type):
            for raw in _raw_iter:
                # Apply type filters before normalization to prevent heading stack mutation
                if include_types is not None and not any(
                    raw.element_type == t or raw.element_type.startswith(t.rstrip("*"))
                    for t in include_types
                ):
                    seq += 1
                    continue
                if exclude_types is not None and any(
                    raw.element_type == t or raw.element_type.startswith(t.rstrip("*"))
                    for t in exclude_types
                ):
                    seq += 1
                    continue
                element = _normalize(
                    raw, seq, source_sha256, source_filename, source_media_type, heading_stack
                )
                seq += 1
                # Apply table text serialization
                if element.table is not None and table_text_mode != "markdown":
                    element = element.model_copy(
                        update={"text": _serialize_table(element.table, table_text_mode)}
                    )
                # Auto-detect language if not already set
                if element.metadata.language is None and len(element.text or "") >= 50:
                    lang = _detect_lang(element.text)
                    if lang is not None:
                        element = element.model_copy(
                            update={
                                "metadata": element.metadata.model_copy(update={"language": lang})
                            }
                        )
                yield element

    reassembled = _reassemble_table_continuations(_normalize_stream())

    # NER annotation step (runs on all elements before chunking)
    if entities:
        try:
            from .ner import NERStep

            ner_step = NERStep(
                labels=entities_types,
                threshold=entities_threshold,
                model_id=entities_model,
            )
            if ner_step_out is not None:
                ner_step_out.append(ner_step)
            with trace_stage("ner", model=entities_model):
                reassembled = ner_step.annotate(reassembled)
        except ImportError:
            import sys as _sys

            print(
                "extractor: NER skipped — gliner2 not installed. "
                "Run: pip install 'coresdk-extractor[ner]'",
                file=_sys.stderr,
            )
            if ner_warning_callback is not None:
                ner_warning_callback(
                    "W_NER_MODEL_UNAVAILABLE",
                    "GLiNER2 not available (ImportError). Install extractor[ner].",
                )

    # Canon step
    if canonicalize and entities and registry_path:
        try:
            from pathlib import Path as _Path

            from .canon import CanonStep

            _canon_step = CanonStep(
                registry_path=_Path(registry_path),
                enabled=True,
                use_embeddings=use_embeddings,
            )
            _elements_list = list(reassembled) if not isinstance(reassembled, list) else reassembled
            _elements_list = _canon_step.canonicalize(
                _elements_list,
                doc_id=source_filename,
                doc_sha256=source_sha256,
            )
            _canon_step.flush()
            reassembled = iter(_elements_list)
        except ImportError:
            pass

    # Structured extraction step (pre-chunking, targets specific element types)
    if extract_schema:
        try:
            from .structured_extract import StructuredExtractStep

            _struct_step = StructuredExtractStep(
                schema=extract_schema,
                on_types=set(extract_on_types) if extract_on_types else None,
                model_id=entities_model,
            )
            reassembled = _struct_step.annotate(reassembled)
        except ImportError:
            import sys as _sys

            print(
                "extractor: StructuredExtractStep skipped — gliner2 not installed.",
                file=_sys.stderr,
            )

    def _apply_post_chunk_steps(chunk_iter: Iterator[BaseElement]) -> Iterator[BaseElement]:
        """Wrap a chunk iterator with classification and relation steps (post-chunking)."""
        out: Iterator[BaseElement] = chunk_iter

        if classify_labels:
            try:
                from .classify import ClassifyStep

                _classify_step = ClassifyStep(labels=classify_labels, model_id=entities_model)
                out = _classify_step.annotate(out)
            except ImportError:
                import sys as _sys

                print(
                    "extractor: ClassifyStep skipped — gliner2 not installed.",
                    file=_sys.stderr,
                )

        if relation_types:
            try:
                from .relations import RelationStep

                _relation_step = RelationStep(
                    relation_types=relation_types, model_id=entities_model
                )
                out = _relation_step.annotate(out)
            except ImportError:
                import sys as _sys

                print(
                    "extractor: RelationStep skipped — gliner2 not installed.",
                    file=_sys.stderr,
                )

        yield from out

    if mode == "chunks":
        with trace_stage("chunk", mode=mode):
            yield from _apply_post_chunk_steps(
                chunk_elements(
                    elements=reassembled,
                    source_filename=source_filename,
                    source_sha256=source_sha256,
                    source_media_type=source_media_type,
                    max_tokens=chunk_size,
                    overlap_tokens=overlap,
                    tokenizer=tokenizer,
                    inject_context_prefix=inject_context_prefix,
                    emit_quality=emit_quality,
                )
            )
    elif mode == "dual-chunks":
        with trace_stage("chunk", mode=mode):
            yield from _apply_post_chunk_steps(
                chunk_elements_dual(
                    elements=reassembled,
                    source_filename=source_filename,
                    source_sha256=source_sha256,
                    source_media_type=source_media_type,
                    child_size=child_size,
                    parent_size=parent_size,
                    overlap=overlap,
                    tokenizer=tokenizer,
                    inject_context_prefix=inject_context_prefix,
                    emit_quality=emit_quality,
                )
            )
    else:
        yield from reassembled
