#!/usr/bin/env python3
"""Apply an approved kb_stage proposal into canonical public.* rows.

Stage 2 of the KB apply pipeline:  approve -> APPLY -> render -> surface.

Governance boundary
-------------------
This tool connects to Postgres as the narrow ``kb_apply`` login role, never as
superuser. ``kb_apply`` can write only ``public.strategies``,
``public.strategy_nodes``, ``public.claim_evidence``, ``public.claim_edges`` and
update the ``kb_stage.kb_proposals`` ledger. It holds no DELETE, no DDL, and no
access to any other table (agents, budgets, personas, ...). This enforces
"agents propose, but do not self-apply" at the database boundary rather than by
convention. The tool is an operator, invoked by a human/reviewer after approval;
it is deliberately not part of the agent runtime.

Flow
----
1. Read the proposal (must be ``status='approved'``).
2. Dispatch by ``proposal_type`` to build the canonical write from a strict
   ``apply_payload`` contract carried on the proposal payload.
3. Emit a single transaction: canonical writes, then a DO block that flips the
   ledger to ``applied`` asserting rowcount=1 (the concurrent double-apply
   guard), stamps ``applied_by_agent_id`` as a real FK, and verifies invariants
   (RAISE on violation -> rollback), COMMIT.

Prerequisites (one-time, superuser): scripts/kb_apply_prereqs.sql bootstraps the
``kb-apply`` service-agent row, grants ``kb_apply`` SELECT on public.agents, and
ensures the one-active-strategy unique index. Run it before the first apply.

``--dry-run`` prints the exact SQL and does not connect for writes.

Contract (v1)
-------------
Freeform Leo eval packets are NOT applied directly. Each proposal must carry a
strict ``apply_payload`` object; a separate normalizer produces it. Shapes:

  revise_strategy:
    {"apply_payload": {
        "agent_id": "<uuid>",
        "strategy": {"diagnosis": str, "guiding_policy": str,
                     "proximate_objectives": [ ... ]},
        "strategy_nodes": [{"node_type": "diagnosis|policy|objective",
                            "title": str, "body": str, "rank": int}, ...]}}

  add_edge:
    {"apply_payload": {"from_claim": "<uuid>", "to_claim": "<uuid>",
                       "edge_type": "<edge_type>", "weight": <0..1|null>}}

  attach_evidence:
    {"apply_payload": {"evidence": [{"claim_id": "<uuid>",
                                     "source_id": "<uuid>",
                                     "role": "<evidence_role>",
                                     "weight": <0..1|null>}, ...]}}

Note: attach_evidence requires an existing ``source_id`` per item. kb_apply has
no INSERT on ``public.sources``, so source creation/dedup is intentionally out of
scope for this tool and handled upstream (a follow-up + explicit grant).
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

DEFAULT_SECRETS_FILE = "/home/teleo/.hermes/profiles/leoclean/secrets/kb-apply.env"
DEFAULT_CONTAINER = "teleo-pg"
DEFAULT_DB = "teleo"
DEFAULT_HOST = "127.0.0.1"
DEFAULT_ROLE = "kb_apply"

# Handle of the service-agent row that performs canonical writes. Bootstrapped
# once by superuser (scripts/kb_apply_prereqs.sql); kb_apply holds SELECT-only on
# public.agents so it can resolve this into applied_by_agent_id, never INSERT.
SERVICE_AGENT_HANDLE = "kb-apply"

APPLYABLE_TYPES = ("revise_strategy", "add_edge", "attach_evidence")


# --------------------------------------------------------------------------- #
# SQL helpers                                                                  #
# --------------------------------------------------------------------------- #
def sql_literal(value: Any) -> str:
    """Render a Python value as a safe SQL literal."""
    if value is None:
        return "null"
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float)):
        return repr(value)
    return "'" + str(value).replace("'", "''") + "'"


def _jsonb(value: Any) -> str:
    return sql_literal(json.dumps(value, sort_keys=True)) + "::jsonb"


def _ledger_and_verify(proposal_id: str, applied_by: str, extra_checks: str) -> str:
    """Flip the proposal to applied and verify invariants, inside the txn.

    The flip runs in a DO block that asserts exactly one row moved from
    'approved' to 'applied' (``GET DIAGNOSTICS ... = row_count``). This closes
    the concurrent double-apply race: ``load_proposal`` (read) and this flip
    (write) are separate statements, so a ``SELECT ... FOR UPDATE`` row lock
    would not span them -- but only one concurrent apply can match
    ``status='approved'``, so rowcount=1 is the authoritative guard. A loser
    sees 0 rows and RAISEs -> the whole transaction rolls back.

    ``applied_by_agent_id`` is stamped as a real FK resolved from
    ``public.agents`` by handle (requires SELECT public.agents; the service-agent
    row is bootstrapped once by superuser -- see scripts/kb_apply_prereqs.sql).
    The resolve is a hard lookup: an unresolved handle RAISEs -> rollback, never
    a silent NULL FK (the column is ON DELETE SET NULL, so NULL is legal and
    would otherwise pass unnoticed).
    """
    return f"""
