"""Persistence + cache lookup for canonical-projection answers.

#000027 — implementation.

Math/logic π* answers (`arithmetic@v1`, `logic-kernel@v1`,
`time-series-quantized@v1`, …) used to short-circuit RAG and
return without ever touching the audit chain. This module makes
them first-class providence_cache rows: cache_key minted from the
8-dim invariant with synthetic-but-deterministic values for the
three RAG-shaped dimensions (`source_root`, `model_profile_hash`,
`conversation_hash`); persisted with `audit_mode =
CANONICAL_PROJECTION` and `verifier_method = canonical_projection`;
audit-chain entry under event_type ``providence_canonical``.

The hard-constraint contract from the ticket: the canonical bytes
ARE the answer. Round-trip MUST be byte-identical — re-canonicalizing
the stored input through the same `pi_star_ref` produces the
stored output, every time. We trust the row on cache hit (do not
re-run the kernel); kernel-version drift is handled by `pi_star_ref`
bumping (`@v1` → `@v2`), which mints a new synthetic source_root
and naturally orphans prior records.

Synthetic-dim choices (per ticket §2.2):

- ``source_root``         = sha256("pi_star_source:" + pi_star_ref)
- ``model_profile_hash``  = sha256("pi_star_model:"  + pi_star_ref)
- ``conversation_hash``   = sha256("pi_star_conv:"   + canonical_question + ":" + pi_star_ref)
- ``chunking_version``    = literal "n/a-canonical" — chunking is
  irrelevant to math; using the live CHUNKING_VERSION would
  mass-stale every canonical row whenever the wikipedia chunker
  changes (unrelated to math correctness).

Other dims (`question_hash`, `governance_policy_hash`,
`schema_version`, `canonicalization_version`) are real and shared
with the RAG path.
"""

from __future__ import annotations

import hashlib
import json
import sqlite3
import time
from typing import Callable

from arborist import (
    CANONICALIZATION_VERSION,
    CHUNKING_VERSION,  # noqa: F401 — referenced by docstring; not used
    SCHEMA_VERSION,
)
from arborist.qa.dag import _canonical_json, _sha256_hex
from arborist.qa.keys import (
    cache_key as _cache_key,
    canonical_question as _canonical_question,
    governance_policy_hash,
    question_hash as _question_hash,
)
from arborist.store import append_audit, transaction


CANONICAL_CHUNKING_VERSION = "n/a-canonical"
CANONICAL_AUDIT_EVENT_TYPE = "providence_canonical"


def canonical_synthetic_source_root(pi_star_ref: str) -> str:
    return _sha256_hex(f"pi_star_source:{pi_star_ref}")


def canonical_synthetic_model_hash(pi_star_ref: str) -> str:
    return _sha256_hex(f"pi_star_model:{pi_star_ref}")


def canonical_synthetic_conversation_hash(
    canonical_question_str: str, pi_star_ref: str
) -> str:
    return _sha256_hex(
        f"pi_star_conv:{canonical_question_str}:{pi_star_ref}"
    )


def canonical_cache_key(
    *,
    question: str,
    pi_star_ref: str,
    policy: dict,
    mode: str = "equivalence_class",
) -> str:
    """8-dim cache_key for a canonical projection answer.

    Reuses the standard `cache_key` function from `arborist.qa.keys`
    so the canonical path lives in the same identity space as RAG
    rows. The three RAG-shaped dimensions get synthetic values
    pinned to `pi_star_ref`.
    """
    qhash = _question_hash(question, mode=mode)
    src = canonical_synthetic_source_root(pi_star_ref)
    mhash = canonical_synthetic_model_hash(pi_star_ref)
    chash = canonical_synthetic_conversation_hash(
        _canonical_question(question, mode=mode), pi_star_ref
    )
    ghash = governance_policy_hash(policy)
    return _cache_key(
        src,
        qhash,
        mhash,
        chash,
        ghash,
        SCHEMA_VERSION,
        CANONICALIZATION_VERSION,
        CANONICAL_CHUNKING_VERSION,
    )


