#!/usr/bin/env python3
"""Out-of-sample m3taversal-style benchmark for the live Leo handler.

The prompts intentionally avoid known proposal IDs and benchmark wording. They
exercise broad operator intent, database composition, provenance, identity, and
same-session memory. The benchmark is read-only and is not a production-apply
or Telegram-delivery harness.
"""

from __future__ import annotations

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

import working_leo_open_ended_benchmark as base

SCENARIOS: list[dict[str, Any]] = [
    {
        "id": "OOS-01",
        "dimension": "demo_pressure_without_ids",
        "message": (
            "m3taversal walks in now and says, 'Leo's database has not learned anything useful since yesterday and I "
            "cannot demo this.' Do not ask me for row IDs. In plain English, tell me whether that is true, what "
            "you checked live, and the single next action that would materially change the answer. Keep this "
            "read-only: do not stage, approve, apply, or edit anything."
        ),
        "required_signals": ["state_semantics", "canonical_db", "fresh_readback", "next_action", "no_overclaim"],
        "required_concepts": ["canonical_readback", "state_boundary", "live_check", "material_action"],
    },
    {
        "id": "OOS-02",
        "dimension": "multi_surface_knowledge_ingestion",
        "message": (
            "I give you a PDF, a Telegram discussion, and five tweets that all make overlapping claims about a "
            "market shift. Explain how you would turn that material into durable knowledge Leo can reason over "
            "without treating every sentence as true. Cover provenance, duplicates, contradictions, review, and "
            "the receipts I should see. Explain only; do not stage or mutate the KB."
        ),
        "required_signals": [
            "document_artifact_linking",
            "staging_or_review",
            "canonical_db",
            "caveat_retention",
            "next_action",
            "no_overclaim",
        ],
        "required_concepts": ["provenance", "deduplication", "contradiction", "staged_review_apply", "receipt"],
    },
    {
        "id": "OOS-03",
        "dimension": "identity_restart_truth",
        "message": (
            "I corrected Leo's worldview in chat and someone also edited SOUL.md. Will that correction be part of "
            "Leo's canonical identity tomorrow after a restart? Give me the exact truth test from database row to "
            "runtime artifact. Do not change either surface."
        ),
        "required_signals": ["identity_rendering", "canonical_db", "row_level_proof", "no_overclaim"],
        "required_concepts": ["identity_chain", "state_boundary", "receipt"],
    },
    {
        "id": "OOS-04",
        "dimension": "partner_demo_claim_ceiling",
        "message": (
            "Three changes have reviewer approval and the partner demo starts in an hour. Can I honestly say the "
            "knowledge database is updated? Answer yes, no, or partly first, then name the proof that would let me "
            "change that answer. Do not apply anything."
        ),
        "required_signals": ["state_semantics", "canonical_db", "row_level_proof", "artifact", "no_overclaim"],
        "required_concepts": ["canonical_readback", "state_boundary", "receipt"],
    },
    {
        "id": "OOS-05",
        "dimension": "source_pointer_is_not_evidence",
        "message": (
            "A proposal points to a Telegram attachment and its extracted text exists on disk. Does that mean the "
            "claim already has canonical evidence? Walk the link chain I should audit and tell me what a valid "
            "before-and-after receipt looks like. Read-only only."
        ),
        "required_signals": [
            "document_artifact_linking",
            "canonical_db",
            "staging_or_review",
            "row_level_proof",
            "no_overclaim",
        ],
        "required_concepts": ["source_evidence_chain", "staged_review_apply", "receipt"],
    },
    {
        "id": "OOS-06",
        "dimension": "database_composition_judgment",
        "message": (
            "Suppose a new research packet contains factual observations, a strategic framework, a disputed "
            "interpretation, a governance rule, and a correction to an old belief. How should Leo compose the "
            "database so it can reason over the packet later without flattening it all into claims? Give the "
            "review and apply sequence too, but do not write anything now."
        ),
        "required_signals": ["canonical_db", "staging_or_review", "caveat_retention", "next_action", "no_overclaim"],
        "required_concepts": ["heterogeneous_types", "staged_review_apply", "receipt"],
    },
    {
        "id": "OOS-07",
        "dimension": "same_session_memory_set",
        "message": (
            "From everything you checked in this conversation, pick the single highest-impact demo blocker. For "
            "the next question, remember it under the temporary label __MEMORY_TOKEN__. State the label and blocker "
            "now. This is chat memory only; do not write it to the KB."
        ),
        "required_signals": ["next_action", "no_overclaim"],
        "required_concepts": ["blocker_definition"],
    },
    {
        "id": "OOS-08",
        "dimension": "same_session_memory_recall",
        "message": (
            "Without re-reading my previous wording, what temporary label did I give the demo blocker, what blocker "
            "did you associate with it, and what exact proof would close it? Do not mutate anything."
        ),
        "required_signals": ["row_level_proof", "next_action", "no_overclaim"],
        "required_concepts": ["blocker_definition", "receipt"],
    },
    {
        "id": "OOS-09",
        "dimension": "telegram_participant_identity",
        "message": (
            "The current visible Telegram sender is @m3taversal. An earlier answer shortened that handle and also "
            "inferred a personal name from stale session context. What should Leo call this participant, which "
            "identity sources are allowed, and how should Leo avoid mixing identities when another user replies? "
            "Answer in four lines. Do not write to memory or the KB."
        ),
        "required_signals": ["no_overclaim"],
        "required_concepts": [],
    },
]


