"""Preflight validation — run before starting any agent experiment."""

from __future__ import annotations

import inspect
import math
import subprocess
from typing import Any, Callable

import httpx

from ._runtime import resolve_public
from .health import InfraCheck
from .provision import (
    OR_BASE,
    get_account_balance as provision_get_account_balance,
    get_available_models as provision_get_available_models,
    provision_key as provision_key_impl,
    get_usage as provision_get_usage,
)


class PreflightError(RuntimeError):
    """Raised when one or more preflight checks fail."""

    def __init__(self, failures: list[str]) -> None:
        self.failures = failures
        super().__init__(f"Preflight failed ({len(failures)} issue(s)): {failures[0]}")


CACHE_SMOKE_MIN_LIMIT_USD = 0.02
CACHE_SMOKE_FALLBACK_PROMPT_PRICE_PER_M = 20.0
CACHE_SMOKE_OUTPUT_HEADROOM_USD = 0.01
get_account_balance = provision_get_account_balance
get_available_models = provision_get_available_models
provision_key = provision_key_impl
get_usage = provision_get_usage


def _cache_smoke_limit_usd(prompt_tokens: int, prompt_price_per_token: float | None) -> float:
    """Estimate a safe spend cap for the two-request cache smoke test."""
    if prompt_price_per_token is None or prompt_price_per_token <= 0:
        prompt_price_per_m = CACHE_SMOKE_FALLBACK_PROMPT_PRICE_PER_M
    else:
        prompt_price_per_m = prompt_price_per_token * 1_000_000

    # Budget for two full prompt evaluations so the test still completes even
    # if the second request misses cache. Add a little headroom for completions.
    prompt_cost = (prompt_tokens * 2 / 1_000_000) * prompt_price_per_m
    limit = max(
        CACHE_SMOKE_MIN_LIMIT_USD,
        (prompt_cost * 1.25) + CACHE_SMOKE_OUTPUT_HEADROOM_USD,
    )
    return math.ceil(limit * 100) / 100


def _check(name: str, ok: bool | None, detail: str = "", hint: str = "") -> bool:
    symbol = "v" if ok is True else ("!" if ok is None else "x")
    detail_str = f"  {detail}" if detail else ""
    hint_str = f" -> {hint}" if hint and ok is not True else ""
    print(f"  [{symbol}] {name}{detail_str}{hint_str}")
    return ok is True


def _accepts_context(fn: Callable[..., Any]) -> bool:
    """Return True if a custom check can accept a context dict argument."""
    try:
        sig = inspect.signature(fn)
    except (TypeError, ValueError):
        return False

    for p in sig.parameters.values():
        if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD):
            return True
        if p.kind is inspect.Parameter.VAR_POSITIONAL:
            return True
    return False


def _run_custom_check(
    fn: Callable[..., tuple[str, bool | None, str]],
    ctx: dict[str, Any],
) -> tuple[str, bool | None, str]:
    """Run a custom check, passing context when the callable accepts it."""
    if _accepts_context(fn):
        return fn(ctx)
    return fn()


