"""Claim-pack source.

Ingests Grok-4-style companion JSON bundles (axiom + theorem packs) into
arborist as one Document per record. Each pillar holds an `axioms` or
`theorems` array; each entry is dual-threaded (`delta` LaTeX symbolic +
`nablaVerbose` prose expansion) and carries explicit citations to its
classical source (Mendelson, Enderton, Hilbert, …).

Grain: one Document per axiom / theorem record. Pillar-level
`provenance.references` arrays (cross-bundle pointers like
`theoremg4.json:pillar.I.logic.excludedMiddle`) become outbound `Edge`
objects on the first record of the referencing pillar.

Lenient JSON: bundles arrive wrapped in markdown ``` ```json ``` ``` fences and may
contain unescaped LaTeX backslashes (`\\Theta`, `\\Sigma`, `\\heart`). The
parser strips fences and double-escapes lone backslashes before
delegating to ``json.loads``. A genuinely malformed bundle raises;
silently returning zero docs would be a footgun.

See ticket #000029 for the full design (kind=surface choice, hard
constraints, cross-reference grammar).
"""

from __future__ import annotations

import json
import re
import unicodedata
from pathlib import Path
from typing import Iterator

from arborist.document import Document, Edge
from arborist.source import Source


_FENCE_OPEN = re.compile(r"^\s*```(?:json)?\s*", re.IGNORECASE)
_FENCE_CLOSE = re.compile(r"\s*```\s*$")

# JSON's legal single-char escapes after a backslash.
_LEGAL_ESCAPE = frozenset('\\"/bfnrtu')


def _escape_lone_backslashes(text: str) -> str:
    """Double up backslashes that aren't already part of a JSON escape.

    Walks left-to-right. When a backslash is followed by a legal escape
    character (one of ``\\ " / b f n r t u``), pass both through as-is —
    they form a single legal JSON escape sequence. When a backslash is
    followed by anything else (LaTeX `\\Theta`, `\\heart`, `\\vec`), double
    it so the JSON parser sees a literal backslash + letter.

    This handles both already-correct double-escaped strings and lone-
    backslash LaTeX without corrupting the former.
    """
    out: list[str] = []
    i = 0
    n = len(text)
    while i < n:
        c = text[i]
        if c == "\\" and i + 1 < n and text[i + 1] in _LEGAL_ESCAPE:
            out.append(text[i:i + 2])
            i += 2
        elif c == "\\":
            out.append("\\\\")
            i += 1
        else:
            out.append(c)
            i += 1
    return "".join(out)


def _parse_lenient(raw: str) -> dict:
    """Parse a JSON bundle that may carry markdown fences + LaTeX backslashes.

    Strategy:
      1. Strip a leading ```json (or ```) fence and trailing ``` fence.
      2. Double every lone backslash so unescaped LaTeX (`\\Theta`) becomes
         legal JSON. Already-escaped pairs (`\\\\to`) pass through.
      3. Delegate to ``json.loads``.

    Raises ``json.JSONDecodeError`` if the bundle still won't parse — a
    silently-empty result would be worse than a loud failure.
    """
    text = raw
    text = _FENCE_OPEN.sub("", text)
    text = _FENCE_CLOSE.sub("", text)
    text = _escape_lone_backslashes(text)
    return json.loads(text)


# --- URI + slug helpers -----------------------------------------------------

_NON_SLUG = re.compile(r"[^a-z0-9]+")


def _slugify(name: str) -> str:
    """Stable lowercase-hyphen slug from a record name.

    Strips diacritics so "Pasch's" becomes "paschs" and Greek-prefixed names
    survive as ASCII-only. Used in document URIs — must be deterministic.
    """
    if not name:
        return "unnamed"
    n = unicodedata.normalize("NFKD", name)
    n = n.encode("ascii", "ignore").decode("ascii")
    n = n.lower()
    n = _NON_SLUG.sub("-", n).strip("-")
    return n or "unnamed"


def _bundle_basename(path: Path) -> str:
    """Stable bundle id from filename: strip extension."""
    return path.stem


