from __future__ import annotations

import argparse
import json
import time
from dataclasses import asdict
from pathlib import Path

from .config import ExperimentConfig
from .events import EventWriter
from .mutation_eval import write_mutation_report
from .policy import PolicySet
from .providers import AnthropicAdapter, GeminiAdapter, MockAdapter, OpenAIAdapter
from .runner import ExperimentRunner, build_cost_guard
from .sandbox import SyntheticSandbox
from .scenarios import load_scenarios
from .safety import AssetManifest
from .validation import validate_project
from .study import freeze_schedule, read_ledger, study_status


def project_root() -> Path:
    return Path(__file__).resolve().parents[2]


def _components(
    root: Path,
    provider: str,
    event_path: Path | None = None,
    study: bool = False,
):
    config = ExperimentConfig.load(root)
    manifest = AssetManifest.load(root / "manifests" / "owned_assets.json")
    policy = PolicySet.load(root / "policies" / "reference_policy.json")
    sandbox = SyntheticSandbox(manifest)
    if provider == "mock":
        adapter = MockAdapter()
    else:
        if study:
            config.assert_study_ready(provider)
        else:
            config.assert_live_ready(provider)
        kwargs = {
            "model": config.model(provider),
            "api_key": config.api_key(provider),
            "timeout_seconds": int(config.raw["request_timeout_seconds"]),
        }
        adapters = {
            "openai": OpenAIAdapter,
            "anthropic": AnthropicAdapter,
            "gemini": GeminiAdapter,
        }
        adapter = adapters[provider](**kwargs)
    writer = EventWriter(
        event_path
        or (root / "runs" / f"{config.experiment_id}-{provider}.jsonl")
    )
    runner = ExperimentRunner(
        policy=policy,
        sandbox=sandbox,
        adapter=adapter,
        event_writer=writer,
        cost_guard=build_cost_guard(config.raw),
    )
    return config, runner


def command_validate(args: argparse.Namespace) -> int:
    report = validate_project(args.root)
    print(json.dumps(asdict(report), indent=2, sort_keys=True))
    return 0


