#!/usr/bin/env python3
"""Generate a guarded production apply packet for a mapped rich KB plan.

This is a packet builder, not an executor. It turns the reviewed/mapped rich
proposal creation plan into reproducible SQL files:

* preflight SQL: proposal status, required existing claim IDs, generated-row
  absence, and expected counts;
* rollback rehearsal SQL: the write transaction with a final ROLLBACK;
* optional commit SQL: the same transaction with COMMIT, for separately
  authorized execution only;
* postflight SQL: exact generated-row count/readback;
* rollback-delete SQL: exact reverse deletes for the generated rows.

The packet marks original rich proposals as applied only when explicitly
requested and the plan has no remaining blocked fragments or unresolved design
decisions. Mapped subsets must leave the original proposals approved, because
marking a partial packet applied would overstate DB state.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import kb_rich_proposal_creation_plan as creation_plan  # noqa: E402


def _uuid_values(values: list[str], *, column_name: str = "id") -> str:
    if not values:
        return f"(select null::uuid as {column_name} where false)"
    rows = ",\n    ".join(f"('{value}'::uuid)" for value in values)
    return f"(values\n    {rows}\n  ) as v({column_name})"


def _edge_values(rows: list[dict[str, Any]]) -> str:
    if not rows:
        return "(select null::uuid as from_claim, null::uuid as to_claim, null::edge_type as edge_type where false)"
    values = ",\n    ".join(
        f"('{row['from_claim']}'::uuid, '{row['to_claim']}'::uuid, '{row['edge_type']}'::edge_type)"
        for row in rows
    )
    return f"(values\n    {values}\n  ) as v(from_claim, to_claim, edge_type)"


def _evidence_values(rows: list[dict[str, Any]]) -> str:
    if not rows:
        return "(select null::uuid as claim_id, null::uuid as source_id, null::evidence_role as role where false)"
    values = ",\n    ".join(
        f"('{row['claim_id']}'::uuid, '{row['source_id']}'::uuid, '{row['role']}'::evidence_role)"
        for row in rows
    )
    return f"(values\n    {values}\n  ) as v(claim_id, source_id, role)"


def load_creation_plans(path: Path) -> list[dict[str, Any]]:
    data = json.loads(path.read_text(encoding="utf-8"))
    plans = data.get("creation_plans") if isinstance(data, dict) else None
    if not isinstance(plans, list):
        raise SystemExit(f"{path} does not contain creation_plans")
    return plans


def collect_rows(plans: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]:
    rows = {"reasoning_tools": [], "claims": [], "sources": [], "claim_evidence": [], "claim_edges": []}
    for plan in plans:
        planned = plan.get("planned_rows") if isinstance(plan.get("planned_rows"), dict) else {}
        for key in rows:
            value = planned.get(key)
            if isinstance(value, list):
                rows[key].extend(row for row in value if isinstance(row, dict))
    return rows


def collect_required_existing_claims(rows: dict[str, list[dict[str, Any]]]) -> list[dict[str, str]]:
    generated_claim_ids = {row["id"] for row in rows["claims"]}
    required: dict[str, str] = {}
    for edge in rows["claim_edges"]:
        for side in ("from_claim", "to_claim"):
            claim_id = edge.get(side)
            if claim_id and claim_id not in generated_claim_ids:
                required.setdefault(str(claim_id), f"claim_edges.{side}")
    for evidence in rows["claim_evidence"]:
        claim_id = evidence.get("claim_id")
        if claim_id and claim_id not in generated_claim_ids:
            required.setdefault(str(claim_id), "claim_evidence.claim_id")
    return [{"id": claim_id, "reason": reason} for claim_id, reason in sorted(required.items())]


def collect_blocked_fragments(plans: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [
        fragment
        for plan in plans
        for fragment in plan.get("blocked_fragments", [])
        if isinstance(fragment, dict)
    ]


def collect_unresolved_design_decisions(plans: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [
        item
        for plan in plans
        for item in plan.get("unresolved_design_decisions", [])
        if isinstance(item, dict)
    ]


def assert_can_mark_original_proposals_applied(plans: list[dict[str, Any]]) -> None:
    blocked = collect_blocked_fragments(plans)
    unresolved = collect_unresolved_design_decisions(plans)
    if blocked or unresolved:
        raise ValueError(
            "cannot mark original proposals applied for a partial packet: "
            f"{len(blocked)} blocked fragment(s), {len(unresolved)} unresolved design decision(s)"
        )


def build_generated_counts_sql(rows: dict[str, list[dict[str, Any]]]) -> str:
    reasoning_tool_ids = [row["id"] for row in rows["reasoning_tools"]]
    claim_ids = [row["id"] for row in rows["claims"]]
    source_ids = [row["id"] for row in rows["sources"]]
    return f"""select
  (select count(*) from public.reasoning_tools where id in (select id from {_uuid_values(reasoning_tool_ids)})) as generated_reasoning_tool_rows,
  (select count(*) from public.claims where id in (select id from {_uuid_values(claim_ids)})) as generated_claim_rows,
  (select count(*) from public.sources where id in (select id from {_uuid_values(source_ids)})) as generated_source_rows,
  (select count(*) from public.claim_edges e join {_edge_values(rows['claim_edges'])} using (from_claim, to_claim, edge_type)) as generated_edge_rows,
  (select count(*) from public.claim_evidence ce join {_evidence_values(rows['claim_evidence'])} using (claim_id, source_id, role)) as generated_evidence_rows;
