"""
Extractor Protocol Schema — Version 1.1.0

All element records emitted by extractor conform to these models.
The element_type field uses the canonical enum defined in ELEMENT_TYPES.
"""

from __future__ import annotations

import hashlib
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator

SCHEMA_VERSION = "1.2.0"

# ---------------------------------------------------------------------------
# Element type registry — single source of truth
# ---------------------------------------------------------------------------

ELEMENT_TYPES: frozenset[str] = frozenset(
    {
        # Structural
        "structural:title",
        "structural:subtitle",
        "structural:section_header",
        "structural:page_header",
        "structural:page_footer",
        "structural:page_number",
        "structural:divider",
        # Text
        "text:narrative",
        "text:paragraph",
        "text:caption",
        "text:footnote",
        "text:endnote",
        "text:pull_quote",
        "text:admonition",
        "text:sidebar",
        "text:abstract",
        "text:note",
        "text:transcript_segment",
        "text:uncategorized",
        # List
        "list:item",
        "list:item_ordered",
        "list:item_definition",
        # Table
        "table:simple",
        "table:complex",
        "table:continuation",
        # Code
        "code:block",
        "code:inline",
        "code:cell",
        # Media
        "media:figure",
        "media:image",
        "media:audio",
        "media:video",
        # Scientific
        "scientific:equation_display",
        "scientific:equation_inline",
        "scientific:citation",
        "scientific:reference_entry",
        "scientific:theorem",
        "scientific:proof",
        "scientific:definition",
        # Form
        "form:field",
        "form:label",
        "form:checkbox",
        # Meta
        "meta:document_title",
        "meta:author",
        "meta:date",
        "meta:address",
        "meta:email",
        "meta:url",
        "meta:page_number",
        "meta:extraction_error",
        # Composite
        "composite:chunk",
        "composite:parent_chunk",
    }
)

SUPPORTED_SCHEMA_VERSIONS: frozenset[str] = frozenset({"1.0.0", "1.1.0", "1.2.0"})

# Types that are NEVER split across chunk boundaries
ATOMIC_TYPES: frozenset[str] = frozenset(
    {
        "table:simple",
        "table:complex",
        "table:continuation",
        "media:figure",
        "media:image",
        "code:block",
        "scientific:equation_display",
    }
)


# ---------------------------------------------------------------------------
# ID derivation helpers
# ---------------------------------------------------------------------------


def element_id(
    text: str,
    page: int | None,
    sequence_index: int,
    source_sha256: str,
) -> str:
    """Derive a stable content-addressed element ID (16 hex chars = 64-bit space)."""
    raw = f"{source_sha256}:{page}:{sequence_index}:{text}"
    digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
    return f"el_{digest}"


def chunk_id(
    element_ids: list[str],
    source_sha256: str,
    sequence_index: int,
) -> str:
    """Derive a stable content-addressed chunk ID."""
    raw = f"{source_sha256}:{sequence_index}:{','.join(element_ids)}"
    digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
    return f"ch_{digest}"


def list_id(source_sha256: str, first_item_sequence_index: int) -> str:
    """Derive a list group ID from the first item's position."""
    raw = f"{source_sha256}:list:{first_item_sequence_index}"
    digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
    return f"li_{digest}"


# ---------------------------------------------------------------------------
# Sub-object models
# ---------------------------------------------------------------------------


class TableStructured(BaseModel):
    headers: list[str] = Field(default_factory=list)
    rows: list[list[Any]] = Field(default_factory=list)


class MergedCell(BaseModel):
    row: int
    col: int
    rowspan: int = 1
    colspan: int = 1
    value: str | None = None


class HeaderLevel(BaseModel):
    level: int
    labels: list[str]


class TableData(BaseModel):
    markdown: str
    structured: TableStructured
    has_header_row: bool = True
    row_count: int = 0  # data rows only (excludes header rows)
    header_row_count: int = 1  # number of header rows
    col_count: int = 0
    caption: str | None = None
    merged_continuation_ids: list[str] = Field(default_factory=list)
    merge_map: list[MergedCell] = Field(default_factory=list)
    multi_level_headers: list[HeaderLevel] = Field(default_factory=list)
    footnotes: list[str] = Field(default_factory=list)
    html: str | None = None
    table_extraction_method: str | None = None


class EquationData(BaseModel):
    latex: str
    plain_text: str | None = None
    mathml: str | None = None


class CitationData(BaseModel):
    reference_entry_id: str | None = None
    citation_text: str


