"""Wikidata (CC0) and CoDEx (MIT) adapters for the Wikidata tier (PLAN section 73).

Wikidata is acquired as bounded, randomly sampled per-property SPARQL extracts
(see ``scripts/download_wikidata_data.py``).  The raw snapshot is a set of
gzip TSV files:

``properties/<PID>.tsv.gz`` and ``expansion/batch-*.tsv.gz``
    ``subject<TAB>property<TAB>object`` with bare Q/P identifiers.
``labels/entities-*.tsv.gz``
    ``qid<TAB>english label<TAB>description<TAB>instance-of QIDs (|-joined)``.
``property_labels.tsv``
    ``pid<TAB>english label<TAB>description``.

Relation phrases come from the real English property labels (for example
``the "located in the administrative territorial entity" relation``), so the
tier adds hundreds of human-named relation types instead of the synthetic
ontology's 18.  Queries never contain the raw property identifiers.

Held-out design (all deterministic):

* a fixed list of properties is removed from every train graph and only emitted
  into a relation-held-out evaluation shard;
* entities are split by ``sha256(QID) % 100``; held-out-bucket entities are never
  the subject of a train edge, and evaluation queries are anchored only on them;
* CoDEx is an evaluation-only, source-held-out shard set.
"""

from __future__ import annotations

import csv
import gzip
import hashlib
import json
import re
from collections import Counter, defaultdict
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass, replace
from pathlib import Path
from typing import IO, Any

from graph.graph import EdgeRecord, GraphStore
from graph.real.corpus import RealGraph, SourceMetadata, generate_real_examples
from graph.synthetic.generator import ReasoningExample

WIKIDATA_ENTITY_PREFIX = "http://www.wikidata.org/entity/"
_QID = re.compile(r"^Q[1-9][0-9]*$")
_PID = re.compile(r"^P[1-9][0-9]*$")


# --------------------------------------------------------------------------- #
# Configuration and query text
# --------------------------------------------------------------------------- #


def load_wikidata_config(path: str | Path) -> dict[str, Any]:
    """Load ``configs/wikidata_sources.yaml`` and validate the allow-list."""
    import yaml

    payload = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
    if not isinstance(payload, dict) or "wikidata" not in payload:
        raise ValueError(f"{path}: missing 'wikidata' section")
    wikidata = payload["wikidata"]
    properties = wikidata.get("properties") or {}
    if not isinstance(properties, dict) or not properties:
        raise ValueError(f"{path}: wikidata.properties must be a non-empty mapping")
    for pid in properties:
        if not _PID.match(str(pid)):
            raise ValueError(f"{path}: invalid property id {pid!r}")
    heldout = [str(pid) for pid in wikidata.get("heldout_properties") or ()]
    missing = [pid for pid in heldout if pid not in properties]
    if missing:
        raise ValueError(f"{path}: held-out properties not allow-listed: {missing}")
    return payload


def property_query(pid: str, limit: int) -> str:
    """Bounded random-sample SPARQL for one property (entity objects only)."""
    if not _PID.match(pid):
        raise ValueError(f"invalid property id {pid!r}")
    if limit < 1:
        raise ValueError("limit must be positive")
    return (
        "SELECT ?s ?o WHERE {\n"
        "  SERVICE bd:sample {\n"
        f"    ?s wdt:{pid} ?o .\n"
        f"    bd:serviceParam bd:sample.limit {int(limit)} .\n"
        '    bd:serviceParam bd:sample.sampleType "RANDOM" .\n'
        "  }\n"
        f'  FILTER(STRSTARTS(STR(?s), "{WIKIDATA_ENTITY_PREFIX}Q") &&\n'
        f'         STRSTARTS(STR(?o), "{WIKIDATA_ENTITY_PREFIX}Q"))\n'
        "}"
    )