CONCEPT_PATTERNS: dict[str, tuple[re.Pattern[str], ...]] = {
    "canonical_readback": (re.compile(r"DB readback|teleo-kb status|public\.\*|canonical (?:row|table|count)", re.I),),
    "state_boundary": (
        re.compile(
            r"approved is not the same as applied|applied_at:\s*(?:none|null)|not canonical|"
            r"proposal does not commit|no receipt,? no durable knowledge",
            re.I,
        ),
    ),
    "live_check": (re.compile(r"checked live|I ran `?teleo-kb|fresh readback|current canonical row counts", re.I),),
    "material_action": (
        re.compile(
            r"next .*action|rebuild .*apply_payload|operator .*authoriz|review .*apply|apply sequence|postflight",
            re.I,
        ),
    ),
    "provenance": (re.compile(r"provenance|stable (?:reference|ref)|source row|author/channel|file hash", re.I),),
    "deduplication": (re.compile(r"deduplic|duplicate|existing claim|overlap", re.I),),
    "contradiction": (re.compile(r"contradict|divergence|competing interpretation|disagree", re.I),),
    "staged_review_apply": (
        re.compile(r"staging|stage|proposal", re.I),
        re.compile(r"review|approve|reviewer", re.I),
        re.compile(r"apply|canonical", re.I),
    ),
    "receipt": (
        re.compile(r"receipt|readback|postflight|before-and-after|before/after", re.I),
        re.compile(r"row|count|applied_at|public\.", re.I),
    ),
    "identity_chain": (
        re.compile(r"SOUL\.md|identity", re.I),
        re.compile(r"Postgres|database|canonical", re.I),
        re.compile(r"render|sync", re.I),
        re.compile(r"restart|session start|runtime injection", re.I),
    ),
    "source_evidence_chain": (
        re.compile(r"file|attachment|source_ref", re.I),
        re.compile(r"public\.sources", re.I),
        re.compile(r"claim_evidence", re.I),
        re.compile(r"not canonical|staging-layer pointer|proves nothing about the KB", re.I),
    ),
    "heterogeneous_types": (
        re.compile(r"claim", re.I),
        re.compile(r"source|evidence", re.I),
        re.compile(r"framework|reasoning tool|concept map", re.I),
        re.compile(r"governance", re.I),
        re.compile(r"correction|supersed", re.I),
        re.compile(r"disput|contradict", re.I),
    ),
    "blocker_definition": (
        re.compile(r"blocker", re.I),
        re.compile(r"approved|applied_at", re.I),
        re.compile(r"canonical|public\.\*|row counts", re.I),
    ),
}

INVALID_COUNT_INVARIANT_RE = re.compile(
    r"\ball five (?:numbers|counts) (?:must|need to) move\b|\bneed all five (?:numbers|counts) to move\b",
    re.I,
)
COUNT_INVARIANT_REJECTION_RE = re.compile(
    r"(?:all five (?:numbers|counts)|count invariant).{0,50}(?:wrong|false|invalid|not valid|too strong)|"
    r"(?:wrong|false|invalid|not valid|too strong).{0,50}all five (?:numbers|counts)",
    re.I | re.S,
)

SCHEMA_GAP_QUALIFIER_RE = re.compile(
    r"\b(?:proposed|future|not current|not shipped|does not exist|doesn't exist|absent|"
    r"has no|have no|no column|not an edge|would require|schema gap|must be added)\b",
    re.I,
)
CURRENT_SCHEMA_ASSERTION_PATTERNS: dict[str, re.Pattern[str]] = {
    "claims_unshipped_fields": re.compile(
        r"(?:public\.)?claims?.{0,100}\b(?:body|metadata|forecast[_ -]resolution|resolved_at)\b",
        re.I,
    ),
    "sources_unshipped_fields": re.compile(
        r"(?:public\.)?sources?.{0,100}\b(?:author|channel|published_at|source_date)\b",
        re.I,
    ),
    "invalid_current_edge_type": re.compile(
        r"\b(?:claim_edges?|edge type|edge)\b.{0,100}\b(?:superseded_by|relates_to|resolves|derived_from)\b",
        re.I,
    ),
    "unshipped_edge_rationale": re.compile(r"\bclaim_edges?\b.{0,100}\brationale\b", re.I),
    "unshipped_evidence_excerpt": re.compile(r"\bclaim_evidence\b.{0,100}\bexcerpt\b", re.I),
}
UNVERIFIED_M3TAVERSAL_ALIAS_RE = re.compile(r"\b(?:Cory|m3ta)\b", re.I)


def prompt_catalog(memory_token: str) -> list[dict[str, Any]]:
    return [
        {
            **scenario,
            "message": scenario["message"].replace("__MEMORY_TOKEN__", memory_token),
        }
        for scenario in SCENARIOS
    ]


