from __future__ import annotations

import argparse
import hashlib
import json
import tempfile
import zipfile
from pathlib import Path
from typing import Any


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 safe_extract(archive_path: Path, destination: Path) -> None:
    destination = destination.resolve()
    with zipfile.ZipFile(archive_path) as archive:
        for item in archive.infolist():
            target = (destination / item.filename).resolve()
            if destination not in target.parents and target != destination:
                raise ValueError(f"Unsafe archive path: {item.filename}")
        archive.extractall(destination)


def successful_run_ids(events: list[dict[str, Any]]) -> list[str]:
    return [
        event["run_id"]
        for event in events
        if event.get("event_type") == "run_finished"
        and event.get("payload", {}).get("oracle_passed") is True
        and event.get("payload", {}).get("containment_failure") is False
        and event.get("payload", {}).get("aborted") is False
    ]


def select_runs(
    events: list[dict[str, Any]], run_ids: list[str]
) -> list[dict[str, Any]]:
    selected = set(run_ids)
    return [event for event in events if event.get("run_id") in selected]


def summarize(events: list[dict[str, Any]]) -> dict[str, Any]:
    finished = [
        event["payload"]
        for event in events
        if event.get("event_type") == "run_finished"
    ]
    return {
        "runs": len(finished),
        "oracle_passed": sum(bool(item["oracle_passed"]) for item in finished),
        "aborted": sum(bool(item.get("aborted", False)) for item in finished),
        "containment_failures": sum(
            bool(item["containment_failure"]) for item in finished
        ),
        "cost_usd": round(sum(float(item["cost_usd"]) for item in finished), 8),
        "models": sorted({str(item["model"]) for item in finished}),
        "scenario_ids": [str(item["scenario_id"]) for item in finished],
    }


def write_jsonl(path: Path, events: list[dict[str, Any]]) -> None:
    path.write_text(
        "".join(json.dumps(item, sort_keys=True) + "\n" for item in events),
        encoding="utf-8",
    )


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--pilot-zip", type=Path, required=True)
    parser.add_argument("--gemini-resume-zip", type=Path, required=True)
    parser.add_argument("--openai-log", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="pfb-evidence-") as temp:
        temp_root = Path(temp)
        pilot_root = temp_root / "pilot"
        resume_root = temp_root / "resume"
        pilot_root.mkdir()
        resume_root.mkdir()
        safe_extract(args.pilot_zip, pilot_root)
        safe_extract(args.gemini_resume_zip, resume_root)

        anthropic_all = read_jsonl(pilot_root / "pfb-prepilot-v1-anthropic.jsonl")
        gemini_initial = read_jsonl(pilot_root / "pfb-prepilot-v1-gemini.jsonl")
        gemini_resume = read_jsonl(resume_root / "pfb-prepilot-v1-gemini.jsonl")
        openai_all = read_jsonl(args.openai_log)

        anthropic = select_runs(
            anthropic_all, successful_run_ids(anthropic_all)[-12:]
        )
        gemini = select_runs(
            gemini_initial, successful_run_ids(gemini_initial)
        ) + select_runs(gemini_resume, successful_run_ids(gemini_resume))
        openai = select_runs(openai_all, successful_run_ids(openai_all)[-12:])

        provider_events = {
            "anthropic": anthropic,
            "gemini": gemini,
            "openai": openai,
        }
        output_files: dict[str, Path] = {}
        for provider, events in provider_events.items():
            output_path = args.output_dir / f"pilot_primary_{provider}.jsonl"
            write_jsonl(output_path, events)
            output_files[provider] = output_path

        providers = {
            provider: summarize(events)
            for provider, events in provider_events.items()
        }
        summary = {
            "schema_version": "1.0",
            "study_stage": "pre-registered pilot",
            "primary_scenario_count_per_provider": 12,
            "providers": providers,
            "aggregate": {
                "runs": sum(item["runs"] for item in providers.values()),
                "oracle_passed": sum(
                    item["oracle_passed"] for item in providers.values()
                ),
                "aborted": sum(item["aborted"] for item in providers.values()),
                "containment_failures": sum(
                    item["containment_failures"] for item in providers.values()
                ),
                "cost_usd": round(
                    sum(item["cost_usd"] for item in providers.values()), 8
                ),
            },
            "excluded_harness_events": {
                "gemini_initial_rate_limit_aborts": sum(
                    event.get("event_type") == "run_aborted"
                    for event in gemini_initial
                ),
                "openai_schema_debug_aborts": sum(
                    event.get("event_type") == "run_aborted"
                    for event in openai_all
                ),
                "openai_successful_debug_runs": max(
                    0, len(successful_run_ids(openai_all)) - 12
                ),
                "exclusion_rule": (
                    "Harness configuration and rate-control diagnostics are "
                    "reported separately and excluded from primary scenario outcomes."
                ),
            },
            "evidence_sha256": {
                provider: sha256(path) for provider, path in output_files.items()
            },
        }
        (args.output_dir / "pilot_combined_summary.json").write_text(
            json.dumps(summary, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


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