def expansion_query(subjects: Sequence[str], pids: Sequence[str], limit: int) -> str:
    """Outgoing allow-listed edges for a batch of (frequently referenced) items."""
    for qid in subjects:
        if not _QID.match(qid):
            raise ValueError(f"invalid entity id {qid!r}")
    for pid in pids:
        if not _PID.match(pid):
            raise ValueError(f"invalid property id {pid!r}")
    values_s = " ".join(f"wd:{qid}" for qid in subjects)
    values_p = " ".join(f"wdt:{pid}" for pid in pids)
    return (
        "SELECT ?s ?p ?o WHERE {\n"
        f"  VALUES ?s {{ {values_s} }}\n"
        f"  VALUES ?p {{ {values_p} }}\n"
        "  ?s ?p ?o .\n"
        f'  FILTER(STRSTARTS(STR(?o), "{WIKIDATA_ENTITY_PREFIX}Q"))\n'
        f"}} LIMIT {int(limit)}"
    )


def label_query(entities: Sequence[str]) -> str:
    """English label, description and instance-of classes for a batch of items."""
    for qid in entities:
        if not _QID.match(qid):
            raise ValueError(f"invalid entity id {qid!r}")
    values = " ".join(f"wd:{qid}" for qid in entities)
    return (
        "SELECT ?e ?l ?d ?t WHERE {\n"
        f"  VALUES ?e {{ {values} }}\n"
        '  OPTIONAL { ?e rdfs:label ?l . FILTER(LANG(?l) = "en") }\n'
        '  OPTIONAL { ?e schema:description ?d . FILTER(LANG(?d) = "en") }\n'
        "  OPTIONAL { ?e wdt:P31 ?t }\n"
        "}"
    )


def property_label_query(pids: Sequence[str]) -> str:
    for pid in pids:
        if not _PID.match(pid):
            raise ValueError(f"invalid property id {pid!r}")
    values = " ".join(f"wd:{pid}" for pid in pids)
    return (
        "SELECT ?p ?l ?d WHERE {\n"
        f"  VALUES ?p {{ {values} }}\n"
        '  OPTIONAL { ?p rdfs:label ?l . FILTER(LANG(?l) = "en") }\n'
        '  OPTIONAL { ?p schema:description ?d . FILTER(LANG(?d) = "en") }\n'
        "}"
    )


# --------------------------------------------------------------------------- #
# SPARQL result parsing (TSV or JSON)
# --------------------------------------------------------------------------- #


def _term(value: str) -> str:
    """Normalise one SPARQL TSV/JSON term to a bare ID or plain string."""
    value = value.strip()
    if value.startswith("<") and value.endswith(">"):
        value = value[1:-1]
    for prefix in (WIKIDATA_ENTITY_PREFIX, "http://www.wikidata.org/prop/direct/"):
        if value.startswith(prefix):
            return value[len(prefix) :]
    if value.startswith('"'):
        match = re.match(r'^"(.*)"(?:@[A-Za-z-]+|\^\^<[^>]*>)?$', value, re.DOTALL)
        if match:
            text = match.group(1)
            return (
                text.replace('\\"', '"')
                .replace("\\t", " ")
                .replace("\\n", " ")
                .replace("\\\\", "\\")
            )
    return value


def parse_sparql_tsv(text: str) -> list[dict[str, str]]:
    """Parse a WDQS ``text/tab-separated-values`` response."""
    lines = text.splitlines()
    if not lines:
        return []
    header = [name.lstrip("?") for name in lines[0].split("\t")]
    rows: list[dict[str, str]] = []
    for line in lines[1:]:
        if not line.strip():
            continue
        fields = line.split("\t")
        fields += [""] * (len(header) - len(fields))
        rows.append(
            {name: _term(field) for name, field in zip(header, fields, strict=False)}
        )
    return rows


def parse_sparql_json(payload: Mapping[str, Any] | str) -> list[dict[str, str]]:
    """Parse a SPARQL 1.1 JSON results document."""
    data = json.loads(payload) if isinstance(payload, str) else payload
    variables = list(data.get("head", {}).get("vars", ()))
    rows: list[dict[str, str]] = []
    for binding in data.get("results", {}).get("bindings", ()):
        row = {}
        for name in variables:
            cell = binding.get(name)
            if not cell:
                row[name] = ""
            elif cell.get("type") == "uri":
                row[name] = _term("<" + str(cell.get("value", "")) + ">")
            else:
                row[name] = " ".join(str(cell.get("value", "")).split())
        rows.append(row)
    return rows


