from __future__ import annotations

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

from policyfaultbench.conformance import evaluate_proposal_conformance
from policyfaultbench.models import ToolCall
from policyfaultbench.providers.anthropic import AnthropicAdapter
from policyfaultbench.providers.gemini import GeminiAdapter
from policyfaultbench.scenarios import load_scenarios


GATE_SCENARIOS = (
    "S01-E",
    "S03-E",
    "S04-D",
    "S05-E",
    "S06-B",
    "S06-D",
    "S07-B",
    "S07-D",
    "S07-I",
    "S09-B",
    "S09-D",
    "S09-I",
)


def is_daily_quota_error(exc: Exception) -> bool:
    message = str(exc).lower()
    return (
        type(exc).__name__ == "ProviderError"
        and "http 429" in message
        and any(
            marker in message
            for marker in (
                "generate_content_free_tier_requests",
                "requests per day",
                "daily quota",
                "rpd",
            )
        )
    )


def is_transient_provider_error(exc: Exception) -> bool:
    message = str(exc).lower()
    return (
        type(exc).__name__ == "ProviderError"
        and any(
            marker in message
            for marker in (
                "http 429",
                "http 500",
                "http 502",
                "http 503",
                "http 504",
                "timeout",
                "connection",
            )
        )
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "provider",
        choices=("anthropic", "gemini"),
    )
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--delay-seconds", type=float, default=0.0)
    parser.add_argument("--max-attempts", type=int, default=4)
    parser.add_argument("--retry-wait-seconds", type=float, default=90.0)
    args = parser.parse_args()
    if args.max_attempts <= 0:
        raise SystemExit("--max-attempts must be positive")

    env_names = {
        "anthropic": "ANTHROPIC_API_KEY",
        "gemini": "GEMINI_API_KEY",
    }
    api_key = os.environ.get(env_names[args.provider], "")
    if len(api_key) < 20:
        raise SystemExit(f"{env_names[args.provider]} is unavailable")
    config = json.loads(
        (args.root / "config" / "experiment.json").read_text(
            encoding="utf-8"
        )
    )
    model = str(config["models"][args.provider])
    adapter_types = {
        "anthropic": AnthropicAdapter,
        "gemini": GeminiAdapter,
    }
    adapter = adapter_types[args.provider](
        model,
        api_key,
        timeout_seconds=int(config["request_timeout_seconds"]),
    )
    scenarios = {
        item.scenario_id: item
        for item in load_scenarios(
            args.root / "scenarios" / "core_40.json"
        )
    }

    prior_rows: dict[str, dict[str, Any]] = {}
    if args.output.is_file():
        try:
            prior = json.loads(args.output.read_text(encoding="utf-8"))
            prior_rows = {
                str(row["scenario_id"]): dict(row)
                for row in prior.get("results", [])
                if row.get("scenario_id")
            }
        except (OSError, json.JSONDecodeError):
            prior_rows = {}

    results: list[dict[str, Any]] = []
    invocation_attempts_this_run = 0
    daily_quota_exhausted = False
    for index, scenario_id in enumerate(GATE_SCENARIOS):
        prior_row = prior_rows.get(scenario_id, {})
        if prior_row.get("passed") is True:
            results.append(prior_row)
            continue
        if index and args.delay_seconds:
            time.sleep(args.delay_seconds)
        scenario = scenarios[scenario_id]
        row: dict[str, Any] = {
            "scenario_id": scenario_id,
            "prior_failed_gate_record": bool(prior_row),
        }
        for attempt in range(1, args.max_attempts + 1):
            row["attempts_in_resume"] = attempt
            invocation_attempts_this_run += 1
            try:
                response = adapter.propose(scenario)
                if len(response.proposals) != 1:
                    raise ValueError("exactly one proposal is required")
                proposal = response.proposals[0]
                if proposal.name != "propose_action":
                    raise ValueError("unexpected proposal name")
                raw = proposal.arguments
                call = ToolCall(
                    subject=str(raw["subject"]),
                    action=str(raw["action"]),
                    resource=str(raw["resource"]),
                    arguments=dict(raw.get("arguments", {})),
                    context=dict(raw.get("context", {})),
                    call_id=proposal.call_id,
                )
                result = evaluate_proposal_conformance(scenario, call)
                row.update(result.as_dict())
                row["input_tokens"] = response.input_tokens
                row["output_tokens"] = response.output_tokens
                row.pop("error_type", None)
                row.pop("error_message", None)
                row.pop("quota_daily_exhausted", None)
                break
            except Exception as exc:
                message = str(exc).replace("\n", " ")[:400]
                quota_daily_exhausted = is_daily_quota_error(exc)
                row.update(
                    {
                        "passed": False,
                        "mismatch_fields": ["provider_or_schema_error"],
                        "error_type": type(exc).__name__,
                        "error_message": message,
                        "quota_daily_exhausted": quota_daily_exhausted,
                    }
                )
                if quota_daily_exhausted:
                    daily_quota_exhausted = True
                    break
                transient = is_transient_provider_error(exc)
                if not transient or attempt == args.max_attempts:
                    break
                time.sleep(args.retry_wait_seconds)
        results.append(row)

    failures = [row for row in results if not row.get("passed", False)]
    report = {
        "schema_version": "1.0",
        "gate_id": (
            f"{args.provider}-schema-conformance-v1p2"
        ),
        "provider": args.provider,
        "model": model,
        "scenario_count": len(results),
        "passed": len(results) - len(failures),
        "failed": len(failures),
        "gate_passed": not failures,
        "invocation_attempts_this_run": invocation_attempts_this_run,
        "daily_quota_exhausted": daily_quota_exhausted,
        "results": results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(report, indent=2, sort_keys=True))
    if failures:
        raise SystemExit(2)


if __name__ == "__main__":
    main()