"""


def build_mark_original_proposals_sql(plans: list[dict[str, Any]], *, commit: bool, applied_by_handle: str) -> str:
    proposal_ids = [str(plan.get("proposal_id")) for plan in plans if plan.get("proposal_id")]
    if not proposal_ids:
        raise ValueError("cannot mark original proposals applied without proposal ids")
    terminator = "commit;" if commit else "rollback;"
    return f"""begin;

do $$
declare
  flipped int;
begin
  update kb_stage.kb_proposals
     set status = 'applied',
         applied_by_handle = {creation_plan._sql_literal(applied_by_handle)},
         applied_at = now(),
         updated_at = now()
   where id in (select id from {_uuid_values(proposal_ids)})
     and status = 'approved';

  get diagnostics flipped = row_count;
  if flipped <> {len(proposal_ids)} then
    raise exception 'rich proposal packet ledger flip affected % row(s), expected {len(proposal_ids)}; original proposals were not all approved', flipped;
  end if;
end $$;

{terminator}
"""


def build_preflight_sql(plans: list[dict[str, Any]], rows: dict[str, list[dict[str, Any]]]) -> str:
    proposal_ids = [str(plan.get("proposal_id")) for plan in plans if plan.get("proposal_id")]
    required_claims = collect_required_existing_claims(rows)
    required_values = (
        ",\n    ".join(f"('{item['id']}'::uuid, {creation_plan._sql_literal(item['reason'])})" for item in required_claims)
        or "(null::uuid, null::text)"
    )
    expected = {
        "reasoning_tools": len(rows["reasoning_tools"]),
        "claims": len(rows["claims"]),
        "sources": len(rows["sources"]),
        "claim_edges": len(rows["claim_edges"]),
        "claim_evidence": len(rows["claim_evidence"]),
    }
    return f"""\\set ON_ERROR_STOP on
-- Preflight for mapped rich proposal apply packet. Read-only.

select id::text, status, proposal_type, reviewed_by_handle, reviewed_at
from kb_stage.kb_proposals
where id in (select id from {_uuid_values(proposal_ids)})
order by id;

with required(id, reason) as (
  values
    {required_values}
)
select 'missing_required_existing_claims' as check_name, id::text, reason
from required
where id is not null
  and not exists (select 1 from public.claims where public.claims.id = required.id)
order by id;

-- Expected generated rows after authorized commit:
select {expected['reasoning_tools']}::int as expected_reasoning_tools,
       {expected['claims']}::int as expected_claims,
       {expected['sources']}::int as expected_sources,
       {expected['claim_edges']}::int as expected_claim_edges,
       {expected['claim_evidence']}::int as expected_claim_evidence;