do $flip$
declare
  flipped int;
  resolved_agent_id uuid;
begin
  select id into resolved_agent_id
    from public.agents
   where handle = {sql_literal(applied_by)};
  if resolved_agent_id is null then
    raise exception 'apply_proposal: applied_by handle % does not resolve to a public.agents row; refusing to stamp a NULL applied_by_agent_id', {sql_literal(applied_by)};
  end if;
  update kb_stage.kb_proposals
     set status = 'applied',
         applied_by_handle = {sql_literal(applied_by)},
         applied_by_agent_id = resolved_agent_id,
         applied_at = now(),
         updated_at = now()
   where id = {sql_literal(proposal_id)}::uuid
     and status = 'approved';
  get diagnostics flipped = row_count;
  if flipped <> 1 then
    raise exception 'apply_proposal: ledger flip affected % row(s), expected 1; proposal % was not approved (already applied or a concurrent apply won)', flipped, {sql_literal(proposal_id)};
  end if;
{extra_checks}
end
$flip$;
""".rstrip()


# --------------------------------------------------------------------------- #
# Per-type SQL builders (pure; unit-tested without a DB)                       #
# --------------------------------------------------------------------------- #
def build_revise_strategy_sql(
    apply_payload: Dict[str, Any], proposal_id: str, reviewer: Optional[str]
) -> str:
    agent_id = apply_payload.get("agent_id")
    strategy = apply_payload.get("strategy") or {}
    nodes: List[Dict[str, Any]] = apply_payload.get("strategy_nodes") or []
    if not agent_id:
        raise ValueError("revise_strategy apply_payload requires 'agent_id'")
    for key in ("diagnosis", "guiding_policy", "proximate_objectives"):
        if strategy.get(key) is None:
            raise ValueError(f"revise_strategy apply_payload.strategy requires '{key}'")
    if not nodes:
        raise ValueError("revise_strategy apply_payload requires non-empty 'strategy_nodes'")

    node_values = []
    for n in nodes:
        nt = n.get("node_type")
        if nt not in ("diagnosis", "policy", "objective"):
            raise ValueError(f"invalid strategy node_type: {nt!r}")
        if not n.get("title") or n.get("body") is None:
            raise ValueError("each strategy_node requires 'title' and 'body'")
        node_values.append(
            f"({sql_literal(agent_id)}::uuid, {sql_literal(nt)}, "
            f"{sql_literal(n['title'])}, {sql_literal(n['body'])}, "
            f"{int(n.get('rank', 1))})"
        )

    node_values_sql = ",\n  ".join(node_values)
    canonical = f"""
-- deactivate the current active strategy for this agent (one_active invariant)
update public.strategies
   set active = false
 where agent_id = {sql_literal(agent_id)}::uuid and active;

-- insert the new strategy as the next version, active
insert into public.strategies
       (agent_id, diagnosis, guiding_policy, proximate_objectives, version, active)
select {sql_literal(agent_id)}::uuid,
       {sql_literal(strategy['diagnosis'])},
       {sql_literal(strategy['guiding_policy'])},
       {_jsonb(strategy['proximate_objectives'])},
       coalesce(max(version), 0) + 1,
       true
  from public.strategies
 where agent_id = {sql_literal(agent_id)}::uuid;

-- retire the agent's current strategy nodes (shared NULL-agent nodes untouched)
update public.strategy_nodes
   set status = 'retired', updated_at = now()
 where agent_id = {sql_literal(agent_id)}::uuid and status <> 'retired';