def preflight(
    admin_key: str,
    or_key: str,
    models: list[str],
    budget_per_model: float,
    checks: list[str] | None = None,
    docker_image: str | None = None,
    custom_checks: list[Callable[..., tuple[str, bool | None, str]]] | None = None,
    required_params: dict[str, list[str]] | None = None,
    common_required_params: list[str] | None = None,
    site_name: str = "agent-preflight",
    raise_on_failure: bool = True,
    infra: list[InfraCheck] | None = None,
) -> list[str]:
    """Run preflight validation before launching a multi-model experiment.

    Checks (all enabled by default unless ``checks`` is specified):

    - ``"balance"``       — OpenRouter account balance > $0.10
    - ``"math"``          — balance covers ``len(models) * budget_per_model``
    - ``"models"``        — each model exists on OpenRouter (with pricing info)
    - ``"tool_support"``  — models in ``required_params`` / ``common_required_params``
                            support those request params
    - ``"key_smoke"``     — provision a $0.01 test key and make a real API call
    - ``"cache_smoke"``   — verify Anthropic prompt caching via OpenRouter
    - ``"docker_image"``  — Docker image exists (only if ``docker_image`` is set)

    You can also pass ``custom_checks``: a list of callables that return
    ``(name, ok, detail)`` tuples, same as the built-in checks. Checks may
    accept either zero arguments or a single context dict argument.

    Args:
        admin_key: OpenRouter admin/management key for provisioning sub-keys.
        or_key: Regular OpenRouter API key for balance/model checks.
        models: List of model IDs (e.g. ``["meta-llama/llama-3.3-70b-instruct"]``).
        budget_per_model: USD budget per model (used for math check).
        checks: Subset of checks to run. Defaults to all applicable checks.
        docker_image: Docker image name to verify exists (enables ``"docker_image"`` check).
        custom_checks: Additional check functions returning ``(name, ok, detail)``.
            A check may optionally accept a single context dict with keys like
            ``available_models``, ``models``, ``or_key``, and ``required_params``.
        required_params: Optional mapping of model ID -> required OpenRouter
            request parameters, e.g. ``{"anthropic/claude-sonnet-4": ["tools"]}``.
            If provided, the built-in ``"tool_support"`` check validates these
            against each model's ``supported_parameters``.
        common_required_params: Optional list of request parameters required for
            every model in ``models`` (e.g. ``["tools"]``).
        site_name: Value for the ``HTTP-Referer`` / ``X-Title`` headers on the smoke test.
        raise_on_failure: If True (default), raise :class:`PreflightError` on any failure.
        infra: Optional list of :class:`InfraCheck` instances for validating
            downstream infrastructure (health endpoints, env vars, Docker
            containers, git versions). See :mod:`agent_preflight.health`.

    Returns:
        List of failure strings (empty if all passed).

    Raises:
        PreflightError: If ``raise_on_failure`` is True and any checks fail.
    """
    unique_models = list(dict.fromkeys(models))
    common_required_params = list(dict.fromkeys(common_required_params or []))
    merged_required_params: dict[str, list[str]] = {}

    if common_required_params:
        for model_id in unique_models:
            merged_required_params[model_id] = list(common_required_params)

    for model_id, params in (required_params or {}).items():
        existing = merged_required_params.setdefault(model_id, [])
        merged_required_params[model_id] = list(dict.fromkeys(existing + params))

    required_params = merged_required_params

    all_checks = {"balance", "math", "models", "key_smoke", "cache_smoke"}
    if required_params:
        all_checks.add("tool_support")
    if docker_image:
        all_checks.add("docker_image")
    enabled = set(checks) if checks is not None else all_checks

    print(f"\n  {'='*60}")
    print(f"  PREFLIGHT CHECKS")
    print(f"  {'='*60}")
    failures: list[str] = []

    # 1. Account balance
    balance: float | None = None
    get_account_balance_fn = resolve_public(
        "get_account_balance",
        get_account_balance,
        original=provision_get_account_balance,
    )
    if "balance" in enabled:
        balance, total, used = get_account_balance_fn(or_key)
        if balance is not None:
            ok = balance > 0.10
            _check("Account balance", ok,
                   f"${balance:.2f} available (${total:.2f} purchased, ${used:.2f} used)")
            if not ok:
                failures.append(f"Account balance too low (${balance:.2f})")
        else:
            _check("Account balance", False, "could not fetch")
            failures.append("Could not fetch account balance")
    else:
        # Still need balance value for math check — fetch quietly
        balance, _, _ = get_account_balance_fn(or_key)

    # 2. Provisioning math
    if "math" in enabled:
        n = len(unique_models)
        needed = n * budget_per_model
        if balance is not None:
            ok = balance >= needed
            detail = (
                f"{n} models x ${budget_per_model:.2f} = ${needed:.2f} "
                f"({'OK' if ok else 'INSUFFICIENT'}: balance ${balance:.2f})"
            )
            _check("Provisioning budget", ok, detail)
            if not ok:
                failures.append(f"Insufficient balance: need ${needed:.2f}, have ${balance:.2f}")
        else:
            _check("Provisioning budget", None,
                   f"{n} x ${budget_per_model:.2f} = ${needed:.2f} (balance unknown)")

    available_models: dict[str, dict] | None = None
    available_models_loaded = False

    def _load_available_models() -> dict[str, dict] | None:
        nonlocal available_models, available_models_loaded
        if not available_models_loaded:
            get_available_models_fn = resolve_public(
                "get_available_models",
                get_available_models,
                original=provision_get_available_models,
            )
            available_models = get_available_models_fn(or_key)
            available_models_loaded = True
        return available_models

    # 3. Model availability + pricing
    if "models" in enabled:
        available_models = _load_available_models()
        if available_models is None:
            _check("Model availability", False, "could not fetch model list")
            failures.append("Could not fetch model list from OpenRouter")
        else:
            for model_id in unique_models:
                if model_id in available_models:
                    p = available_models[model_id].get("pricing", {})
                    inp = float(p.get("prompt", 0)) * 1_000_000
                    out = float(p.get("completion", 0)) * 1_000_000
                    _check(f"Model {model_id}", True, f"${inp:.2f}/${out:.2f} per M tokens in/out")
                else:
                    _check(f"Model {model_id}", False, "not found on OpenRouter")
                    failures.append(f"Model {model_id} not on OpenRouter")

    # 3b. Parameter / tool support
    if "tool_support" in enabled:
        if not required_params:
            _check(
                "Parameter support",
                None,
                "no required_params/common_required_params specified — skipped",
            )
        else:
            available_models = _load_available_models()
            if available_models is None:
                _check("Parameter support", False, "could not fetch model list")
                failures.append("Could not fetch model list from OpenRouter for parameter support")
            else:
                for model_id, params in required_params.items():
                    if model_id not in available_models:
                        _check(f"Params {model_id}", False, "model not found on OpenRouter")
                        failures.append(f"Model {model_id} not on OpenRouter for parameter support check")
                        continue

                    model_info = available_models[model_id]
                    supported = set(model_info.get("supported_parameters") or [])
                    missing = [p for p in params if p not in supported]
                    if missing:
                        _check(
                            f"Params {model_id}",
                            False,
                            f"missing {', '.join(missing)}; supports {len(supported)} parameters",
                        )
                        failures.append(
                            f"Model {model_id} missing required params: {', '.join(missing)}"
                        )
                    else:
                        _check(
                            f"Params {model_id}",
                            True,
                            f"supports {', '.join(params)}",
                        )

    # 4. Key provisioning smoke test (1¢)
    if "key_smoke" in enabled:
        try:
            provision_key_fn = resolve_public(
                "provision_key",
                provision_key,
                original=provision_key_impl,
            )
            test_key = provision_key_fn(admin_key, f"{site_name}:preflight-smoke", limit_usd=0.01)
            resp = httpx.post(
                f"{OR_BASE}/chat/completions",
                headers={
                    "Authorization": f"Bearer {test_key}",
                    "Content-Type": "application/json",
                    "HTTP-Referer": f"https://{site_name}",
                    "X-Title": f"{site_name} preflight",
                },
                json={
                    "model": "openai/gpt-4.1-nano",
                    "messages": [{"role": "user", "content": "Say OK"}],
                    "max_tokens": 16,
                },
                timeout=30.0,
            )
            if resp.status_code != 200:
                raise RuntimeError(f"API call failed ({resp.status_code}): {resp.text[:120]}")
            reply = (
                resp.json().get("choices", [{}])[0]
                .get("message", {}).get("content", "")
            )
            get_usage_fn = resolve_public("get_usage", get_usage, original=provision_get_usage)
            usage = get_usage_fn(test_key)
            usage_str = f", usage=${usage:.6f}" if usage is not None else ""
            _check("Key provisioning (1¢ smoke)", True,
                   f"key={test_key[:16]}… reply={reply!r}{usage_str}")
        except Exception as exc:
            _check("Key provisioning (1¢ smoke)", False, str(exc))
            failures.append(f"Key provisioning failed: {exc}")

    # 5. Cache smoke test — verify prompt caching works via OpenRouter
    if "cache_smoke" in enabled:
        anthropic_models = [m for m in models if "anthropic/" in m.lower() or "claude" in m.lower()]
        if anthropic_models:
            try:
                import time as _time
                from .caching import min_cache_tokens

                cache_model = anthropic_models[0]
                min_tokens = min_cache_tokens(cache_model)

                # Pad system prompt to exceed the model's minimum cacheable size.
                # ~13 tokens per rule line, so 400 lines ≈ 5200 tokens (covers 4096).
                n_rules = max(100, (min_tokens // 13) + 50)
                est_prompt_tokens = n_rules * 13
                padding = (
                    "You are a meticulous coding assistant. "
                    "Follow these rules:\n" +
                    "\n".join(f"- Rule {i}: Always write clean, well-documented code." for i in range(1, n_rules + 1))
                )

                if available_models is None and not available_models_loaded:
                    available_models = _load_available_models()
                model_info = available_models.get(cache_model, {}) if available_models else {}
                pricing = model_info.get("pricing", {}) if isinstance(model_info, dict) else {}
                prompt_price = pricing.get("prompt")
                limit_usd = _cache_smoke_limit_usd(
                    prompt_tokens=est_prompt_tokens,
                    prompt_price_per_token=float(prompt_price) if prompt_price is not None else None,
                )
                test_key = provision_key(
                    admin_key,
                    f"{site_name}:cache-smoke",
                    limit_usd=limit_usd,
                )

                # Use explicit Anthropic cache_control on the system block.
                # This is the most portable Anthropic caching mode on
                # OpenRouter's current routes.
                request_body = {
                    "model": cache_model,
                    "messages": [
                        {
                            "role": "system",
                            "content": [
                                {
                                    "type": "text",
                                    "text": padding,
                                    "cache_control": {"type": "ephemeral"},
                                }
                            ],
                        },
                        {"role": "user", "content": "Say OK"},
                    ],
                    "max_tokens": 4,
                }
                headers = {
                    "Authorization": f"Bearer {test_key}",
                    "Content-Type": "application/json",
                    "HTTP-Referer": f"https://{site_name}",
                    "X-Title": f"{site_name} cache-smoke",
                }

                # First request: populates the cache
                resp1 = httpx.post(
                    f"{OR_BASE}/chat/completions",
                    headers=headers,
                    json=request_body,
                    timeout=30.0,
                )
                if resp1.status_code != 200:
                    raise RuntimeError(f"Cache write request failed ({resp1.status_code})")
                u1 = resp1.json().get("usage", {})
                d1 = u1.get("prompt_tokens_details", {}) or {}
                # OpenRouter normalized: cache_write_tokens
                # Fallback: Anthropic native cache_creation_input_tokens
                creation = d1.get("cache_write_tokens", 0) or u1.get("cache_creation_input_tokens", 0)

                # Brief pause then identical request — should hit cache
                # OpenRouter sticky routing should route to the same provider.
                _time.sleep(1)
                resp2 = httpx.post(
                    f"{OR_BASE}/chat/completions",
                    headers=headers,
                    json=request_body,
                    timeout=30.0,
                )
                if resp2.status_code != 200:
                    raise RuntimeError(f"Cache read request failed ({resp2.status_code})")
                u2 = resp2.json().get("usage", {})
                d2 = u2.get("prompt_tokens_details", {}) or {}
                # OpenRouter normalized: cached_tokens
                # Fallback: Anthropic native cache_read_input_tokens
                cache_read = d2.get("cached_tokens", 0) or u2.get("cache_read_input_tokens", 0)

                if cache_read > 0:
                    _check("Prompt caching", True,
                           f"model={cache_model} created={creation} read={cache_read} tokens "
                           f"(min={min_tokens}, limit=${limit_usd:.2f})")
                elif creation > 0:
                    # Cache was created but second request didn't hit.
                    # With sticky routing this is unusual but possible if the
                    # provider instance changed between requests.
                    _check("Prompt caching", None,
                           f"model={cache_model} cache created ({creation} tokens) but no read hit "
                           f"— sticky routing may not have activated yet "
                           f"(limit=${limit_usd:.2f})")
                else:
                    # No cache fields at all.
                    _check("Prompt caching", False,
                           f"model={cache_model} — no cache tokens in response. "
                           f"Prompt may be below {min_tokens} token minimum, or "
                           f"cache_control not supported for this model/route. "
                           f"(limit=${limit_usd:.2f})")
                    failures.append(
                        f"Cache smoke failed for {cache_model}: no cache tokens in response"
                    )
            except Exception as exc:
                _check("Prompt caching", False, f"cache smoke error: {exc}")
                failures.append(f"Cache smoke failed: {exc}")
        else:
            _check("Prompt caching", None, "no Anthropic models in list — skipped")

    # 6. Docker image
    if "docker_image" in enabled and docker_image:
        r = subprocess.run(
            ["docker", "image", "inspect", docker_image],
            capture_output=True, text=True,
        )
        if r.returncode == 0:
            try:
                import json
                created = json.loads(r.stdout)[0].get("Created", "")[:19]
                _check(f"Docker image {docker_image!r}", True, f"created {created}")
            except Exception:
                _check(f"Docker image {docker_image!r}", True)
        else:
            _check(f"Docker image {docker_image!r}", False, "not found — build it first")
            failures.append(f"Docker image '{docker_image}' not found")

    # 7. Custom checks
    if custom_checks:
        custom_checks_need_ctx = any(_accepts_context(fn) for fn in custom_checks)
        if custom_checks_need_ctx and available_models is None and not available_models_loaded:
            available_models = _load_available_models()

        ctx: dict[str, Any] = {
            "admin_key": admin_key,
            "or_key": or_key,
            "models": unique_models,
            "budget_per_model": budget_per_model,
            "checks": sorted(enabled),
            "docker_image": docker_image,
            "site_name": site_name,
            "required_params": required_params,
            "balance": balance,
            "available_models": available_models,
        }
        for fn in custom_checks:
            try:
                name, ok, detail = _run_custom_check(fn, ctx)
                _check(name, ok, detail)
                if ok is False:
                    failures.append(f"{name}: {detail}")
            except Exception as exc:
                _check(f"Custom check ({fn.__name__})", False, str(exc))
                failures.append(f"Custom check {fn.__name__} raised: {exc}")

    # 8. Infrastructure readiness checks
    if infra:
        for ic in infra:
            try:
                result = ic.check()
                _check(ic.name, result.ok, result.detail, hint=result.hint)
                if result.ok is False:
                    msg = f"{ic.name}: {result.detail}"
                    if result.hint:
                        msg += f" -> {result.hint}"
                    failures.append(msg)
            except Exception as exc:
                _check(f"Infra check ({ic.name})", False, str(exc))
                failures.append(f"Infra check {ic.name} raised: {exc}")

    print()
    if failures:
        print(f"  PREFLIGHT FAILED ({len(failures)} issue(s)):")
        for f in failures:
            print(f"    x {f}")
        print()
        if raise_on_failure:
            raise PreflightError(failures)
    else:
        print("  All preflight checks passed.\n")

    return failures