-- Must be all zero before production apply.
{build_generated_counts_sql(rows)}"""


def build_write_sql(
    plans: list[dict[str, Any]],
    *,
    commit: bool,
    mark_original_proposals_applied: bool,
    applied_by_handle: str,
) -> str:
    sql = "\n\n".join(plan.get("rollback_rehearsal_sql", "") for plan in plans)
    terminator = "commit;" if commit else "rollback;"
    write_sql = sql.replace("rollback;", terminator)
    if mark_original_proposals_applied:
        write_sql = (
            write_sql.rstrip()
            + "\n\n-- Full-coverage packet ledger update: mark original approved proposal(s) applied only after generated rows exist.\n"
            + build_mark_original_proposals_sql(plans, commit=commit, applied_by_handle=applied_by_handle)
        )
    return write_sql


def build_postflight_sql(
    rows: dict[str, list[dict[str, Any]]],
    *,
    proposal_ids: list[str],
    include_original_proposal_readback: bool,
) -> str:
    reasoning_tool_ids = [row["id"] for row in rows["reasoning_tools"]]
    claim_ids = [row["id"] for row in rows["claims"]]
    source_ids = [row["id"] for row in rows["sources"]]
    return f"""\\set ON_ERROR_STOP on
-- Postflight row-level readback for mapped rich proposal apply packet.

{build_generated_counts_sql(rows)}

select id::text, agent_id::text, name, category, left(description, 500) as description
from public.reasoning_tools
where id in (select id from {_uuid_values(reasoning_tool_ids)})
order by id;

select id::text, type, confidence, left(text, 220) as text
from public.claims
where id in (select id from {_uuid_values(claim_ids)})
order by id;

select id::text, source_type, left(coalesce(excerpt, ''), 220) as excerpt
from public.sources
where id in (select id from {_uuid_values(source_ids)})
order by id;

select e.from_claim::text, e.to_claim::text, e.edge_type::text
from public.claim_edges e
join {_edge_values(rows['claim_edges'])} using (from_claim, to_claim, edge_type)
order by e.from_claim, e.to_claim, e.edge_type;

select ce.claim_id::text, ce.source_id::text, ce.role::text
from public.claim_evidence ce
join {_evidence_values(rows['claim_evidence'])} using (claim_id, source_id, role)
order by ce.claim_id, ce.source_id, ce.role;

select id::text, superseded_by::text
from public.claims
where superseded_by in (select id from {_uuid_values(claim_ids)})
order by id;
""" + (
        f"""
select id::text, status, applied_by_handle, applied_at
from kb_stage.kb_proposals
where id in (select id from {_uuid_values(proposal_ids)})
order by id;
"""
        if include_original_proposal_readback
        else ""
    )


def build_delete_rollback_sql(
    rows: dict[str, list[dict[str, Any]]],
    *,
    proposal_ids: list[str],
    restore_original_proposals_approved: bool,
) -> str:
    reasoning_tool_ids = [row["id"] for row in rows["reasoning_tools"]]
    claim_ids = [row["id"] for row in rows["claims"]]
    source_ids = [row["id"] for row in rows["sources"]]
    return f"""\\set ON_ERROR_STOP on
-- Emergency rollback for the mapped generated rows only.
-- Run only before any later production data depends on these generated rows.
begin;

delete from public.claim_evidence ce
using {_evidence_values(rows['claim_evidence'])}
where ce.claim_id = v.claim_id and ce.source_id = v.source_id and ce.role = v.role;

delete from public.claim_edges e
using {_edge_values(rows['claim_edges'])}
where e.from_claim = v.from_claim and e.to_claim = v.to_claim and e.edge_type = v.edge_type;

update public.claims
set superseded_by = null,
    updated_at = now()
where superseded_by in (select id from {_uuid_values(claim_ids)});

delete from public.sources
where id in (select id from {_uuid_values(source_ids)});

delete from public.claims
where id in (select id from {_uuid_values(claim_ids)});

delete from public.reasoning_tools
where id in (select id from {_uuid_values(reasoning_tool_ids)});

{build_restore_original_proposals_sql(proposal_ids) if restore_original_proposals_approved else ""}

commit;
"""


def build_restore_original_proposals_sql(proposal_ids: list[str]) -> str:
    if not proposal_ids:
        return ""
    return f"""do $$
declare
  restored int;
begin
  update kb_stage.kb_proposals
     set status = 'approved',
         applied_by_handle = null,
         applied_by_agent_id = null,
         applied_at = null,
         updated_at = now()
   where id in (select id from {_uuid_values(proposal_ids)})
     and status = 'applied';

  get diagnostics restored = row_count;
  if restored <> {len(proposal_ids)} then
    raise exception 'rich proposal packet ledger rollback restored % row(s), expected {len(proposal_ids)}; original proposals were not all applied', restored;
  end if;
