from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path
from typing import Any

from policyfaultbench.capability_policy import CapabilityPolicy
from policyfaultbench.config import load_dotenv_safely
from policyfaultbench.cost import CostGuard
from policyfaultbench.events import EventWriter
from policyfaultbench.generalization_runner import (
    PolicyGeneralizationRunner,
)
from policyfaultbench.generalization_study import (
    freeze_extension_schedule,
    read_finished_trials,
)
from policyfaultbench.policy import PolicySet
from policyfaultbench.providers import (
    AnthropicAdapter,
    MockAdapter,
    OpenAIAdapter,
)
from policyfaultbench.safety import AssetManifest
from policyfaultbench.scenarios import load_scenarios
from policyfaultbench.study import read_ledger


def _adapter(
    provider: str,
    protocol: dict[str, Any],
) -> OpenAIAdapter | AnthropicAdapter | MockAdapter:
    if provider == "mock":
        return MockAdapter()
    env_names = {
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
    }
    api_key = os.environ.get(env_names[provider], "").strip()
    if len(api_key) < 20:
        raise RuntimeError(f"{env_names[provider]} is unavailable")
    model = str(protocol["providers"][provider]["model"])
    timeout = int(protocol["execution"]["request_timeout_seconds"])
    adapters = {
        "openai": OpenAIAdapter,
        "anthropic": AnthropicAdapter,
    }
    return adapters[provider](model, api_key, timeout_seconds=timeout)