def build_canonical_run_dag(
    *,
    question_hash_value: str,
    pi_star_ref: str,
    canonical_input_sha256: str,
    canonical_output_sha256: str,
) -> dict:
    """3-stage run-DAG for canonical projections.

    question → canonical_projection → final_label

    Each node carries an opaque ``hash`` (stage payload sha256). The
    Merkle root of the three hashes is the run_dag_root we persist.
    """
    canonical_payload = {
        "stage": "canonical_projection",
        "pi_star_ref": pi_star_ref,
        "canonical_input_sha256": canonical_input_sha256,
        "canonical_output_sha256": canonical_output_sha256,
    }
    canonical_hash = _sha256_hex(_canonical_json(canonical_payload))
    final_payload = {
        "stage": "final_label",
        "audit_mode": "CANONICAL_PROJECTION",
        "verifier_method": "canonical_projection",
    }
    final_hash = _sha256_hex(_canonical_json(final_payload))
    nodes = [
        {"stage": "question", "hash": question_hash_value},
        {"stage": "canonical_projection", "hash": canonical_hash,
         "payload": canonical_payload},
        {"stage": "final_label", "hash": final_hash,
         "payload": final_payload},
    ]
    # Root: sha256 of concatenated hashes (cheap, deterministic, and
    # consistent with the lighter run-DAG shape used by other paths
    # that don't need full Merkle proofs).
    root = hashlib.sha256(
        "".join(n["hash"] for n in nodes).encode("utf-8")
    ).hexdigest()
    return {"root": root, "nodes": nodes}


def lookup_canonical(
    qa_conn: sqlite3.Connection, cache_key_value: str
) -> sqlite3.Row | None:
    """Return the live providence_cache row for `cache_key_value`,
    or None on miss. Filters on falsification_state='live' the same
    way the RAG path does.
    """
    row = qa_conn.execute(
        "SELECT cache_key, source_root, document_uri, question_text, "
        " answer_text, audit_mode, verifier_method, n_quotes, n_verified, "
        " falsification_state, run_dag_root, run_dag_blob, "
        " audit_event_hash, created_at "
        "FROM providence_cache "
        "WHERE cache_key = ? AND falsification_state = 'live'",
        (cache_key_value,),
    ).fetchone()
    return row


def increment_hit_count(
    qa_conn: sqlite3.Connection, cache_key_value: str
) -> None:
    """Mirror the RAG path's last_hit_at + hit_count bump on cache hit."""
    qa_conn.execute(
        "UPDATE providence_cache "
        "SET last_hit_at = ?, hit_count = hit_count + 1 "
        "WHERE cache_key = ?",
        (int(time.time()), cache_key_value),
    )