-- insert the new node set
insert into public.strategy_nodes (agent_id, node_type, title, body, rank)
values
  {node_values_sql};
""".rstrip()

    checks = f"""  if (select count(*) from public.strategies
       where agent_id = {sql_literal(agent_id)}::uuid and active) <> 1 then
    raise exception 'apply_proposal: expected exactly one active strategy for agent %', {sql_literal(agent_id)};
  end if;"""

    return _wrap_txn(canonical, _ledger_and_verify(proposal_id, reviewer, checks))


def build_add_edge_sql(
    apply_payload: Dict[str, Any], proposal_id: str, reviewer: Optional[str]
) -> str:
    f = apply_payload.get("from_claim")
    t = apply_payload.get("to_claim")
    et = apply_payload.get("edge_type")
    w = apply_payload.get("weight")
    if not (f and t and et):
        raise ValueError("add_edge apply_payload requires 'from_claim', 'to_claim', 'edge_type'")
    if f == t:
        raise ValueError("add_edge from_claim and to_claim must differ")

    canonical = f"""
-- insert the edge unless an identical (from,to,type) edge already exists
insert into public.claim_edges (from_claim, to_claim, edge_type, weight)
select {sql_literal(f)}::uuid, {sql_literal(t)}::uuid,
       {sql_literal(et)}::edge_type, {sql_literal(w)}
 where not exists (
   select 1 from public.claim_edges
    where from_claim = {sql_literal(f)}::uuid
      and to_claim = {sql_literal(t)}::uuid
      and edge_type = {sql_literal(et)}::edge_type);
""".rstrip()

    checks = f"""  if not exists (select 1 from public.claim_edges
       where from_claim = {sql_literal(f)}::uuid
         and to_claim = {sql_literal(t)}::uuid
         and edge_type = {sql_literal(et)}::edge_type) then
    raise exception 'apply_proposal: claim_edge not present after apply';
  end if;"""

    return _wrap_txn(canonical, _ledger_and_verify(proposal_id, reviewer, checks))


def build_attach_evidence_sql(
    apply_payload: Dict[str, Any], proposal_id: str, reviewer: Optional[str]
) -> str:
    evidence: List[Dict[str, Any]] = apply_payload.get("evidence") or []
    if not evidence:
        raise ValueError("attach_evidence apply_payload requires non-empty 'evidence'")

    statements = []
    checks = []
    for i, ev in enumerate(evidence):
        claim_id = ev.get("claim_id")
        source_id = ev.get("source_id")
        role = ev.get("role") or "grounds"
        weight = ev.get("weight")
        if not claim_id:
            raise ValueError(f"evidence[{i}] requires 'claim_id'")
        if not source_id:
            raise ValueError(
                f"evidence[{i}] requires an existing 'source_id'. "
                "kb_apply cannot create public.sources; mint the source upstream first."
            )
        statements.append(
            f"""insert into public.claim_evidence (claim_id, source_id, role, weight)