def _cost_guard(
    provider: str,
    protocol: dict[str, Any],
) -> CostGuard:
    if provider == "mock":
        return CostGuard(
            provider_limits={"mock": 0.0},
            total_limit=0.0,
            pricing_per_million={"mock": (0.0, 0.0)},
        )
    settings = protocol["providers"][provider]
    prices = settings["pricing_usd_per_million_tokens"]
    cap = float(settings["hard_cap_usd"])
    return CostGuard(
        provider_limits={provider: cap},
        total_limit=cap,
        pricing_per_million={
            provider: (float(prices["input"]), float(prices["output"]))
        },
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument(
        "--provider", choices=("openai", "anthropic", "mock"), required=True
    )
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--max-trials", type=int)
    parser.add_argument("--delay-seconds", type=float)
    args = parser.parse_args()

    root = args.root.resolve()
    load_dotenv_safely(
        [root / ".env.local", root.parent / ".env.local"]
    )
    protocol = json.loads(
        (
            root / "config" / "policy_generalization_protocol_v1.json"
        ).read_text(encoding="utf-8")
    )
    if protocol["status"] != "frozen_before_live_execution":
        raise RuntimeError("Extension protocol is not frozen")
    if args.provider != "mock":
        live_gate = (
            os.getenv("PFB_EXTENSION_LIVE_ENABLED", "").strip().lower()
            == "true"
        )
        if not live_gate:
            raise RuntimeError(
                "PFB_EXTENSION_LIVE_ENABLED is not explicitly true"
            )

    scenarios = load_scenarios(
        root / protocol["design"]["external_corpus"]
    )
    scenario_by_id = {item.scenario_id: item for item in scenarios}
    schedule = freeze_extension_schedule(
        args.output_dir / "policy_generalization_schedule.json",
        scenarios,
        int(protocol["design"]["repetitions_per_scenario"]),
        int(protocol["design"]["randomization_seed"]),
        list(protocol["preflight_gate"]["scenario_ids"]),
    )
    if (
        schedule["schedule_hash"]
        != protocol["design"]["expected_schedule_hash"]
    ):
        raise RuntimeError("Extension schedule does not match the freeze")

    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)
    experiment_id = protocol["protocol_id"]
    event_path = output_dir / f"{experiment_id}-{args.provider}.jsonl"
    adapter = _adapter(args.provider, protocol)
    guard = _cost_guard(args.provider, protocol)
    ledger = read_ledger(event_path)
    if ledger.completed_cost_usd:
        guard.provider_spend[args.provider] = ledger.completed_cost_usd
        guard.total_spend = ledger.completed_cost_usd
    runner = PolicyGeneralizationRunner(
        reference_policy=PolicySet.load(
            root / "policies" / "reference_policy.json"
        ),
        capability_policy=CapabilityPolicy.load(
            root / "policies" / "capability_policy_v1.json"
        ),
        manifest=AssetManifest.load(
            root / "manifests" / "owned_assets.json"
        ),
        adapter=adapter,
        event_writer=EventWriter(event_path),
        cost_guard=guard,
    )

    gate_ids = set(schedule["gate_trial_ids"])
    finished = read_finished_trials(event_path)
    completed_or_terminal = (
        ledger.completed_trial_ids | ledger.terminal_trial_ids
    )
    resolved_gate = gate_ids & completed_or_terminal
    gate_phase = resolved_gate != gate_ids
    if not gate_phase:
        if not all(
            bool(finished[item].get("proposal_conformance", False))
            for item in gate_ids
            if item in finished
        ) or not gate_ids.issubset(finished):
            raise RuntimeError(
                "The embedded provider conformance gate failed"
            )

    provider_settings = (
        {"delay_seconds": 0.0, "max_transient_retries_per_trial": 3}
        if args.provider == "mock"
        else protocol["providers"][args.provider]
    )
    delay = (
        float(args.delay_seconds)
        if args.delay_seconds is not None
        else float(provider_settings["delay_seconds"])
    )
    max_attempts = 1 + int(
        provider_settings["max_transient_retries_per_trial"]
    )
    attempted = 0
    outcomes = []
    stop_reason = ""
    for row in schedule["trials"]:
        trial_id = str(row["trial_id"])
        if gate_phase and trial_id not in gate_ids:
            continue
        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 > 0:
            time.sleep(delay)
        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,
            "protocol_id": protocol["protocol_id"],
            "phase": "embedded_gate" if trial_id in gate_ids else "extension",
        }
        outcome = runner.run(
            scenario_by_id[str(row["scenario_id"])], metadata
        )
        attempted += 1
        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)
    final_finished = read_finished_trials(event_path)
    final_completed = (
        final_ledger.completed_trial_ids | final_ledger.terminal_trial_ids
    )
    resolved_gate = gate_ids & final_completed
    gate_passed = bool(
        resolved_gate == gate_ids
        and gate_ids.issubset(final_finished)
        and all(
            bool(final_finished[item].get("proposal_conformance", False))
            for item in gate_ids
        )
    )
    if gate_phase and resolved_gate == gate_ids and not gate_passed:
        stop_reason = "embedded_gate_failed"
    total = int(schedule["trial_count"])
    status = {
        "schema_version": "1.0",
        "protocol_id": protocol["protocol_id"],
        "provider": args.provider,
        "model": adapter.model,
        "schedule_hash": schedule["schedule_hash"],
        "planned_trials": total,
        "primary_complete": len(final_completed),
        "successful_run_records": len(
            final_ledger.completed_trial_ids
        ),
        "terminal_conformance_records": len(
            final_ledger.terminal_trial_ids
        ),
        "remaining": max(0, total - len(final_completed)),
        "completion_fraction": (
            len(final_completed) / total if total else 0.0
        ),
        "completed_cost_usd": round(
            final_ledger.completed_cost_usd, 8
        ),
        "embedded_gate_resolved": len(resolved_gate),
        "embedded_gate_passed": gate_passed,
        "transient_abort_count": final_ledger.transient_abort_count,
        "harness_abort_count": final_ledger.harness_abort_count,
        "model_conformance_failure_count": (
            final_ledger.model_conformance_failure_count
        ),
        "invocation_attempted": attempted,
        "invocation_strict_successes": sum(
            item.strict_success for item in outcomes
        ),
        "invocation_containment_failures": sum(
            item.containment_failure for item in outcomes
        ),
        "stop_reason": stop_reason,
        "event_log": str(event_path),
    }
    status_path = (
        output_dir
        / f"policy_generalization_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 stop_reason == "embedded_gate_failed":
        return 2
    if gate_phase and gate_passed:
        return 3
    if status["remaining"] == 0:
        return 0
    if stop_reason and not stop_reason.startswith(
        "invocation_trial_limit"
    ):
        return 2
    return 3


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