def persist_canonical(
    qa_conn: sqlite3.Connection,
    *,
    cache_key_value: str,
    question: str,
    pi_star_ref: str,
    canonical_input_bytes: bytes,
    canonical_output_bytes: bytes,
    policy: dict,
    chain: str = "private",
    mode: str = "equivalence_class",
) -> tuple[str, str, dict]:
    """Write one canonical providence_cache row + matching audit event.

    Returns (event_hash, run_dag_root, run_dag).
    """
    canonical_input_sha = _sha256_hex(
        canonical_input_bytes.decode("utf-8", errors="surrogatepass")
    )
    canonical_output_sha = _sha256_hex(
        canonical_output_bytes.decode("utf-8", errors="surrogatepass")
    )
    qhash = _question_hash(question, mode=mode)
    src = canonical_synthetic_source_root(pi_star_ref)
    mhash = canonical_synthetic_model_hash(pi_star_ref)
    chash = canonical_synthetic_conversation_hash(
        _canonical_question(question, mode=mode), pi_star_ref
    )
    ghash = governance_policy_hash(policy)

    run_dag = build_canonical_run_dag(
        question_hash_value=qhash,
        pi_star_ref=pi_star_ref,
        canonical_input_sha256=canonical_input_sha,
        canonical_output_sha256=canonical_output_sha,
    )
    run_dag_blob = json.dumps(run_dag, separators=(",", ":"))

    answer_text = canonical_output_bytes.decode("utf-8", errors="replace")
    document_uri = f"pi_star://{pi_star_ref}"
    merkle_proof = json.dumps(
        {"context_root": None, "sources": []}, separators=(",", ":")
    )
    now = int(time.time())

    audit_body = {
        "pi_star_ref": pi_star_ref,
        "question_text": question,
        "canonical_input_sha256": canonical_input_sha,
        "canonical_output_bytes_hex": canonical_output_bytes.hex(),
        "canonical_output_text": answer_text,
        "kernel_audit_mode": "CANONICAL_PROJECTION",
    }

    with transaction(qa_conn):
        event_hash = append_audit(
            qa_conn,
            event_type=CANONICAL_AUDIT_EVENT_TYPE,
            subject_root=cache_key_value,
            body=audit_body,
            ts=now,
        )
        qa_conn.execute(
            # ON CONFLICT(cache_key) DO NOTHING: canonical-projection answers
            # are deterministic, so a concurrent peer persisting the same
            # cache_key writes an identical row — the loser no-ops instead of
            # raising UNIQUE constraint failed.
            "INSERT INTO providence_cache "
            "(cache_key, source_root, document_uri, question_hash, question_text, "
            " answer_text, merkle_proof, model_profile_hash, conversation_hash, "
            " governance_policy_hash, schema_version, canonicalization_version, "
            " chunking_version, falsification_state, chain, audit_event_hash, "
            " created_at, hit_count, audit_mode, n_quotes, n_verified, "
            " unverified_quotes, verifier_method, run_dag_root, run_dag_blob) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'live', ?, ?, ?, 0, "
            " ?, ?, ?, ?, ?, ?, ?) "
            "ON CONFLICT(cache_key) DO NOTHING",
            (
                cache_key_value,
                src,
                document_uri,
                qhash,
                question,
                answer_text,
                merkle_proof,
                mhash,
                chash,
                ghash,
                SCHEMA_VERSION,
                CANONICALIZATION_VERSION,
                CANONICAL_CHUNKING_VERSION,
                chain,
                event_hash,
                now,
                "CANONICAL_PROJECTION",
                1,  # n_quotes — the kernel output IS the one "quote"
                1,  # n_verified — kernel-verified by definition
                None,  # unverified_quotes
                "canonical_projection",
                run_dag["root"],
                run_dag_blob,
            ),
        )
    return event_hash, run_dag["root"], run_dag


def lookup_or_persist(
    qa_conn: sqlite3.Connection,
    *,
    question: str,
    pi_star_ref: str,
    canonical_output_bytes: bytes,
    policy: dict,
    chain: str = "private",
    mode: str = "equivalence_class",
) -> tuple[str, sqlite3.Row | None, bool, dict]:
    """One-shot helper: mint cache_key, look up, and persist on miss.

    Returns (cache_key, row_or_None, was_hit, run_dag).
    On hit, ``row`` is the cached row and ``was_hit=True``;
    ``run_dag`` is reconstructed from the row's run_dag_blob (or {}
    if absent on legacy rows).
    On miss, the row is written; ``row`` is None (caller has the
    canonical bytes already), ``was_hit=False``; ``run_dag`` is
    the freshly-built DAG.
    """
    ckey = canonical_cache_key(
        question=question,
        pi_star_ref=pi_star_ref,
        policy=policy,
        mode=mode,
    )
    row = lookup_canonical(qa_conn, ckey)
    if row is not None:
        increment_hit_count(qa_conn, ckey)
        try:
            run_dag = json.loads(row["run_dag_blob"]) if row["run_dag_blob"] else {}
        except (json.JSONDecodeError, KeyError):
            run_dag = {}
        return ckey, row, True, run_dag
    _, _, run_dag = persist_canonical(
        qa_conn,
        cache_key_value=ckey,
        question=question,
        pi_star_ref=pi_star_ref,
        canonical_input_bytes=question.encode("utf-8"),
        canonical_output_bytes=canonical_output_bytes,
        policy=policy,
        chain=chain,
        mode=mode,
    )
    return ckey, None, False, run_dag


# Public exports.
__all__ = [
    "CANONICAL_AUDIT_EVENT_TYPE",
    "CANONICAL_CHUNKING_VERSION",
    "build_canonical_run_dag",
    "canonical_cache_key",
    "canonical_synthetic_conversation_hash",
    "canonical_synthetic_model_hash",
    "canonical_synthetic_source_root",
    "increment_hit_count",
    "lookup_canonical",
    "lookup_or_persist",
    "persist_canonical",
]