select {sql_literal(claim_id)}::uuid, {sql_literal(source_id)}::uuid,
       {sql_literal(role)}::evidence_role, {sql_literal(weight)}
 where not exists (
   select 1 from public.claim_evidence
    where claim_id = {sql_literal(claim_id)}::uuid
      and source_id = {sql_literal(source_id)}::uuid
      and role = {sql_literal(role)}::evidence_role);"""
        )
        checks.append(
            f"""  if not exists (select 1 from public.claim_evidence
       where claim_id = {sql_literal(claim_id)}::uuid
         and source_id = {sql_literal(source_id)}::uuid
         and role = {sql_literal(role)}::evidence_role) then
    raise exception 'apply_proposal: evidence row % not present after apply', {i};
  end if;"""
        )

    canonical = "\n".join(statements)
    return _wrap_txn(canonical, _ledger_and_verify(proposal_id, reviewer, "\n".join(checks)))


def _wrap_txn(canonical_sql: str, ledger_sql: str) -> str:
    return f"begin;\n{canonical_sql}\n{ledger_sql}\ncommit;\n"


BUILDERS = {
    "revise_strategy": build_revise_strategy_sql,
    "add_edge": build_add_edge_sql,
    "attach_evidence": build_attach_evidence_sql,
}


def build_apply_sql(proposal: Dict[str, Any], applied_by: Optional[str]) -> str:
    ptype = proposal["proposal_type"]
    if ptype not in BUILDERS:
        raise ValueError(f"apply_proposal cannot apply proposal_type {ptype!r}")
    payload = proposal.get("payload") or {}
    apply_payload = payload.get("apply_payload")
    if apply_payload is None:
        raise ValueError(
            "proposal payload has no 'apply_payload' (contract v1). "
            "Run the normalizer to produce apply_payload before applying."
        )
    return BUILDERS[ptype](apply_payload, proposal["id"], applied_by or SERVICE_AGENT_HANDLE)


def assert_applyable(proposal: Dict[str, Any]) -> None:
    status = proposal.get("status")
    if status == "applied":
        raise SystemExit(f"proposal {proposal['id']} is already applied (idempotent no-op)")
    if status != "approved":
        raise SystemExit(
            f"proposal {proposal['id']} has status {status!r}; only 'approved' proposals apply"
        )


# --------------------------------------------------------------------------- #
# Execution (docker exec as the kb_apply role)                                 #
# --------------------------------------------------------------------------- #
def load_password(secrets_file: str) -> str:
    path = Path(secrets_file)
    if not path.is_file():
        raise SystemExit(f"secrets file not found: {secrets_file}")
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, val = line.partition("=")
        if key.strip() in ("KB_APPLY_PASSWORD", "PGPASSWORD"):
            return val.strip().strip('"').strip("'")
    raise SystemExit(f"no KB_APPLY_PASSWORD/PGPASSWORD entry in {secrets_file}")


def run_psql(args: argparse.Namespace, sql: str, password: str) -> str:
    command = [
        "docker", "exec", "-e", "PGPASSWORD", "-i", args.container,
        "psql", "-U", args.role, "-h", args.host, "-d", args.db,
        "-v", "ON_ERROR_STOP=1", "-At", "-q",
    ]
    result = subprocess.run(
        command, input=sql, text=True, capture_output=True,
        env={"PGPASSWORD": password, "PATH": "/usr/bin:/bin:/usr/local/bin"},
        check=False,
    )
    if result.returncode != 0:
        raise SystemExit(
            f"psql failed ({result.returncode})\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
        )
    return result.stdout


def load_proposal(args: argparse.Namespace, password: str) -> Dict[str, Any]:
    sql = f"""select jsonb_build_object(
        'id', id::text,
        'proposal_type', proposal_type,
        'status', status,
        'payload', payload)::text
      from kb_stage.kb_proposals
     where id = {sql_literal(args.proposal_id)}::uuid;"""
    out = run_psql(args, sql, password).strip()
    if not out:
        raise SystemExit(f"proposal not found: {args.proposal_id}")
    return json.loads(out)


def parse_args(argv: List[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("proposal_id", help="UUID of the approved proposal to apply")
    p.add_argument(
        "--applied-by",
        default=SERVICE_AGENT_HANDLE,
        help="agent handle recorded as applied_by_handle and resolved to the "
        "applied_by_agent_id FK (default: the kb-apply service agent)",
    )
    p.add_argument("--dry-run", action="store_true", help="print the SQL, do not execute")
    p.add_argument("--secrets-file", default=DEFAULT_SECRETS_FILE)
    p.add_argument("--container", default=DEFAULT_CONTAINER)
    p.add_argument("--db", default=DEFAULT_DB)
    p.add_argument("--host", default=DEFAULT_HOST)
    p.add_argument("--role", default=DEFAULT_ROLE)
    return p.parse_args(argv)


def main(argv: Optional[List[str]] = None) -> int:
    args = parse_args(sys.argv[1:] if argv is None else argv)

    if args.dry_run:
        # Dry-run still needs the proposal to build SQL; read it as kb_apply.
        password = load_password(args.secrets_file)
        proposal = load_proposal(args, password)
        assert_applyable(proposal)
        print(build_apply_sql(proposal, args.applied_by))
        return 0

    password = load_password(args.secrets_file)
    proposal = load_proposal(args, password)
    assert_applyable(proposal)
    sql = build_apply_sql(proposal, args.applied_by)
    run_psql(args, sql, password)
    print(f"applied proposal {proposal['id']} ({proposal['proposal_type']})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())