class ReferenceData(BaseModel):
    authors: list[str] = Field(default_factory=list)
    year: str | None = None
    title: str | None = None
    venue: str | None = None
    doi: str | None = None
    url: str | None = None
    raw: str = ""


class FigureData(BaseModel):
    caption: str | None = None
    caption_id: str | None = None
    image_ref: str | None = None
    image_sha256: str | None = None
    ocr_text: str | None = None
    width_px: int | None = None
    height_px: int | None = None
    asset_path: str | None = None  # absolute path to exported image file
    asset_format: str | None = None  # "png", "jpg", "svg"
    context_window: str | None = None  # up to 200 chars of surrounding text


class ImageData(BaseModel):
    image_ref: str | None = None
    image_sha256: str | None = None
    format: str | None = None
    width_px: int | None = None
    height_px: int | None = None
    alt_text: str | None = None
    ocr_text: str | None = None
    storage: str = "ref"  # "ref" | "inline" | "skip"


class AdmonitionData(BaseModel):
    kind: str  # "note" | "warning" | "tip" | "caution" | "important" | "danger" | "custom"
    title: str | None = None


class TranscriptData(BaseModel):
    speaker: str | None = None
    start_time_s: float
    end_time_s: float
    word_timestamps: list[dict[str, Any]] | None = None


# ---------------------------------------------------------------------------
# Entity annotation model (schema v1.1.0)
# ---------------------------------------------------------------------------


class EntityAnnotation(BaseModel):
    """
    A single named entity span extracted by the NER step.

    Absent means: NER did not run (or element type was skipped).
    Empty list means: NER ran but found no entities above threshold.
    Never null — use absence (None on the parent field) to signal "not run".
    """

    model_config = ConfigDict(extra="forbid")

    entity_type: str
    text: str
    start: int | None = None  # char offset in element.text (element-level only)
    end: int | None = None  # char offset end (exclusive)
    score: float
    extraction_method: str = "gliner2"
    source_element_id: str | None = None  # populated on chunk-level rollup annotations

    # Coreference fields (set by CorefStep — schema v1.2.0)
    coref_cluster_id: str | None = None  # intra-doc cluster: "coref_3_a1b2c3d4"
    is_pronoun: bool = False  # True for "he", "she", "they", etc.
    is_alias: bool = False  # True for "Cook", "AAPL", "the company"

    # Canonicalization fields (set by CanonStep — schema v1.2.0)
    canonical_id: str | None = None  # e.g. "person:tim-cook"
    canonical_text: str | None = None  # e.g. "Tim Cook"
    merge_confidence: float | None = None  # 1.0=exact, 0.92=embedding, 0.85=string
    merge_decision: str | None = None  # "exact"|"rule"|"similarity"|"embedding"|"new"|"review"


class CategoryScore(BaseModel):
    model_config = ConfigDict(extra="forbid")
    label: str
    score: float


class ClassificationResult(BaseModel):
    model_config = ConfigDict(extra="forbid")
    top_label: str
    top_score: float
    scores: list[CategoryScore] | None = None
    extraction_method: str = "gliner2"


class RelationTriple(BaseModel):
    model_config = ConfigDict(extra="forbid")
    subject: str
    predicate: str
    object: str
    score: float | None = None
    subject_type: str | None = None
    object_type: str | None = None
    source_text: str | None = None
    extraction_method: str = "gliner2"


class ChunkMetadata(BaseModel):
    element_ids: list[str] = Field(default_factory=list)
    chunking_strategy: str = "section-aware-v1"
    max_tokens: int = 512
    overlap_tokens: int = 0
    overlap_source: str = "none"  # "none" | "element" | "sentence_split"
    has_table: bool = False
    has_code: bool = False
    has_equation: bool = False
    prefix_tokens: int = 0
    granularity: Literal["fine", "coarse", "single"] = "single"
    parent_size: int | None = None  # token budget of the parent pass
    # NER rollup (schema v1.1.0): sorted deduped entity_type values from chunk entities
    entity_types: list[str] | None = None


class ElementMetadata(BaseModel):
    """Optional metadata bag — parser-specific extras."""

    page_number: int | None = None
    bbox: list[float] | None = None
    style_name: str | None = None
    language: str | None = None
    is_speaker_note: bool | None = None
    merged_cell: bool | None = None
    parent_id: str | None = None
    continuation_of: str | None = None


# ---------------------------------------------------------------------------
# Base element model
# ---------------------------------------------------------------------------