def _record_uri(bundle: str, pillar: str, kind: str, idx: int, slug: str) -> str:
    """Stable URI for one axiom/theorem record."""
    return f"claim-pack://{bundle}/pillar/{pillar}/{kind}/{idx:03d}/{slug}"


# --- Cross-reference resolution --------------------------------------------

# Maps short-form bundle keys used inside `provenance.references` strings to
# canonical bundle basenames. Both directions covered. The "v2" / "v3" suffix
# inside an actual filename is preserved by the file's stem; the JSON
# references use the unversioned short key (theoremg4 / axiomsg4).
_REF_BUNDLE_ALIASES = {
    "theoremg4": ("theoremsg4", "theoremg4"),
    "axiomsg4": ("axiomsg4",),
}


def _resolve_reference(
    ref: str,
    bundles: dict[str, dict],
) -> tuple[str | None, str]:
    """Resolve `theoremg4.json:pillar.I.logic.excludedMiddle` to a URI.

    Returns ``(resolved_basename, uri)``. ``resolved_basename`` is the
    actual bundle stem matched against ``bundles``; ``None`` if no match
    (the URI is still a best-effort string pointer the edge can carry).

    Reference grammar (observed in the v2 packs):
      ``<bundle_short>.json:pillar.<pillar>.<area>.<slug>`` — by-slug match
      ``<bundle_short>.json:pillar.<pillar>``                — pillar-level
    """
    body = ref
    if ":" in body:
        head, body = body.split(":", 1)
        # head is "axiomsg4.json" or "theoremg4.json"
        short = head.removesuffix(".json")
    else:
        short = ""
    parts = [p for p in body.split(".") if p]
    if len(parts) < 2 or parts[0] != "pillar":
        return None, ref
    pillar = parts[1]
    leaf = parts[-1] if len(parts) >= 3 else None

    candidates = _REF_BUNDLE_ALIASES.get(short, (short,))
    matched_basename: str | None = None
    matched_bundle: dict | None = None
    for b_name, b_data in bundles.items():
        for cand in candidates:
            if b_name.startswith(cand):
                matched_basename = b_name
                matched_bundle = b_data
                break
        if matched_bundle is not None:
            break

    if matched_bundle is None or leaf is None:
        # Pillar-level pointer or unknown bundle. Best-effort URI.
        if matched_basename and pillar:
            uri = f"claim-pack://{matched_basename}/pillar/{pillar}"
            return matched_basename, uri
        return None, ref

    pillar_obj = (matched_bundle.get("pillars") or {}).get(pillar) or {}
    items = pillar_obj.get("axioms") or pillar_obj.get("theorems") or []
    leaf_slug = _slugify(leaf)
    item_kind = "axioms" if "axioms" in pillar_obj else "theorems"
    for idx, item in enumerate(items):
        if _slugify(item.get("name", "")) == leaf_slug:
            uri = _record_uri(
                matched_basename, pillar, item_kind, idx, leaf_slug
            )
            return matched_basename, uri
    # Bundle matched but slug didn't — fall back to pillar-level pointer.
    return matched_basename, f"claim-pack://{matched_basename}/pillar/{pillar}"


# --- Document construction --------------------------------------------------

def _format_record(record: dict) -> str:
    """Project one axiom/theorem record into canonical Document.content.

    Layout (see ticket #000029 §2.2):

        <name>

        <delta>

        <nablaConcise>

        <nablaVerbose>

        Role: <role>
        Source: <source_reference>
        Group: <foundational_group> · Category: <category> · Subfield: <subfield>

    Empty fields are omitted. Whitespace is preserved as-is — canonicalize()
    in arborist/document.py runs later in the ingest pipeline.
    """
    name = record.get("name", "").strip()
    delta = record.get("delta", "").strip()
    concise = record.get("nablaConcise", "").strip()
    verbose = record.get("nablaVerbose", "").strip()
    role = record.get("role", "").strip()
    source_ref = record.get("source_reference", "").strip()
    group = record.get("foundational_group", "").strip()
    category = record.get("category", "").strip()
    subfield = record.get("subfield", "").strip()

    parts: list[str] = []
    if name:
        parts.append(name)
    if delta:
        parts.append(delta)
    if concise:
        parts.append(concise)
    if verbose:
        parts.append(verbose)
    tail: list[str] = []
    if role:
        tail.append(f"Role: {role}")
    if source_ref:
        tail.append(f"Source: {source_ref}")
    classification = " · ".join(
        f"{label}: {value}"
        for label, value in (
            ("Group", group), ("Category", category), ("Subfield", subfield),
        )
        if value
    )
    if classification:
        tail.append(classification)
    if tail:
        parts.append("\n".join(tail))
    return "\n\n".join(parts)


