from __future__ import annotations

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

from policyfaultbench.conformance import evaluate_proposal_conformance
from policyfaultbench.events import sha256_json
from policyfaultbench.scenarios import load_scenarios


def audit(root: Path, event_log: Path) -> dict[str, Any]:
    scenarios = {
        item.scenario_id: item
        for item in load_scenarios(root / "scenarios" / "core_40.json")
    }
    events: list[dict[str, Any]] = []
    hash_failures: list[int] = []
    event_ids: set[str] = set()
    duplicate_event_ids: list[str] = []
    for line_number, line in enumerate(
        event_log.read_text(encoding="utf-8").splitlines(), start=1
    ):
        if not line.strip():
            continue
        event = json.loads(line)
        events.append(event)
        if sha256_json(event.get("payload", {})) != event.get("payload_hash"):
            hash_failures.append(line_number)
        event_id = str(event.get("event_id", ""))
        if event_id in event_ids:
            duplicate_event_ids.append(event_id)
        event_ids.add(event_id)

    by_run: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for event in events:
        by_run[str(event["run_id"])].append(event)

    conformance_failures: list[dict[str, Any]] = []
    oracle_failures: list[dict[str, Any]] = []
    containment_failures: list[dict[str, Any]] = []
    audited_proposals = 0
    complete_runs = 0
    for run_id, run_events in by_run.items():
        started = next(
            (x for x in run_events if x["event_type"] == "run_started"),
            None,
        )
        proposed = next(
            (x for x in run_events if x["event_type"] == "tool_proposed"),
            None,
        )
        finished = next(
            (x for x in run_events if x["event_type"] == "run_finished"),
            None,
        )
        if finished is not None:
            complete_runs += 1
            payload = finished["payload"]
            if not bool(payload.get("oracle_passed", False)):
                row = {
                    "run_id": run_id,
                    "trial_id": payload.get("trial_id", ""),
                    "scenario_id": payload.get("scenario_id", ""),
                    "decision": payload.get("decision", ""),
                    "rule_id": payload.get("rule_id", ""),
                    "state_changed": bool(payload.get("state_changed", False)),
                }
                oracle_failures.append(row)
                if bool(payload.get("containment_failure", False)):
                    containment_failures.append(row)
        if started is None or proposed is None:
            continue
        scenario_id = str(started["payload"]["scenario_id"])
        trial = started["payload"].get("trial", {})
        result = evaluate_proposal_conformance(
            scenarios[scenario_id],
            proposed["payload"]["call"],
        )
        audited_proposals += 1
        if not result.passed:
            conformance_failures.append(
                {
                    "run_id": run_id,
                    "trial_id": trial.get("trial_id", ""),
                    "scenario_id": scenario_id,
                    **result.as_dict(),
                }
            )

    return {
        "schema_version": "1.0",
        "event_log": event_log.name,
        "event_count": len(events),
        "run_count": len(by_run),
        "complete_run_count": complete_runs,
        "audited_proposal_count": audited_proposals,
        "proposal_conformance_passed": (
            audited_proposals - len(conformance_failures)
        ),
        "proposal_conformance_failed": len(conformance_failures),
        "conformance_failures": conformance_failures,
        "oracle_failure_count": len(oracle_failures),
        "oracle_failures": oracle_failures,
        "containment_failure_count": len(containment_failures),
        "containment_failures": containment_failures,
        "payload_hash_failure_lines": hash_failures,
        "duplicate_event_ids": duplicate_event_ids,
        "integrity_ok": not hash_failures and not duplicate_event_ids,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--event-log", type=Path, required=True)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--require-conformant", action="store_true")
    args = parser.parse_args()
    report = audit(args.root, args.event_log)
    rendered = json.dumps(report, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered, encoding="utf-8")
    print(rendered, end="")
    if not report["integrity_ok"]:
        raise SystemExit(2)
    if args.require_conformant and report["proposal_conformance_failed"]:
        raise SystemExit(3)


if __name__ == "__main__":
    main()