class BaseElement(BaseModel):
    id: str
    element_type: str
    text: str
    section_path: list[str] = Field(default_factory=list)
    section_path_tier: int = Field(ge=1, le=4)
    sequence_index: int = Field(ge=0)
    extraction_method: str
    confidence: float | None = Field(default=None, ge=0.0, le=1.0)
    schema_version: str = SCHEMA_VERSION
    page: int | None = None
    heading_level: int | None = None  # 1-6 for headings, None otherwise

    # Provenance (strongly recommended, required on chunks)
    source_filename: str | None = None
    source_sha256: str | None = None
    source_media_type: str | None = None
    source_uri: str | None = None  # original cloud URI
    source_etag: str | None = None  # cloud ETag for dedup

    # Optional sub-objects — populated based on element_type
    table: TableData | None = None
    equation: EquationData | None = None
    citation: CitationData | None = None
    reference: ReferenceData | None = None
    figure: FigureData | None = None
    image: ImageData | None = None
    admonition: AdmonitionData | None = None
    transcript: TranscriptData | None = None

    # List grouping
    list_id: str | None = None
    list_index: int | None = None
    list_depth: int | None = None

    # Metadata bag
    metadata: ElementMetadata = Field(default_factory=ElementMetadata)

    # Entity annotations (schema v1.1.0)
    # None = NER did not run (or element type was skipped)
    # []   = NER ran, no entities found above threshold
    entities: list[EntityAnnotation] | None = None

    # GLiNER2 extended capabilities (schema v1.2.0)
    classification: ClassificationResult | None = None
    relations: list[RelationTriple] | None = None
    structured_fields: dict[str, str | int | float | bool | list[str]] | None = None

    @field_validator("element_type")
    @classmethod
    def validate_element_type(cls, v: str) -> str:
        if v not in ELEMENT_TYPES:
            raise ValueError(
                f"Unknown element_type: {v!r}. Must be one of: {sorted(ELEMENT_TYPES)}"
            )
        return v

    @field_validator("schema_version")
    @classmethod
    def validate_schema_version(cls, v: str) -> str:
        if v not in SUPPORTED_SCHEMA_VERSIONS:
            raise ValueError(
                f"Unsupported schema_version {v!r}. Supported: {sorted(SUPPORTED_SCHEMA_VERSIONS)}"
            )
        return v

    @property
    def is_heading(self) -> bool:
        """True for title, section_header, page_header elements."""
        return self.element_type in (
            "structural:title",
            "structural:subtitle",
            "structural:section_header",
            "structural:page_header",
        )

    @property
    def is_semantic_heading(self) -> bool:
        """True for semantic section headings only (excludes page_header layout artifacts)."""
        return self.element_type in (
            "structural:title",
            "structural:subtitle",
            "structural:section_header",
        )

    @property
    def is_table(self) -> bool:
        """True for any table:* element."""
        return self.element_type.startswith("table:")

    @property
    def is_atomic(self) -> bool:
        """True for element types that are never split across chunk boundaries."""
        return self.element_type in ATOMIC_TYPES

    @property
    def category(self) -> str:
        """The namespace prefix of the element type (e.g. 'structural', 'table', 'text')."""
        return self.element_type.split(":")[0]


# ---------------------------------------------------------------------------
# Chunk record
# ---------------------------------------------------------------------------


class ChunkQuality(BaseModel):
    model_config = ConfigDict(extra="forbid")
    lexical_density: float  # unique content words / total words
    entity_density: float  # entities per 100 tokens
    compression_ratio: float  # information density via zlib
    heading_coverage: float  # section_path filled fraction (max tier=4)
    text_length: int  # character count


class ChunkRecord(BaseElement):
    element_type: str = "composite:chunk"
    token_count: int = Field(ge=1)
    chunk_metadata: ChunkMetadata = Field(default_factory=ChunkMetadata)

    # Provenance required on chunks
    source_filename: str
    source_sha256: str
    source_media_type: str
    source_uri: str | None = None
    source_etag: str | None = None

    # Context prefix injection
    context_prefix: str | None = None
    raw_text: str | None = None

    # Dual-granularity: links a fine chunk to its coarse parent
    parent_chunk_id: str | None = None

    # Chunk quality scoring (optional, enabled via --quality)
    quality: ChunkQuality | None = None

    # Entity rollup (schema v1.1.0)
    # None = NER did not run; [] = NER ran, no entities found
    entities: list[EntityAnnotation] | None = None

    # GLiNER2 extended capabilities (schema v1.2.0)
    classification: ClassificationResult | None = None
    relations: list[RelationTriple] | None = None
    doc_category: str | None = None  # = classification.top_label (filter shortcut)
    doc_category_score: float | None = None  # = classification.top_score (filter shortcut)
    relation_subjects: list[str] | None = None  # rollup of all triple subjects
    relation_objects: list[str] | None = None  # rollup of all triple objects