# --------------------------------------------------------------------------- #
# Raw snapshot readers
# --------------------------------------------------------------------------- #


def _open_text(path: str | Path) -> IO[str]:
    file_path = Path(path)
    if file_path.name.endswith(".gz"):
        return gzip.open(file_path, "rt", encoding="utf-8", newline="")
    return file_path.open(encoding="utf-8", newline="")


def _rows(path: str | Path) -> Iterator[list[str]]:
    with _open_text(path) as handle:
        reader = csv.reader(handle, delimiter="\t", quoting=csv.QUOTE_NONE)
        header = next(reader, None)
        if header is None:
            return
        yield from reader


def iter_wikidata_triples(
    paths: Iterable[str | Path],
) -> Iterator[tuple[str, str, str]]:
    """Yield ``(subject, property, object)`` from raw extract files, validated."""
    for path in paths:
        for fields in _rows(path):
            if len(fields) < 3:
                continue
            subject, pid, obj = (field.strip() for field in fields[:3])
            if (
                _QID.match(subject)
                and _PID.match(pid)
                and _QID.match(obj)
                and subject != obj
            ):
                yield subject, pid, obj


@dataclass(frozen=True)
class PropertyInfo:
    pid: str
    label: str
    description: str = ""


def load_property_labels(path: str | Path) -> dict[str, PropertyInfo]:
    result: dict[str, PropertyInfo] = {}
    for fields in _rows(path):
        if not fields or not _PID.match(fields[0]):
            continue
        label = fields[1].strip() if len(fields) > 1 else ""
        description = fields[2].strip() if len(fields) > 2 else ""
        if label:
            result[fields[0]] = PropertyInfo(fields[0], label, description)
    return result


@dataclass(frozen=True)
class EntityTable:
    labels: Mapping[str, str]
    descriptions: Mapping[str, str]
    classes: Mapping[str, tuple[str, ...]]


def load_entity_labels(paths: Iterable[str | Path]) -> EntityTable:
    labels: dict[str, str] = {}
    descriptions: dict[str, str] = {}
    classes: dict[str, tuple[str, ...]] = {}
    for path in paths:
        for fields in _rows(path):
            if not fields or not _QID.match(fields[0]):
                continue
            qid = fields[0]
            label = fields[1].strip() if len(fields) > 1 else ""
            if label:
                labels[qid] = label
            if len(fields) > 2 and fields[2].strip():
                descriptions[qid] = fields[2].strip()
            if len(fields) > 3 and fields[3].strip():
                classes[qid] = tuple(
                    sorted(
                        {value for value in fields[3].split("|") if _QID.match(value)}
                    )
                )
    return EntityTable(labels, descriptions, classes)


# --------------------------------------------------------------------------- #
# Graph construction
# --------------------------------------------------------------------------- #


def relation_phrase_from_label(label: str) -> str:
    """Natural-language phrase built from a real property label."""
    clean = " ".join(label.replace('"', "'").split())
    return f'the "{clean}" relation'


def entity_bucket(qid: str) -> int:
    """Stable 0-99 bucket from the QID (not Python's salted hash)."""
    return int.from_bytes(hashlib.sha256(qid.encode("utf-8")).digest()[:8], "big") % 100


def is_heldout_entity(qid: str, percent: int) -> bool:
    return entity_bucket(qid) < percent


def _node_type(
    qid: str,
    classes: Mapping[str, tuple[str, ...]],
    labels: Mapping[str, str],
    class_frequency: Mapping[str, int],
) -> str:
    candidates = [cls for cls in classes.get(qid, ()) if cls in labels]
    if not candidates:
        return "entity"
    best = min(candidates, key=lambda cls: (-class_frequency.get(cls, 0), cls))
    return " ".join(labels[best].split())[:80] or "entity"