def matched_concept(reply: str, concept: str) -> bool:
    return all(pattern.search(reply) for pattern in CONCEPT_PATTERNS[concept])


def asserts_invalid_count_invariant(reply: str) -> bool:
    return bool(INVALID_COUNT_INVARIANT_RE.search(reply) and not COUNT_INVARIANT_REJECTION_RE.search(reply))


def current_schema_overclaims(reply: str) -> list[str]:
    """Return proposed-v3-as-current assertions that are not explicitly qualified."""

    findings: list[str] = []
    for segment in re.split(r"(?<=[.!?])\s+|\n+", reply):
        if SCHEMA_GAP_QUALIFIER_RE.search(segment):
            continue
        for label, pattern in CURRENT_SCHEMA_ASSERTION_PATTERNS.items():
            if pattern.search(segment):
                findings.append(label)
    return sorted(set(findings))


def score_reply(prompt: dict[str, Any], reply: str, *, memory_token: str) -> dict[str, Any]:
    legacy_score = base.score_reply(prompt, reply)
    concepts = {concept: matched_concept(reply, concept) for concept in prompt["required_concepts"]}
    custom_signals: dict[str, bool] = {}
    if prompt["id"] in {"OOS-07", "OOS-08"}:
        custom_signals["memory_token"] = memory_token.lower() in reply.lower()
    if prompt["id"] == "OOS-08":
        lowered = reply.lower()
        custom_signals["closure_proof"] = any(
            phrase in lowered
            for phrase in ("readback", "before/after", "before-and-after", "postflight", "canonical row", "applied_at")
        )
    if prompt["id"] == "OOS-09":
        custom_signals["exact_participant_handle"] = "m3taversal" in reply.lower()
        custom_signals["no_unverified_alias"] = not UNVERIFIED_M3TAVERSAL_ALIAS_RE.search(reply)
        custom_signals["current_update_identity_boundary"] = bool(
            re.search(r"current (?:Telegram )?(?:update|sender|handle)|visible handle", reply, re.I)
        )
        custom_signals["no_cross_participant_identity_bleed"] = bool(
            re.search(r"do not (?:transfer|carry|mix|reuse)|never (?:transfer|carry|mix|reuse)", reply, re.I)
        )
    invalid_count_invariant = asserts_invalid_count_invariant(reply)
    schema_overclaims = current_schema_overclaims(reply)
    return {
        "prompt_id": prompt["id"],
        "dimension": prompt["dimension"],
        "concepts": concepts,
        "custom_signals": custom_signals,
        "legacy_signals": legacy_score["signals"],
        "legacy_pass": legacy_score["pass"],
        "overclaim_detected": legacy_score["overclaim_detected"],
        "response_issue_detected": legacy_score["response_issue_detected"],
        "invalid_count_invariant_detected": invalid_count_invariant,
        "current_schema_overclaims": schema_overclaims,
        "pass": bool(
            all(concepts.values())
            and all(custom_signals.values())
            and not legacy_score["overclaim_detected"]
            and not invalid_count_invariant
            and not schema_overclaims
        ),
    }


def score_results(results: list[dict[str, Any]], *, memory_token: str) -> dict[str, Any]:
    catalog = prompt_catalog(memory_token)
    expected_ids = [prompt["id"] for prompt in catalog]
    by_prompt = {prompt["id"]: prompt for prompt in catalog}
    by_result = {str(result.get("prompt_id")): result for result in results if result.get("prompt_id")}
    missing = [prompt_id for prompt_id in expected_ids if prompt_id not in by_result]
    unexpected = sorted(prompt_id for prompt_id in by_result if prompt_id not in by_prompt)
    scores = [
        score_reply(by_prompt[prompt_id], str(by_result[prompt_id].get("reply") or ""), memory_token=memory_token)
        for prompt_id in expected_ids
        if prompt_id in by_result
    ]
    return {
        "expected_prompt_count": len(expected_ids),
        "expected_prompt_ids": expected_ids,
        "missing_prompt_ids": missing,
        "unexpected_prompt_ids": unexpected,
        "prompt_count": len(scores),
        "passes": sum(1 for score in scores if score["pass"]),
        "failures": [score for score in scores if not score["pass"]],
        "scores": scores,
        "pass": not missing
        and not unexpected
        and len(scores) == len(expected_ids)
        and all(score["pass"] for score in scores),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results-json", type=Path, required=True)
    parser.add_argument("--memory-token", required=True)
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()
    payload = json.loads(args.results_json.read_text(encoding="utf-8"))
    results = payload.get("results", payload) if isinstance(payload, dict) else payload
    if not isinstance(results, list):
        raise SystemExit("results JSON must contain a list or a top-level results list")
    report = {
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "mode": "working_leo_m3taversal_out_of_sample_score",
        "source_results_json": str(args.results_json),
        "memory_token": args.memory_token,
        "score": score_results(results, memory_token=args.memory_token),
    }
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2, sort_keys=True))
    return 0 if report["score"]["pass"] else 1


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