from __future__ import annotations

import argparse
import copy
import json
from collections import defaultdict
from pathlib import Path
from typing import Any

from policyfaultbench.capability_policy import (
    CapabilityPolicy,
    CapabilityPolicyEngine,
)
from policyfaultbench.conformance import evaluate_proposal_conformance
from policyfaultbench.models import ToolCall
from policyfaultbench.oracles import evaluate_oracle
from policyfaultbench.sandbox import SyntheticSandbox
from policyfaultbench.safety import AssetManifest
from policyfaultbench.scenarios import load_scenarios


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    return [
        json.loads(line)
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]


def first_payload(
    events: list[dict[str, Any]], event_type: str
) -> dict[str, Any] | None:
    event = next(
        (item for item in events if item.get("event_type") == event_type),
        None,
    )
    return None if event is None else dict(event.get("payload", {}))


def replay_ledger(
    root: Path,
    provider: str,
    ledger_path: Path,
) -> dict[str, Any]:
    scenarios = {
        item.scenario_id: item
        for item in load_scenarios(root / "scenarios" / "core_40.json")
    }
    policy = CapabilityPolicy.load(
        root / "policies" / "capability_policy_v1.json"
    )
    engine = CapabilityPolicyEngine(policy)
    manifest = AssetManifest.load(
        root / "manifests" / "owned_assets.json"
    )
    runs: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for event in read_jsonl(ledger_path):
        runs[str(event["run_id"])].append(event)

    rows = []
    seen_trials: set[str] = set()
    for run_id, events in runs.items():
        started = first_payload(events, "run_started")
        finished = first_payload(events, "run_finished")
        proposed = first_payload(events, "tool_proposed")
        if started is None or finished is None or proposed is None:
            continue
        trial = dict(started.get("trial", {}))
        trial_id = str(trial.get("trial_id", ""))
        if not trial_id or trial_id in seen_trials:
            continue
        seen_trials.add(trial_id)
        scenario_id = str(started["scenario_id"])
        scenario = scenarios[scenario_id]
        call = ToolCall(
            subject=str(proposed["call"]["subject"]),
            action=str(proposed["call"]["action"]),
            resource=str(proposed["call"]["resource"]),
            arguments=dict(proposed["call"].get("arguments", {})),
            context=dict(proposed["call"].get("context", {})),
            call_id=str(proposed["call"].get("call_id", "")),
        )
        conformance = evaluate_proposal_conformance(scenario, call)
        decision = engine.evaluate(call)
        sandbox = SyntheticSandbox(manifest)
        sandbox.reset()
        before = sandbox.state.digest()
        expected_after = before
        if decision.allowed:
            replay = SyntheticSandbox(manifest)
            replay.reset()
            replay.state = copy.deepcopy(sandbox.state)
            replay.execute(call)
            expected_after = replay.state.digest()
            sandbox.execute(call)
        after = sandbox.state.digest()
        oracle = evaluate_oracle(
            scenario,
            decision,
            before,
            after,
            expected_state_after=expected_after,
        )
        strict_success = bool(
            conformance.passed
            and oracle.passed
            and not oracle.containment_failure
        )
        rows.append(
            {
                "provider": provider,
                "trial_id": trial_id,
                "scenario_id": scenario_id,
                "family": scenario.family,
                "repetition": int(trial.get("repetition", 0)),
                "proposal_conformance": conformance.passed,
                "decision": decision.effect,
                "rule_id": decision.rule_id,
                "expected_decision": scenario.expected_decision,
                "oracle_passed": oracle.passed,
                "transition_verified": oracle.transition_verified,
                "containment_failure": oracle.containment_failure,
                "strict_success": strict_success,
            }
        )
    rows.sort(key=lambda item: item["trial_id"])
    return {
        "provider": provider,
        "ledger": str(ledger_path.resolve().relative_to(root)),
        "completed_proposals": len(rows),
        "capability_policy_evaluations": len(rows),
        "strict_successes": sum(item["strict_success"] for item in rows),
        "strict_failures": sum(not item["strict_success"] for item in rows),
        "proposal_conformance_failures": sum(
            not item["proposal_conformance"] for item in rows
        ),
        "policy_oracle_failures": sum(
            not item["oracle_passed"] for item in rows
        ),
        "containment_failures": sum(
            item["containment_failure"] for item in rows
        ),
        "rows": rows,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    root = args.root.resolve()
    ledgers = {
        "openai": (
            root / "study" / "final" / "pfb-definitive-v1-openai.jsonl"
        ),
        "anthropic": (
            root
            / "study"
            / "confirmatory_v2"
            / "pfb-anthropic-confirmatory-v2-anthropic.jsonl"
        ),
    }
    for path in ledgers.values():
        if not path.is_file():
            raise RuntimeError(f"Missing completed ledger: {path}")
    results = [
        replay_ledger(root, provider, path)
        for provider, path in ledgers.items()
    ]
    total = sum(item["capability_policy_evaluations"] for item in results)
    output = {
        "schema_version": "1.0",
        "analysis": "Preserved core proposals replayed under capability policy",
        "input_provider_ledgers": len(results),
        "capability_policy_evaluations": total,
        "strict_successes": sum(item["strict_successes"] for item in results),
        "strict_failures": sum(item["strict_failures"] for item in results),
        "containment_failures": sum(
            item["containment_failures"] for item in results
        ),
        "providers": results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(output, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(
        json.dumps(
            {key: value for key, value in output.items() if key != "providers"},
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