def build_labelled_graph(
    triples: Iterable[tuple[str, str, str]],
    source: SourceMetadata,
    *,
    relation_labels: Mapping[str, PropertyInfo],
    entities: EntityTable,
    node_prefix: str = "wd:",
    max_edges_per_subject_property: int | None = 24,
    evidence_kind: str = "wikidata_truthy_statement",
) -> RealGraph:
    """Build a ``RealGraph`` from ID triples, dropping unlabelled endpoints.

    Triples are de-duplicated and sorted, so the result is independent of the
    order of the input files.  Relation phrases come from ``relation_labels``.
    """
    unique = sorted(
        {
            triple
            for triple in triples
            if triple[1] in relation_labels
            and triple[0] in entities.labels
            and triple[2] in entities.labels
        }
    )
    per_key: Counter[tuple[str, str]] = Counter()
    kept: list[tuple[str, str, str]] = []
    for subject, pid, obj in unique:
        key = (subject, pid)
        if (
            max_edges_per_subject_property is not None
            and per_key[key] >= max_edges_per_subject_property
        ):
            continue
        per_key[key] += 1
        kept.append((subject, pid, obj))

    endpoints = sorted({qid for s, _, o in kept for qid in (s, o)})
    class_frequency: Counter[str] = Counter(
        cls for qid in endpoints for cls in entities.classes.get(qid, ())
    )
    nodes: dict[str, str] = {}
    node_attributes: dict[str, dict[str, str]] = {}
    labels: dict[str, str] = {}
    for qid in endpoints:
        node_id = node_prefix + qid
        node_type = _node_type(qid, entities.classes, entities.labels, class_frequency)
        nodes[node_id] = node_type
        labels[node_id] = entities.labels[qid]
        attributes = {"label": entities.labels[qid], "wikidata_id": qid}
        description = entities.descriptions.get(qid, "")
        if description:
            attributes["description"] = description[:160]
        node_attributes[node_id] = attributes

    relations = sorted({pid for _, pid, _ in kept})
    phrases = {
        pid: relation_phrase_from_label(relation_labels[pid].label) for pid in relations
    }
    edges = [
        EdgeRecord(
            node_prefix + subject,
            pid,
            node_prefix + obj,
            nodes[node_prefix + subject],
            nodes[node_prefix + obj],
            relation_labels[pid].label,
            {"evidence_kind": evidence_kind},
        )
        for subject, pid, obj in kept
    ]
    relation_attributes = {
        pid: {
            "functional": "false",
            "label": relation_labels[pid].label,
            "description": relation_labels[pid].description[:240] or phrases[pid],
            "wikidata_property": pid,
        }
        for pid in relations
    }
    if not nodes:
        raise ValueError(f"{source.source_id}: no labelled triples to build a graph")
    return RealGraph(
        graph=_fast_graph_store(
            nodes,
            edges,
            graph_version=f"{source.source_id}@{source.version}",
            node_attributes=node_attributes,
            relation_attributes=relation_attributes,
        ),
        node_labels=labels,
        relation_phrases=phrases,
        source=source,
    )


def _fast_graph_store(
    nodes: Mapping[str, str],
    edges: Sequence[EdgeRecord],
    *,
    graph_version: str,
    node_attributes: Mapping[str, Mapping[str, str]],
    relation_attributes: Mapping[str, Mapping[str, str]],
) -> GraphStore:
    """Construct a ``GraphStore`` with an O(E) adjacency build.

    ``GraphStore.__init__`` appends to per-node tuples, which is quadratic in
    out-degree; at Wikidata scale that dominates build time.  The resulting
    object is identical (same edges, same per-node edge order).
    """
    store = GraphStore.__new__(GraphStore)
    store.nodes = dict(nodes)
    if not store.nodes:
        raise ValueError("graph must contain at least one node")
    store.edges = tuple(edges)
    store.graph_version = graph_version
    store.node_attributes = {key: dict(value) for key, value in node_attributes.items()}
    store.relation_attributes = {
        key: dict(value) for key, value in relation_attributes.items()
    }
    outgoing: dict[str, list[EdgeRecord]] = defaultdict(list)
    for edge in store.edges:
        if edge.source_id not in store.nodes or edge.target_id not in store.nodes:
            raise ValueError(f"edge references unknown node: {edge}")
        outgoing[edge.source_id].append(edge)
    store._out = {key: tuple(value) for key, value in outgoing.items()}
    return store