# ---------------------------------------------------------------------------
# Envelope record (first line of JSONL)
# ---------------------------------------------------------------------------


class SourceInfo(BaseModel):
    filename: str
    path: str
    sha256: str
    size_bytes: int
    media_type: str
    subtype: str | None = None  # "digital" | "scientific" | "scanned" for PDFs
    uri: str | None = None  # cloud URI e.g. s3://bucket/key
    etag: str | None = None  # cloud ETag / content hash
    version_id: str | None = None  # S3 VersionId, GCS generation
    last_modified: str | None = None  # ISO-8601 from cloud metadata


class NERStatsRecord(BaseModel):
    """Manifest stats block for the NER step (schema v1.1.0)."""

    elements_processed: int = 0
    elements_skipped: int = 0
    elements_truncated: int = 0
    inference_errors: int = 0
    total_entities: int = 0
    duration_ms: float = 0.0
    model_id: str = ""
    tier: str = ""  # "api" or "local"


class RunConfig(BaseModel):
    mode: str  # "elements" | "chunks" | "dual-chunks"
    chunk_size: int = 512
    overlap: int = 0
    strategy: str = "fast"  # "fast" | "accurate" | "ocr"
    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
    # NER config (schema v1.1.0)
    extract_entities: bool = False
    entities_model: str = "fastino/gliner2-base-v1"
    entities_labels: list[str] | None = None
    entities_threshold: float = 0.50
    classify_labels: list[str] | None = None
    relation_types: list[str] | None = None
    extract_schema: dict[str, Any] | None = None
    extract_on_types: list[str] | None = None
    emit_quality: bool = False
    table_text_mode: str = "markdown"
    config_file: str | None = None
    figures_dir: str | None = None


class EnvelopeRecord(BaseModel):
    type: str = "envelope"
    schema_version: str = SCHEMA_VERSION
    extractor_version: str
    source: SourceInfo
    run_config: RunConfig
    created_at: str  # ISO 8601


# ---------------------------------------------------------------------------
# Manifest record
# ---------------------------------------------------------------------------


class ConfidenceStats(BaseModel):
    mean: float
    p10: float
    p90: float
    below_0_5: int = 0
    below_0_7: int = 0


class SectionPathTierStats(BaseModel):
    t1: int = Field(alias="1", default=0)
    t2: int = Field(alias="2", default=0)
    t3: int = Field(alias="3", default=0)
    t4: int = Field(alias="4", default=0)

    model_config = {"populate_by_name": True}


class TableStats(BaseModel):
    total: int = 0
    with_header: int = 0
    continuation_segments: int = 0
    continuation_pairs_merged: int = 0


class ChunkStats(BaseModel):
    total: int = 0
    min_tokens: int = 0
    max_tokens: int = 0
    mean_tokens: float = 0.0
    atomic_solo: int = 0


class StatsRecord(BaseModel):
    total_elements: int = 0
    total_chunks: int = 0
    total_tokens: int = 0
    elements_by_type: dict[str, int] = Field(default_factory=dict)
    suppressed_by_filter: int = 0
    confidence: ConfidenceStats | None = None
    section_path_tiers: dict[str, int] = Field(default_factory=dict)
    tables: TableStats = Field(default_factory=TableStats)
    chunks: ChunkStats = Field(default_factory=ChunkStats)
    ner: NERStatsRecord | None = None  # populated when NER ran (schema v1.1.0)


class ManifestWarning(BaseModel):
    code: str
    message: str
    element_id: str | None = None
    page: int | None = None


class ManifestRecord(BaseModel):
    type: str = "manifest"
    manifest_id: str
    schema_version: str = SCHEMA_VERSION
    status: str  # "complete" | "partial"
    source: SourceInfo
    run_config: RunConfig
    parser_chain: list[str] = Field(default_factory=list)
    stats: StatsRecord = Field(default_factory=StatsRecord)
    warnings: list[ManifestWarning] = Field(default_factory=list)
    errors: list[dict[str, Any]] = Field(default_factory=list)
    toc: list[dict[str, Any]] = Field(default_factory=list)
    output_files: list[str] = Field(default_factory=list)
    quality_gate_results: list[dict[str, Any]] | None = None
    created_at: str
    completed_at: str | None = None