def command_pilot(args: argparse.Namespace) -> int:
    _, runner = _components(args.root, args.provider)
    scenarios = load_scenarios(args.root / "scenarios" / "core_40.json")
    if args.live_core:
        scenarios = [item for item in scenarios if item.live_core]
    offset = int(getattr(args, "offset", 0))
    scenarios = scenarios[offset : offset + args.limit]
    outcomes = []
    for index, scenario in enumerate(scenarios):
        if index and float(getattr(args, "delay_seconds", 0.0)) > 0:
            time.sleep(float(args.delay_seconds))
        outcome = runner.run(scenario)
        outcomes.append(outcome)
        if getattr(args, "stop_on_failure", False) and (
            outcome.aborted or outcome.containment_failure
        ):
            break
    summary = {
        "provider": args.provider,
        "runs": len(outcomes),
        "oracle_passed": sum(item.oracle_passed for item in outcomes),
        "aborted": sum(item.aborted for item in outcomes),
        "containment_failures": sum(item.containment_failure for item in outcomes),
        "cost_usd": round(sum(item.cost_usd for item in outcomes), 8),
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0 if summary["aborted"] == 0 else 2


def command_mutation_eval(args: argparse.Namespace) -> int:
    output_dir = args.output_dir or (args.root / "results")
    report = write_mutation_report(args.root, output_dir)
    summary = {
        "reference_cases": report["reference_cases"],
        "reference_failures": len(report["reference_failures"]),
        "operator_count": report["operator_count"],
        "core_killed": report["core_killed"],
        "expanded_killed": report["expanded_killed"],
        "core_mutation_score": report["core_mutation_score"],
        "expanded_mutation_score": report["expanded_mutation_score"],
        "output_dir": str(output_dir),
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0 if not report["reference_failures"] else 2


def _study_settings(config: ExperimentConfig) -> dict:
    return dict(config.raw["definitive_study"])


def command_study_plan(args: argparse.Namespace) -> int:
    config = ExperimentConfig.load(args.root)
    settings = _study_settings(config)
    repetitions = int(args.repetitions or settings["repetitions"])
    seed = int(args.seed if args.seed is not None else settings["seed"])
    scenarios = load_scenarios(args.root / "scenarios" / "core_40.json")
    output_dir = args.output_dir or (args.root / "study")
    schedule = freeze_schedule(
        output_dir / "study_schedule.json",
        scenarios,
        repetitions,
        seed,
    )
    print(
        json.dumps(
            {
                "schedule_hash": schedule["schedule_hash"],
                "seed": schedule["seed"],
                "repetitions": schedule["repetitions"],
                "scenario_count": schedule["scenario_count"],
                "trial_count": schedule["trial_count"],
                "output": str(output_dir / "study_schedule.json"),
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


def command_study_run(args: argparse.Namespace) -> int:
    config = ExperimentConfig.load(args.root)
    settings = _study_settings(config)
    repetitions = int(args.repetitions or settings["repetitions"])
    seed = int(args.seed if args.seed is not None else settings["seed"])
    output_dir = args.output_dir or (args.root / "study")
    output_dir.mkdir(parents=True, exist_ok=True)
    scenarios = load_scenarios(args.root / "scenarios" / "core_40.json")
    scenario_by_id = {item.scenario_id: item for item in scenarios}
    schedule = freeze_schedule(
        output_dir / "study_schedule.json",
        scenarios,
        repetitions,
        seed,
    )
    experiment_id = str(settings["experiment_id"])
    event_path = output_dir / f"{experiment_id}-{args.provider}.jsonl"
    ledger = read_ledger(event_path)
    config, runner = _components(
        args.root,
        args.provider,
        event_path=event_path,
        study=args.provider != "mock",
    )
    if ledger.completed_cost_usd:
        runner.cost_guard.provider_spend[args.provider] = ledger.completed_cost_usd
        runner.cost_guard.total_spend = ledger.completed_cost_usd

    delay_seconds = (
        float(args.delay_seconds)
        if args.delay_seconds is not None
        else float(settings["provider_delay_seconds"].get(args.provider, 0.0))
    )
    max_attempts = 1 + int(settings["max_transient_retries_per_trial"])
    completed_or_terminal = (
        ledger.completed_trial_ids | ledger.terminal_trial_ids
    )
    attempted = 0
    invocation_outcomes = []
    stop_reason = ""
    for row in schedule["trials"]:
        trial_id = str(row["trial_id"])
        if trial_id in completed_or_terminal:
            continue
        prior_attempts = ledger.abort_attempts.get(trial_id, 0)
        if prior_attempts >= max_attempts:
            stop_reason = f"transient_retry_exhausted:{trial_id}"
            break
        if args.max_trials is not None and attempted >= args.max_trials:
            stop_reason = "invocation_trial_limit"
            break
        if attempted and delay_seconds > 0:
            time.sleep(delay_seconds)
        metadata = {
            "trial_id": trial_id,
            "repetition": int(row["repetition"]),
            "within_repetition_position": int(
                row["within_repetition_position"]
            ),
            "schedule_position": int(row["schedule_position"]),
            "schedule_hash": schedule["schedule_hash"],
            "family": row["family"],
            "variant": row["variant"],
            "expected_decision": row["expected_decision"],
            "attempt": prior_attempts + 1,
            "study_protocol_version": settings["protocol_version"],
        }
        outcome = runner.run(scenario_by_id[row["scenario_id"]], metadata)
        attempted += 1
        invocation_outcomes.append(outcome)
        if outcome.containment_failure:
            stop_reason = f"containment_failure:{trial_id}"
            break
        if outcome.aborted and not outcome.terminal:
            stop_reason = f"transient_abort:{trial_id}"
            break
        if outcome.aborted and outcome.abort_class != "model_conformance":
            stop_reason = f"{outcome.abort_class}:{trial_id}"
            break

    final_ledger = read_ledger(event_path)
    status = study_status(
        args.provider,
        runner.adapter.model,
        schedule,
        final_ledger,
    )
    status.update(
        {
            "invocation_attempted": attempted,
            "invocation_aborted": sum(x.aborted for x in invocation_outcomes),
            "invocation_oracle_passed": sum(
                x.oracle_passed for x in invocation_outcomes
            ),
            "invocation_containment_failures": sum(
                x.containment_failure for x in invocation_outcomes
            ),
            "stop_reason": stop_reason,
            "event_log": str(event_path),
        }
    )
    status_path = output_dir / f"study_status_{args.provider}.json"
    status_path.write_text(
        json.dumps(status, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(status, indent=2, sort_keys=True))
    if status["remaining"] == 0:
        return 0
    if status["invocation_containment_failures"] or (
        stop_reason and not stop_reason.startswith("invocation_trial_limit")
    ):
        return 2
    return 3


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="policyfaultbench")
    parser.add_argument("--root", type=Path, default=project_root())
    sub = parser.add_subparsers(dest="command", required=True)
    validate = sub.add_parser("validate")
    validate.set_defaults(func=command_validate)
    offline = sub.add_parser("offline-pilot")
    offline.add_argument("--limit", type=int, default=12)
    offline.add_argument("--live-core", action="store_true")
    offline.set_defaults(func=command_pilot, provider="mock")
    mutation = sub.add_parser("mutation-eval")
    mutation.add_argument("--output-dir", type=Path)
    mutation.set_defaults(func=command_mutation_eval)
    plan = sub.add_parser("study-plan")
    plan.add_argument("--repetitions", type=int)
    plan.add_argument("--seed", type=int)
    plan.add_argument("--output-dir", type=Path)
    plan.set_defaults(func=command_study_plan)
    study = sub.add_parser("study-run")
    study.add_argument(
        "provider", choices=["openai", "anthropic", "gemini", "mock"]
    )
    study.add_argument("--repetitions", type=int)
    study.add_argument("--seed", type=int)
    study.add_argument("--output-dir", type=Path)
    study.add_argument("--delay-seconds", type=float)
    study.add_argument("--max-trials", type=int)
    study.set_defaults(func=command_study_run)
    live = sub.add_parser("live-debug")
    live.add_argument("provider", choices=["openai", "anthropic", "gemini"])
    live.add_argument("--limit", type=int, default=1)
    live.add_argument("--offset", type=int, default=0)
    live.add_argument("--delay-seconds", type=float, default=0.0)
    live.add_argument("--live-core", action="store_true")
    live.add_argument(
        "--stop-on-failure",
        action="store_true",
        help="Stop immediately after an aborted run or containment failure.",
    )
    live.set_defaults(func=command_pilot)
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    return int(args.func(args))


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