end $$;"""


def build_packet(
    plan_path: Path,
    *,
    emit_commit_sql: bool,
    mark_original_proposals_applied: bool = False,
    applied_by_handle: str = "codex-packet",
) -> dict[str, Any]:
    plans = load_creation_plans(plan_path)
    if mark_original_proposals_applied:
        assert_can_mark_original_proposals_applied(plans)
    rows = collect_rows(plans)
    proposal_ids = [str(plan.get("proposal_id")) for plan in plans if plan.get("proposal_id")]
    remaining_blocked_fragments = collect_blocked_fragments(plans)
    unresolved_design_decisions = collect_unresolved_design_decisions(plans)
    claim_ceiling = (
        "Production apply packet for full covered proposal contract; authorized commit SQL marks original proposal(s) "
        "applied after generated rows are written."
        if mark_original_proposals_applied
        else (
            "Production apply packet for mapped supported subset only; original rich proposals are not marked applied "
            "unless the packet fully covers the approved reasoning-tool/claim/source contract."
        )
    )
    return {
        "packet_version": 1,
        "plan_path": str(plan_path),
        "proposal_ids": proposal_ids,
        "claim_ceiling": claim_ceiling,
        "row_counts": {key: len(value) for key, value in rows.items()},
        "required_existing_claims": collect_required_existing_claims(rows),
        "remaining_blocked_fragments": remaining_blocked_fragments,
        "unresolved_design_decisions": unresolved_design_decisions,
        "marks_original_proposals_applied": mark_original_proposals_applied,
        "applied_by_handle": applied_by_handle if mark_original_proposals_applied else None,
        "sql": {
            "preflight": build_preflight_sql(plans, rows),
            "apply_rollback_rehearsal": build_write_sql(
                plans,
                commit=False,
                mark_original_proposals_applied=mark_original_proposals_applied,
                applied_by_handle=applied_by_handle,
            ),
            "apply_commit": (
                build_write_sql(
                    plans,
                    commit=True,
                    mark_original_proposals_applied=mark_original_proposals_applied,
                    applied_by_handle=applied_by_handle,
                )
                if emit_commit_sql
                else None
            ),
            "postflight": build_postflight_sql(
                rows,
                proposal_ids=proposal_ids,
                include_original_proposal_readback=mark_original_proposals_applied,
            ),
            "delete_rollback": build_delete_rollback_sql(
                rows,
                proposal_ids=proposal_ids,
                restore_original_proposals_approved=mark_original_proposals_applied,
            ),
        },
    }


def write_packet_files(packet: dict[str, Any], output_dir: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    sql = packet["sql"]
    (output_dir / "production-preflight.sql").write_text(sql["preflight"], encoding="utf-8")
    (output_dir / "production-apply-rollback-rehearsal.sql").write_text(
        sql["apply_rollback_rehearsal"], encoding="utf-8"
    )
    if sql["apply_commit"]:
        (output_dir / "production-apply-authorized-commit.sql").write_text(sql["apply_commit"], encoding="utf-8")
    (output_dir / "production-postflight.sql").write_text(sql["postflight"], encoding="utf-8")
    (output_dir / "production-delete-rollback.sql").write_text(sql["delete_rollback"], encoding="utf-8")

    summary = {key: value for key, value in packet.items() if key != "sql"}
    (output_dir / "production-apply-packet.json").write_text(json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8")


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--plan-json", required=True, type=Path, help="mapped rich proposal creation plan JSON")
    parser.add_argument("--output-dir", type=Path, help="directory to write packet files")
    parser.add_argument("--emit-commit-sql", action="store_true", help="also write the authorized COMMIT SQL file")
    parser.add_argument(
        "--mark-original-proposals-applied",
        action="store_true",
        help="include guarded ledger SQL to mark original approved proposals applied; requires full coverage",
    )
    parser.add_argument(
        "--applied-by-handle",
        default="codex-packet",
        help="applied_by_handle value for guarded full-coverage ledger SQL",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(sys.argv[1:] if argv is None else argv)
    try:
        packet = build_packet(
            args.plan_json,
            emit_commit_sql=args.emit_commit_sql,
            mark_original_proposals_applied=args.mark_original_proposals_applied,
            applied_by_handle=args.applied_by_handle,
        )
    except ValueError as exc:
        raise SystemExit(str(exc)) from exc
    if args.output_dir:
        write_packet_files(packet, args.output_dir)
    print(json.dumps({key: value for key, value in packet.items() if key != "sql"}, indent=2, sort_keys=True))
    return 0


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