# --------------------------------------------------------------------------- #
# Held-out splits
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class WikidataSplit:
    """ID-level triples for the train graph and the held-out evaluation graph."""

    train: tuple[tuple[str, str, str], ...]
    evaluation: tuple[tuple[str, str, str], ...]
    heldout_properties: frozenset[str]
    heldout_entity_percent: int
    #: held-out-bucket subjects with at least one held-out-property edge
    heldout_relation_starts: tuple[str, ...]
    #: held-out-bucket subjects with at least one train-property edge
    heldout_entity_starts: tuple[str, ...]


def split_wikidata_triples(
    triples: Iterable[tuple[str, str, str]],
    *,
    heldout_properties: Iterable[str],
    heldout_entity_percent: int,
    continuation_layers: int = 2,
) -> WikidataSplit:
    """Split triples into a train set and a held-out evaluation set.

    Train: allow-listed, non-held-out properties whose subject is outside the
    held-out entity bucket.  Evaluation: every edge whose subject is in the
    held-out bucket (both held-out and seen properties, so seen relations act as
    distractors), plus ``continuation_layers`` of seen-property edges reachable
    from those objects so multi-hop chains can continue.  Held-out properties
    therefore never occur in ``train`` and held-out-bucket entities are never a
    train subject.
    """
    if not 0 <= heldout_entity_percent < 100:
        raise ValueError("heldout_entity_percent must be in [0, 100)")
    heldout = frozenset(heldout_properties)
    unique = sorted(set(triples))
    by_subject: dict[str, list[tuple[str, str, str]]] = defaultdict(list)
    for triple in unique:
        by_subject[triple[0]].append(triple)

    train: list[tuple[str, str, str]] = []
    evaluation: set[tuple[str, str, str]] = set()
    relation_starts: set[str] = set()
    entity_starts: set[str] = set()
    frontier: set[str] = set()
    for subject, pid, obj in unique:
        if is_heldout_entity(subject, heldout_entity_percent):
            evaluation.add((subject, pid, obj))
            frontier.add(obj)
            if pid in heldout:
                relation_starts.add(subject)
            else:
                entity_starts.add(subject)
        elif pid not in heldout:
            train.append((subject, pid, obj))

    seen = set(frontier)
    for _ in range(continuation_layers):
        next_frontier: set[str] = set()
        for node in sorted(frontier):
            for triple in by_subject.get(node, ()):
                if triple[1] in heldout and not is_heldout_entity(
                    node, heldout_entity_percent
                ):
                    # held-out relations only ever appear anchored in the bucket
                    continue
                evaluation.add(triple)
                if triple[2] not in seen:
                    seen.add(triple[2])
                    next_frontier.add(triple[2])
        frontier = next_frontier
    return WikidataSplit(
        train=tuple(train),
        evaluation=tuple(sorted(evaluation)),
        heldout_properties=heldout,
        heldout_entity_percent=heldout_entity_percent,
        heldout_relation_starts=tuple(sorted(relation_starts)),
        heldout_entity_starts=tuple(sorted(entity_starts)),
    )


def generate_filtered_examples(
    real_graph: RealGraph,
    *,
    count: int,
    seed: int,
    accept: Callable[[ReasoningExample], bool],
    start_nodes: Sequence[str] | None = None,
    id_prefix: str | None = None,
    max_attempt_factor: int = 20,
    **kwargs: Any,
) -> Iterator[ReasoningExample]:
    """Stream ``count`` examples from ``generate_real_examples`` that pass ``accept``.

    Accepted examples are renumbered so example IDs stay contiguous.  Raises
    ``RuntimeError`` if the acceptance rate is too low to reach ``count``.
    """
    if count < 0:
        raise ValueError("count must be non-negative")
    prefix = id_prefix or real_graph.source.source_id
    produced = 0
    stream = generate_real_examples(
        real_graph,
        count=max(1, count * max_attempt_factor),
        seed=seed,
        start_nodes=start_nodes,
        **kwargs,
    )
    try:
        for example in stream:
            if produced >= count:
                break
            if not accept(example):
                continue
            yield replace(
                example,
                example_id=f"{prefix}-{seed}-{produced}",
                seed=seed + produced,
            )
            produced += 1
            if produced >= count:
                break
    except RuntimeError:
        if produced < count:
            raise
    finally:
        close = getattr(stream, "close", None)
        if close is not None:
            close()
    if produced < count:
        raise RuntimeError(
            f"accepted only {produced}/{count} examples from {prefix}; "
            "the acceptance filter is too strict for this graph"
        )