def _record_extra(
    bundle_basename: str,
    bundle_meta: dict,
    pillar: str,
    kind: str,
    record: dict,
) -> dict:
    """Per-record metadata sidecar — never enters chunk hashes or cache_key."""
    return {
        "bundle": bundle_basename,
        "bundle_id": (bundle_meta.get("artifact_id") or "").strip(),
        "version": (bundle_meta.get("version") or "").strip(),
        "pillar": pillar,
        "kind": kind,
        "name": record.get("name", "").strip(),
        "runic": record.get("runicLabel", "").strip(),
        "category": record.get("category", "").strip(),
        "subfield": record.get("subfield", "").strip(),
        "source_ref": record.get("source_reference", "").strip(),
        "date_intro": record.get("date_of_introduction", "").strip(),
        "formal_lang": record.get("formal_language", "").strip(),
    }


def _record_title(record: dict) -> str:
    """Document title = the human-friendly record name."""
    return record.get("name", "").strip() or "(unnamed)"


# --- Source class -----------------------------------------------------------

class ClaimPackSource(Source):
    """One Document per axiom / theorem record across one or more bundles."""

    source_type = "claim_pack"

    def __init__(self, paths):
        if isinstance(paths, (str, Path)):
            paths = [paths]
        self.paths = [Path(p) for p in paths]
        if not self.paths:
            raise ValueError("ClaimPackSource requires at least one bundle path")
        for p in self.paths:
            if not p.is_file():
                raise FileNotFoundError(f"claim-pack bundle not found: {p}")

    def iter_documents(self) -> Iterator[Document]:
        bundles: dict[str, dict] = {}
        for p in self.paths:
            basename = _bundle_basename(p)
            bundles[basename] = _parse_lenient(p.read_text(encoding="utf-8"))

        for basename, bundle in bundles.items():
            yield from _iter_bundle(basename, bundle, bundles)


def _iter_bundle(
    basename: str,
    bundle: dict,
    all_bundles: dict[str, dict],
) -> Iterator[Document]:
    """Yield one Document per axiom/theorem inside a single bundle."""
    meta = bundle.get("metadata") or {}
    pillars = bundle.get("pillars") or {}
    for pillar_key, pillar_obj in pillars.items():
        # Collect pillar-level provenance refs once; attach to first record.
        refs = (pillar_obj.get("provenance") or {}).get("references") or []
        pillar_edges = [_ref_to_edge(ref, all_bundles) for ref in refs]

        for kind_label, items_key in (("axiom", "axioms"), ("theorem", "theorems")):
            items = pillar_obj.get(items_key) or []
            for idx, record in enumerate(items):
                slug = _slugify(record.get("name", ""))
                uri = _record_uri(basename, pillar_key, items_key, idx, slug)
                content = _format_record(record)
                if not content:
                    # An empty record (no name, no delta, no prose) is a
                    # malformed entry — skip rather than commit nothing.
                    continue
                edges: list[Edge] = []
                if idx == 0 and pillar_edges:
                    edges.extend(pillar_edges)
                doc = Document(
                    uri=uri,
                    content=content,
                    source_type="claim_pack",
                    title=_record_title(record),
                    edges=edges,
                    extra=_record_extra(basename, meta, pillar_key, kind_label, record),
                )
                yield doc


def _ref_to_edge(ref: str, bundles: dict[str, dict]) -> Edge:
    """One pillar-level reference string → one outbound Edge."""
    _, dst_uri = _resolve_reference(ref, bundles)
    return Edge(edge_type="pillar_reference", dst_uri=dst_uri)
