#!/usr/bin/env python3
"""Build a private, row-level replay receipt for one applied KB proposal.

The guarded apply transaction keeps the immutable reviewed payload in
``kb_stage.kb_proposals``. This module binds that payload to the exact canonical
rows observed after commit, including generated UUIDs and timestamps. The
result is sufficient input for a future genesis-plus-ledger replayer without
putting claim bodies or source excerpts in logs or stdout.
"""

from __future__ import annotations

import hashlib
import json
import os
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

RECEIPT_CONTRACT_VERSION = 1
DEFAULT_RECEIPT_DIR = "/opt/teleo-eval/logs/kb-apply-receipts"


def _canonical_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)


def _sha256_json(value: Any) -> str:
    return hashlib.sha256(_canonical_json(value).encode("utf-8")).hexdigest()


def sha256_text(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def _sql_literal(value: Any) -> str:
    if value is None:
        return "null"
    escaped = str(value).replace("\\", "\\\\").replace("'", "''")
    return "E'" + escaped + "'"


def _jsonb(value: Any) -> str:
    return _sql_literal(_canonical_json(value)) + "::jsonb"


def _require_applied_strict_proposal(proposal: dict[str, Any]) -> dict[str, Any]:
    if proposal.get("status") != "applied":
        raise ValueError("replay receipt requires proposal status='applied'")
    if not proposal.get("applied_at"):
        raise ValueError("replay receipt requires a non-null applied_at")
    if not proposal.get("applied_by_handle"):
        raise ValueError("replay receipt requires applied_by_handle")
    payload = proposal.get("payload")
    apply_payload = payload.get("apply_payload") if isinstance(payload, dict) else None
    if not isinstance(apply_payload, dict):
        raise ValueError("replay receipt requires a strict payload.apply_payload object")
    return apply_payload


def _coalesced_json_agg(table: str, alias: str, where: str, order_by: str) -> str:
    return (
        "coalesce((select jsonb_agg(to_jsonb(" + alias + ") order by " + order_by + ") "
        "from " + table + " " + alias + " where " + where + "), '[]'::jsonb)"
    )


def _id_filter(alias: str, values: list[Any]) -> str:
    if not values:
        return "false"
    rendered = ", ".join(f"{_sql_literal(value)}::uuid" for value in values)
    return f"{alias}.id in ({rendered})"


def _or_filter(clauses: list[str]) -> str:
    if not clauses:
        return "false"
    return "(" + ") or (".join(clauses) + ")"


def _build_add_edge_postflight_sql(apply_payload: dict[str, Any]) -> str:
    where = (
        f"e.from_claim = {_sql_literal(apply_payload.get('from_claim'))}::uuid "
        f"and e.to_claim = {_sql_literal(apply_payload.get('to_claim'))}::uuid "
        f"and e.edge_type = {_sql_literal(apply_payload.get('edge_type'))}::edge_type"
    )
    rows = _coalesced_json_agg("public.claim_edges", "e", where, "e.id::text")
    return f"select jsonb_build_object('claim_edges', {rows})::text;"


def _build_attach_evidence_postflight_sql(apply_payload: dict[str, Any]) -> str:
    clauses = []
    for row in apply_payload.get("evidence") or []:
        role = row.get("role") or "grounds"
        clauses.append(
            f"ce.claim_id = {_sql_literal(row.get('claim_id'))}::uuid "
            f"and ce.source_id = {_sql_literal(row.get('source_id'))}::uuid "
            f"and ce.role = {_sql_literal(role)}::evidence_role"
        )
    rows = _coalesced_json_agg(
        "public.claim_evidence",
        "ce",
        _or_filter(clauses),
        "ce.claim_id::text, ce.source_id::text, ce.role::text",
    )
    return f"select jsonb_build_object('claim_evidence', {rows})::text;"


def _build_approve_claim_postflight_sql(apply_payload: dict[str, Any]) -> str:
    claims = apply_payload.get("claims") or []
    sources = apply_payload.get("sources") or []
    evidence = apply_payload.get("evidence") or []
    edges = apply_payload.get("edges") or []
    tools = apply_payload.get("reasoning_tools") or []

    evidence_clauses = []
    for row in evidence:
        role = row.get("role") or "grounds"
        evidence_clauses.append(
            f"ce.claim_id = {_sql_literal(row.get('claim_id'))}::uuid "
            f"and ce.source_id = {_sql_literal(row.get('source_id'))}::uuid "
            f"and ce.role = {_sql_literal(role)}::evidence_role"
        )
    edge_clauses = [
        f"e.from_claim = {_sql_literal(row.get('from_claim'))}::uuid "
        f"and e.to_claim = {_sql_literal(row.get('to_claim'))}::uuid "
        f"and e.edge_type = {_sql_literal(row.get('edge_type'))}::edge_type"
        for row in edges
    ]

    fields = {
        "claims": _coalesced_json_agg(
            "public.claims", "c", _id_filter("c", [row.get("id") for row in claims]), "c.id::text"
        ),
        "sources": _coalesced_json_agg(
            "public.sources", "s", _id_filter("s", [row.get("id") for row in sources]), "s.id::text"
        ),
        "claim_evidence": _coalesced_json_agg(
            "public.claim_evidence",
            "ce",
            _or_filter(evidence_clauses),
            "ce.claim_id::text, ce.source_id::text, ce.role::text",
        ),
        "claim_edges": _coalesced_json_agg("public.claim_edges", "e", _or_filter(edge_clauses), "e.id::text"),
        "reasoning_tools": _coalesced_json_agg(
            "public.reasoning_tools",
            "rt",
            _id_filter("rt", [row.get("id") for row in tools]),
            "rt.id::text",
        ),
    }
    arguments = ", ".join(f"{_sql_literal(name)}, {expression}" for name, expression in fields.items())
    return f"select jsonb_build_object({arguments})::text;"


def _build_revise_strategy_postflight_sql(proposal: dict[str, Any], apply_payload: dict[str, Any]) -> str:
    strategy = apply_payload.get("strategy") or {}
    agent_id = apply_payload.get("agent_id")
    node_clauses = []
    for row in apply_payload.get("strategy_nodes") or []:
        node_clauses.append(
            f"n.node_type = {_sql_literal(row.get('node_type'))} "
            f"and n.title = {_sql_literal(row.get('title'))} "
            f"and n.body = {_sql_literal(row.get('body'))} "
            f"and n.rank = {int(row.get('rank', 1))}"
        )
    node_filter = _or_filter(node_clauses)
    return f"""with matched_strategy as (
  select s.*
    from public.strategies s
   where s.agent_id = {_sql_literal(agent_id)}::uuid
     and s.diagnosis = {_sql_literal(strategy.get("diagnosis"))}
     and s.guiding_policy = {_sql_literal(strategy.get("guiding_policy"))}
     and s.proximate_objectives = {_jsonb(strategy.get("proximate_objectives"))}
     and s.created_at <= {_sql_literal(proposal.get("applied_at"))}::timestamptz
   order by s.version desc, s.id::text desc
   limit 1
), matched_nodes as (
  select n.*
    from public.strategy_nodes n
    join matched_strategy s
      on n.agent_id = s.agent_id and n.created_at = s.created_at
   where {node_filter}
)
select jsonb_build_object(
  'strategies', coalesce((select jsonb_agg(to_jsonb(s) order by s.id::text) from matched_strategy s), '[]'::jsonb),
  'strategy_nodes', coalesce((select jsonb_agg(to_jsonb(n) order by n.rank, n.id::text) from matched_nodes n), '[]'::jsonb)
)::text;"""


def build_postflight_sql(proposal: dict[str, Any]) -> str:
    apply_payload = _require_applied_strict_proposal(proposal)
    proposal_type = proposal.get("proposal_type")
    if proposal_type == "add_edge":
        return _build_add_edge_postflight_sql(apply_payload)
    if proposal_type == "attach_evidence":
        return _build_attach_evidence_postflight_sql(apply_payload)
    if proposal_type == "approve_claim":
        return _build_approve_claim_postflight_sql(apply_payload)
    if proposal_type == "revise_strategy":
        return _build_revise_strategy_postflight_sql(proposal, apply_payload)
    raise ValueError(f"replay receipt does not support proposal_type {proposal_type!r}")


def expected_row_counts(proposal: dict[str, Any]) -> dict[str, int]:
    apply_payload = _require_applied_strict_proposal(proposal)
    proposal_type = proposal.get("proposal_type")
    if proposal_type == "add_edge":
        return {"claim_edges": 1}
    if proposal_type == "attach_evidence":
        return {"claim_evidence": len(apply_payload.get("evidence") or [])}
    if proposal_type == "approve_claim":
        return {
            "claims": len(apply_payload.get("claims") or []),
            "sources": len(apply_payload.get("sources") or []),
            "claim_evidence": len(apply_payload.get("evidence") or []),
            "claim_edges": len(apply_payload.get("edges") or []),
            "reasoning_tools": len(apply_payload.get("reasoning_tools") or []),
        }
    if proposal_type == "revise_strategy":
        return {
            "strategies": 1,
            "strategy_nodes": len(apply_payload.get("strategy_nodes") or []),
        }
    raise ValueError(f"replay receipt does not support proposal_type {proposal_type!r}")


def _normalize_rows(canonical_rows: dict[str, Any], expected_counts: dict[str, int]) -> dict[str, list[dict[str, Any]]]:
    normalized: dict[str, list[dict[str, Any]]] = {}
    for table in expected_counts:
        rows = canonical_rows.get(table)
        if not isinstance(rows, list) or any(not isinstance(row, dict) for row in rows):
            raise ValueError(f"postflight {table} must be a list of row objects")
        normalized[table] = sorted(rows, key=_canonical_json)
    return normalized


def _semantic_projection(row: dict[str, Any], fields: tuple[str, ...]) -> dict[str, Any]:
    return {field: row.get(field) for field in fields}


def _expected_semantic_rows(proposal: dict[str, Any]) -> dict[str, list[dict[str, Any]]]:
    apply_payload = _require_applied_strict_proposal(proposal)
    proposal_type = proposal.get("proposal_type")
    if proposal_type == "add_edge":
        return {
            "claim_edges": [
                {
                    "from_claim": apply_payload.get("from_claim"),
                    "to_claim": apply_payload.get("to_claim"),
                    "edge_type": apply_payload.get("edge_type"),
                    "weight": apply_payload.get("weight"),
                    "created_by": None,
                }
            ]
        }
    if proposal_type == "attach_evidence":
        return {
            "claim_evidence": [
                {
                    "claim_id": row.get("claim_id"),
                    "source_id": row.get("source_id"),
                    "role": row.get("role") or "grounds",
                    "weight": row.get("weight"),
                    "created_by": None,
                }
                for row in apply_payload.get("evidence") or []
            ]
        }
    if proposal_type == "approve_claim":
        agent_id = apply_payload.get("agent_id")
        return {
            "claims": [
                {
                    "id": row.get("id"),
                    "type": row.get("type"),
                    "text": row.get("text"),
                    "status": row.get("status", "open"),
                    "confidence": row.get("confidence"),
                    "tags": row.get("tags") or [],
                    "created_by": row.get("created_by", agent_id),
                    "superseded_by": row.get("superseded_by"),
                }
                for row in apply_payload.get("claims") or []
            ],
            "sources": [
                {
                    "id": row.get("id"),
                    "source_type": row.get("source_type"),
                    "url": row.get("url"),
                    "storage_path": row.get("storage_path"),
                    "excerpt": row.get("excerpt"),
                    "hash": row.get("hash"),
                    "created_by": row.get("created_by", agent_id),
                }
                for row in apply_payload.get("sources") or []
            ],
            "claim_evidence": [
                {
                    "claim_id": row.get("claim_id"),
                    "source_id": row.get("source_id"),
                    "role": row.get("role") or "grounds",
                    "weight": row.get("weight"),
                    "created_by": row.get("created_by", agent_id),
                }
                for row in apply_payload.get("evidence") or []
            ],
            "claim_edges": [
                {
                    "from_claim": row.get("from_claim"),
                    "to_claim": row.get("to_claim"),
                    "edge_type": row.get("edge_type"),
                    "weight": row.get("weight"),
                    "created_by": row.get("created_by", agent_id),
                }
                for row in apply_payload.get("edges") or []
            ],
            "reasoning_tools": [
                {
                    "id": row.get("id"),
                    "agent_id": row.get("agent_id", agent_id),
                    "name": row.get("name"),
                    "description": row.get("description"),
                    "category": row.get("category"),
                }
                for row in apply_payload.get("reasoning_tools") or []
            ],
        }
    if proposal_type == "revise_strategy":
        strategy = apply_payload.get("strategy") or {}
        agent_id = apply_payload.get("agent_id")
        return {
            "strategies": [
                {
                    "agent_id": agent_id,
                    "diagnosis": strategy.get("diagnosis"),
                    "guiding_policy": strategy.get("guiding_policy"),
                    "proximate_objectives": strategy.get("proximate_objectives"),
                    "active": True,
                }
            ],
            "strategy_nodes": [
                {
                    "agent_id": agent_id,
                    "node_type": row.get("node_type"),
                    "title": row.get("title"),
                    "body": row.get("body"),
                    "rank": row.get("rank", 1),
                    "status": "active",
                    "horizon": None,
                    "measure": None,
                    "source_ref": None,
                    "metadata": {},
                }
                for row in apply_payload.get("strategy_nodes") or []
            ],
        }
    raise ValueError(f"replay receipt does not support proposal_type {proposal_type!r}")


def _assert_semantic_rows_match(proposal: dict[str, Any], rows: dict[str, list[dict[str, Any]]]) -> None:
    expected = _expected_semantic_rows(proposal)
    for table, expected_rows in expected.items():
        fields = tuple(expected_rows[0]) if expected_rows else ()
        actual_projected = [_semantic_projection(row, fields) for row in rows.get(table, [])]
        if sorted(expected_rows, key=_canonical_json) != sorted(actual_projected, key=_canonical_json):
            raise ValueError(
                f"row-level postflight semantic mismatch for {table}; "
                "canonical values do not match the immutable strict payload"
            )


def build_replay_receipt(
    proposal: dict[str, Any],
    canonical_rows: dict[str, Any],
    *,
    apply_sql_sha256: str,
    apply_sql_source: str,
    exported_at_utc: str | None = None,
) -> dict[str, Any]:
    if len(apply_sql_sha256) != 64 or any(character not in "0123456789abcdef" for character in apply_sql_sha256):
        raise ValueError("apply_sql_sha256 must be a lowercase SHA-256 hex digest")
    if apply_sql_source not in {"exact_executed_sql", "reconstructed_current_engine"}:
        raise ValueError("invalid apply_sql_source")
    apply_payload = _require_applied_strict_proposal(proposal)
    expected = expected_row_counts(proposal)
    rows = _normalize_rows(canonical_rows, expected)
    actual = {table: len(table_rows) for table, table_rows in rows.items()}
    mismatches = {
        table: {"expected": expected[table], "actual": actual.get(table, 0)}
        for table in expected
        if actual.get(table, 0) != expected[table]
    }
    if mismatches:
        raise ValueError(f"row-level postflight count mismatch: {mismatches}")
    _assert_semantic_rows_match(proposal, rows)

    proposal_envelope = {
        "id": str(proposal.get("id")),
        "proposal_type": proposal.get("proposal_type"),
        "status": proposal.get("status"),
        "payload": proposal.get("payload"),
        "reviewed_by_handle": proposal.get("reviewed_by_handle"),
        "reviewed_by_agent_id": proposal.get("reviewed_by_agent_id"),
        "reviewed_at": proposal.get("reviewed_at"),
        "review_note": proposal.get("review_note"),
        "applied_by_handle": proposal.get("applied_by_handle"),
        "applied_by_agent_id": proposal.get("applied_by_agent_id"),
        "applied_at": proposal.get("applied_at"),
    }
    row_hashes = {table: [_sha256_json(row) for row in table_rows] for table, table_rows in rows.items()}
    table_hashes = {table: _sha256_json(table_rows) for table, table_rows in rows.items()}
    replay_material = {
        "receipt_contract_version": RECEIPT_CONTRACT_VERSION,
        "apply_engine": {
            "apply_sql_sha256": apply_sql_sha256,
            "source": apply_sql_source,
        },
        "proposal": proposal_envelope,
        "canonical_rows": rows,
    }
    return {
        **replay_material,
        "artifact": "kb_apply_replay_receipt",
        "exported_at_utc": exported_at_utc or datetime.now(timezone.utc).isoformat(),
        "expected_row_counts": expected,
        "actual_row_counts": actual,
        "checks": {
            "proposal_applied": True,
            "strict_apply_payload_present": True,
            "applied_at_present": True,
            "row_level_postflight_exact": True,
            "semantic_rows_match_strict_payload": True,
        },
        "hashes": {
            "apply_payload_sha256": _sha256_json(apply_payload),
            "proposal_payload_sha256": _sha256_json(proposal.get("payload")),
            "canonical_rows_sha256": _sha256_json(rows),
            "row_sha256": row_hashes,
            "table_sha256": table_hashes,
            "replay_material_sha256": _sha256_json(replay_material),
        },
        "replay_ready": True,
        "privacy": "private receipt; may contain claim bodies and source excerpts; do not commit",
    }


def resolve_receipt_path(
    proposal_id: str,
    *,
    receipt_out: str | None = None,
    receipt_dir: str | None = None,
) -> Path:
    if receipt_out:
        return Path(receipt_out)
    root = Path(receipt_dir or os.environ.get("KB_APPLY_RECEIPT_DIR") or DEFAULT_RECEIPT_DIR)
    return root / f"{proposal_id}.json"


def write_private_receipt(path: Path, receipt: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    payload = json.dumps(receipt, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
    fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent))
    try:
        os.fchmod(fd, 0o600)
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            handle.write(payload)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary_name, path)
        os.chmod(path, 0o600)
    except Exception:
        try:
            os.close(fd)
        except OSError:
            pass
        try:
            os.unlink(temporary_name)
        except FileNotFoundError:
            pass
        raise