def contains_relation(relations: Iterable[str]) -> Callable[[ReasoningExample], bool]:
    wanted = frozenset(relations)
    return lambda example: any(pid in wanted for pid in example.composition)


def excludes_relations(relations: Iterable[str]) -> Callable[[ReasoningExample], bool]:
    banned = frozenset(relations)
    return lambda example: not any(pid in banned for pid in example.composition)


# --------------------------------------------------------------------------- #
# CoDEx
# --------------------------------------------------------------------------- #


def _read_json(path: str | Path) -> Any:
    with _open_text(path) as handle:
        return json.load(handle)


def load_codex(
    triple_paths: Sequence[str | Path],
    source: SourceMetadata,
    *,
    entities_json: str | Path,
    relations_json: str | Path,
    entity_types_json: str | Path | None = None,
    types_json: str | Path | None = None,
    max_triples: int | None = None,
    max_edges_per_subject_property: int | None = 24,
) -> RealGraph:
    """Load CoDEx triples with English entity/relation labels and entity types."""
    entities_payload = _read_json(entities_json)
    relations_payload = _read_json(relations_json)
    labels: dict[str, str] = {}
    descriptions: dict[str, str] = {}
    for qid, record in entities_payload.items():
        if isinstance(record, dict) and str(record.get("label", "")).strip():
            labels[str(qid)] = " ".join(str(record["label"]).split())
            if str(record.get("description", "")).strip():
                descriptions[str(qid)] = " ".join(str(record["description"]).split())
    relation_labels = {
        str(pid): PropertyInfo(
            str(pid),
            " ".join(str(record.get("label", "")).split()),
            " ".join(str(record.get("description", "")).split()),
        )
        for pid, record in relations_payload.items()
        if isinstance(record, dict) and str(record.get("label", "")).strip()
    }
    classes: dict[str, tuple[str, ...]] = {}
    if entity_types_json is not None:
        for qid, values in _read_json(entity_types_json).items():
            if isinstance(values, list):
                classes[str(qid)] = tuple(sorted(str(value) for value in values))
    if types_json is not None:
        for qid, record in _read_json(types_json).items():
            if isinstance(record, dict) and str(record.get("label", "")).strip():
                labels.setdefault(str(qid), " ".join(str(record["label"]).split()))

    def triples() -> Iterator[tuple[str, str, str]]:
        emitted = 0
        for path in triple_paths:
            with _open_text(path) as handle:
                for line in handle:
                    fields = line.rstrip("\n").split("\t")
                    if len(fields) < 3:
                        continue
                    head, relation, tail = (field.strip() for field in fields[:3])
                    if not head or not relation or not tail or head == tail:
                        continue
                    yield head, relation, tail
                    emitted += 1
                    if max_triples is not None and emitted >= max_triples:
                        return

    return build_labelled_graph(
        triples(),
        source,
        relation_labels=relation_labels,
        entities=EntityTable(labels, descriptions, classes),
        node_prefix="codex:",
        max_edges_per_subject_property=max_edges_per_subject_property,
        evidence_kind="codex_benchmark_triple",
    )


def load_wikidata_snapshot(
    root: str | Path,
) -> tuple[list[tuple[str, str, str]], dict[str, PropertyInfo], EntityTable]:
    """Read every raw extract under ``root`` (a ``data/wikidata/wikidata`` dir)."""
    base = Path(root)
    triple_paths = sorted(base.glob("properties/*.tsv.gz")) + sorted(
        base.glob("expansion/*.tsv.gz")
    )
    triples = list(iter_wikidata_triples(triple_paths))
    relation_labels = load_property_labels(base / "property_labels.tsv")
    entities = load_entity_labels(sorted(base.glob("labels/*.tsv.gz")))
    return triples, relation_labels, entities