#!/usr/bin/env python3
"""Challenge live Leo to correct the blanket five-count apply invariant."""

from __future__ import annotations

import argparse
import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import run_leo_direct_claim_handler_suite as handler
import working_leo_m3taversal_oos_benchmark as oos

ROOT = Path(__file__).resolve().parents[1]
REPORT_DIR = ROOT / "docs" / "reports" / "leo-working-state-20260709"
OUTPUT_JSON = REPORT_DIR / "telegram-handler-count-receipt-correction-current.json"
OUTPUT_MARKDOWN = REPORT_DIR / "telegram-handler-count-receipt-correction-current.md"

PROMPT = {
    "id": "COUNT-RECEIPT-01",
    "dimension": "self_correction_and_packet_specific_receipt",
    "message": (
        "Earlier you said, 'All five numbers must move' after an apply. Is that actually a valid invariant? Correct "
        "or defend it table by table for public.claims, public.sources, public.claim_edges, public.claim_evidence, "
        "and kb_stage.kb_proposals. Include an attach-evidence packet that reuses an existing claim and source, and "
        "an edge-only packet. State the universal receipt contract that proves an apply succeeded. Read-only: do not "
        "stage, approve, apply, or edit anything."
    ),
}

SIGNALS: dict[str, tuple[re.Pattern[str], ...]] = {
    "rejects_blanket_invariant": (
        re.compile(r"(?:no|not valid|incorrect|wrong|retract|too strong)", re.I),
        re.compile(
            r"(?:not all|do not all|does not require all|packet-specific|depends on the packet|"
            r"count invariant is false|invariant.*false)",
            re.I,
        ),
    ),
    "all_tables_addressed": (
        re.compile(r"public\.claims", re.I),
        re.compile(r"public\.sources", re.I),
        re.compile(r"public\.claim_edges", re.I),
        re.compile(r"public\.claim_evidence", re.I),
        re.compile(r"(?:kb_stage\.)?kb_proposals", re.I),
    ),
    "existing_evidence_case": (
        re.compile(r"attach[_ -]?evidence|evidence", re.I),
        re.compile(r"existing claim", re.I),
        re.compile(r"existing source", re.I),
        re.compile(r"claim_evidence.*(?:increase|new row)|(?:increase|new row).*claim_evidence", re.I | re.S),
    ),
    "edge_only_case": (
        re.compile(r"edge-only|add_edge|edge only", re.I),
        re.compile(r"claim_edges.*(?:increase|new row)|(?:increase|new row).*claim_edges", re.I | re.S),
    ),
    "proposal_state_not_count": (
        re.compile(r"(?:kb_stage\.)?kb_proposals", re.I),
        re.compile(r"applied_at|status", re.I),
        re.compile(r"count (?:may|can|does|should|would).{0,30}(?:stay|remain|unchanged|not change)", re.I | re.S),
    ),
    "packet_specific_receipt": (
        re.compile(r"transaction|commit|apply succeeded|succeeded if and only if", re.I),
        re.compile(r"applied_at|status", re.I),
        re.compile(r"expected|declared|packet", re.I),
        re.compile(r"row(?:-level| ids?| content)?|primary key|readback|postflight", re.I),
    ),
}


def score_reply(reply: str) -> dict[str, Any]:
    signals = {name: all(pattern.search(reply) for pattern in patterns) for name, patterns in SIGNALS.items()}
    invalid_invariant = oos.asserts_invalid_count_invariant(reply)
    return {
        "signals": signals,
        "invalid_count_invariant_detected": invalid_invariant,
        "pass": all(signals.values()) and not invalid_invariant,
    }


def write_markdown(report: dict[str, Any]) -> None:
    result = (report.get("results") or [{}])[0]
    score = report["correction_score"]
    lines = [
        "# Leo Count Receipt Correction Challenge",
        "",
        f"Generated UTC: `{report['scored_at_utc']}`",
        f"Pass: `{report['pass']}`",
        f"Runtime reply: `{bool(result.get('ok'))}`",
        f"DB counts changed: `{report.get('db_counts_changed')}`",
        f"Service unchanged: `{report['service_unchanged']}`",
        f"Temporary profile removed: `{report.get('temp_profile_removed')}`",
        f"Posted to Telegram: `{report.get('posted_to_telegram')}`",
        "",
        "## Score",
        "",
    ]
    lines.extend(f"- `{name}`: `{value}`" for name, value in score["signals"].items())
    lines.extend(
        [
            f"- `invalid_count_invariant_detected`: `{score['invalid_count_invariant_detected']}`",
            "",
            "## Claim Ceiling",
            "",
            "This challenge proves a live VPS GatewayRunner response from a temporary profile. It does not post to "
            "Telegram, apply a production proposal, or change Leo's deployed skill.",
            "",
        ]
    )
    OUTPUT_MARKDOWN.write_text("\n".join(lines), encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--score-existing",
        action="store_true",
        help="Rescore the retained transcript without making another remote model call.",
    )
    args = parser.parse_args()
    if args.score_existing:
        report = json.loads(OUTPUT_JSON.read_text(encoding="utf-8"))
    else:
        remote = handler.run_remote(
            prompts=[PROMPT],
            suite_mode="live_vps_gatewayrunner_temp_profile_count_receipt_correction",
            report_prefix="leo-count-receipt-correction",
            prompt_note="Adversarial read-only correction of the blanket five-count apply invariant.",
        )
        report = handler.write_output(remote, output_json=OUTPUT_JSON)
    reply = str(((report.get("results") or [{}])[0]).get("reply") or "")
    report["correction_score"] = score_reply(reply)
    report["service_unchanged"] = bool(
        (report.get("service_before_after") or {}).get("unchanged_from_preexisting_live_readback")
    )
    report["scored_at_utc"] = datetime.now(timezone.utc).isoformat()
    report["pass"] = bool(
        report.get("remote_returncode") == 0
        and report.get("pass_runtime") is True
        and report.get("db_counts_changed") is False
        and report["service_unchanged"]
        and report.get("temp_profile_removed") is True
        and report.get("posted_to_telegram") is False
        and report["correction_score"]["pass"]
    )
    OUTPUT_JSON.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    write_markdown(report)
    print(
        json.dumps(
            {
                "json": str(OUTPUT_JSON),
                "markdown": str(OUTPUT_MARKDOWN),
                "pass": report["pass"],
                "rescored_existing": args.score_existing,
                "score": report["correction_score"],
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0 if report["pass"] else 1


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