"""Pricing data and cost estimation for noxaudit estimate command."""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class ModelPricing:
    """Per-model pricing data."""

    input_per_million: float  # $/M tokens, standard tier
    output_per_million: float  # $/M tokens, standard tier
    tier_threshold: int | None  # Input token count above which tiered rates apply
    input_per_million_high: float | None  # $/M tokens, high tier
    output_per_million_high: float | None  # $/M tokens, high tier output
    batch_discount: float = 0.0  # Fraction discount (0.5 = 50%)
    context_window: int = 200_000
    cache_read_per_million: float = 0.0  # $/M cache-read tokens (flat rate, not tiered)
    cache_write_per_million: float = 0.0  # $/M cache-write tokens (flat rate, not tiered)


MODEL_PRICING: dict[str, ModelPricing] = {
    # --- Anthropic (verified 2026-03-06 from platform.claude.com) ---
    "claude-opus-4-6": ModelPricing(
        input_per_million=5.00,
        output_per_million=25.00,
        tier_threshold=200_000,
        input_per_million_high=10.00,
        output_per_million_high=37.50,
        batch_discount=0.50,
        context_window=200_000,
        cache_read_per_million=0.50,
        cache_write_per_million=6.25,
    ),
    "claude-sonnet-4-6": ModelPricing(
        input_per_million=3.00,
        output_per_million=15.00,
        tier_threshold=200_000,
        input_per_million_high=6.00,
        output_per_million_high=22.50,
        batch_discount=0.50,
        context_window=200_000,
        cache_read_per_million=0.30,
        cache_write_per_million=3.75,
    ),
    "claude-haiku-4-5": ModelPricing(
        input_per_million=1.00,
        output_per_million=5.00,
        tier_threshold=None,
        input_per_million_high=None,
        output_per_million_high=None,
        batch_discount=0.50,
        context_window=200_000,
        cache_read_per_million=0.10,
        cache_write_per_million=1.25,
    ),
    # --- Google Gemini (verified 2026-03-06 from ai.google.dev/gemini-api/docs/pricing) ---
    "gemini-2.5-flash": ModelPricing(
        input_per_million=0.30,
        output_per_million=2.50,
        tier_threshold=None,
        input_per_million_high=None,
        output_per_million_high=None,
        batch_discount=0.50,
        context_window=1_000_000,
    ),
    "gemini-3-flash-preview": ModelPricing(
        input_per_million=0.50,
        output_per_million=3.00,
        tier_threshold=None,
        input_per_million_high=None,
        output_per_million_high=None,
        batch_discount=0.50,
        context_window=1_000_000,
    ),
    "gemini-2.5-pro": ModelPricing(
        input_per_million=1.25,
        output_per_million=10.00,
        tier_threshold=200_000,
        input_per_million_high=2.50,
        output_per_million_high=15.00,
        batch_discount=0.50,
        context_window=1_000_000,
    ),
    # --- OpenAI (verified 2026-03-06 from developers.openai.com/api/docs/pricing) ---
    "gpt-5-nano": ModelPricing(
        input_per_million=0.05,
        output_per_million=0.40,
        tier_threshold=None,
        input_per_million_high=None,
        output_per_million_high=None,
        batch_discount=0.50,
        context_window=400_000,
    ),
    "gpt-5-mini": ModelPricing(
        input_per_million=0.25,
        output_per_million=2.00,
        tier_threshold=None,
        input_per_million_high=None,
        output_per_million_high=None,
        batch_discount=0.50,
        context_window=400_000,
    ),
    "o4-mini": ModelPricing(
        input_per_million=1.10,
        output_per_million=4.40,
        tier_threshold=None,
        input_per_million_high=None,
        output_per_million_high=None,
        batch_discount=0.50,
        context_window=200_000,
    ),
    "gpt-5.4": ModelPricing(
        input_per_million=2.50,
        output_per_million=15.00,
        tier_threshold=272_000,
        input_per_million_high=5.00,
        output_per_million_high=22.50,
        batch_discount=0.50,
        context_window=1_050_000,
    ),
}

# Provider name for each model key
_MODEL_PROVIDER: dict[str, str] = {
    "claude-opus-4-6": "anthropic",
    "claude-sonnet-4-6": "anthropic",
    "claude-haiku-4-5": "anthropic",
    "gemini-2.5-flash": "gemini",
    "gemini-3-flash-preview": "gemini",
    "gemini-2.5-pro": "gemini",
    "gpt-5-nano": "openai",
    "gpt-5-mini": "openai",
    "o4-mini": "openai",
    "gpt-5.4": "openai",
}


def resolve_model_key(provider: str, model: str) -> str:
    """Map a provider + model name to a MODEL_PRICING key.

    Tries direct lookup first, then falls back to provider-based heuristics.
    """
    if model in MODEL_PRICING:
        return model
    model_lower = model.lower()
    if provider == "anthropic":
        if "opus" in model_lower:
            return "claude-opus-4-6"
        if "haiku" in model_lower:
            return "claude-haiku-4-5"
        return "claude-sonnet-4-6"
    elif provider == "gemini":
        if "3" in model_lower and "flash" in model_lower:
            return "gemini-3-flash-preview"
        if "pro" in model_lower:
            return "gemini-2.5-pro"
        if "2.5" in model or "2-5" in model:
            return "gemini-2.5-flash"
        return "gemini-2.5-flash"
    elif provider == "openai":
        if "nano" in model_lower:
            return "gpt-5-nano"
        if "mini" in model_lower and "o4" not in model_lower:
            return "gpt-5-mini"
        if model_lower.startswith("o4"):
            return "o4-mini"
        if "5.4" in model:
            return "gpt-5.4"
        return "gpt-5.4"
    return "gemini-2.5-flash"


def estimate_cost(
    input_tokens: int,
    output_tokens: int,
    pricing: ModelPricing,
    use_batch: bool = True,
    cache_read_tokens: int = 0,
    cache_write_tokens: int = 0,
) -> float:
    """Calculate total cost for given token counts.

    Handles tiered pricing split at threshold. When input exceeds the tier
    threshold, output tokens are also billed at the high tier rate (matching
    Anthropic actual pricing). Cache tokens are billed at flat rates (not
    tiered). Batch discount applied last to the combined total.
    """
    if (
        input_tokens == 0
        and output_tokens == 0
        and cache_read_tokens == 0
        and cache_write_tokens == 0
    ):
        return 0.0

    if pricing.tier_threshold and input_tokens > pricing.tier_threshold:
        standard_input = pricing.tier_threshold
        high_input = input_tokens - pricing.tier_threshold
        input_cost = (standard_input / 1_000_000) * pricing.input_per_million
        input_cost += (high_input / 1_000_000) * (pricing.input_per_million_high or 0.0)
        output_cost = (output_tokens / 1_000_000) * (pricing.output_per_million_high or 0.0)
    else:
        input_cost = (input_tokens / 1_000_000) * pricing.input_per_million
        output_cost = (output_tokens / 1_000_000) * pricing.output_per_million

    # Cache tokens are billed at flat rates, not subject to tiered pricing
    cache_read_cost = (cache_read_tokens / 1_000_000) * pricing.cache_read_per_million
    cache_write_cost = (cache_write_tokens / 1_000_000) * pricing.cache_write_per_million

    total = input_cost + output_cost + cache_read_cost + cache_write_cost

    if use_batch and pricing.batch_discount > 0:
        total *= 1.0 - pricing.batch_discount

    return total


def estimate_output_tokens(input_tokens: int, num_focus_areas: int = 1) -> int:
    """Estimate output tokens: min(16384, input * 0.1) per focus area."""
    per_focus = min(16384, int(input_tokens * 0.1))
    return per_focus * max(1, num_focus_areas)


def estimate_prepass_reduction(files: list, total_tokens: int) -> dict:
    """Estimate what a pre-pass triage would reduce tokens to without running it.

    Sorts files by content size ascending, applies 15% high / 30% medium /
    30% low / 25% skip distribution. Kept files (high + medium) contribute
    to the reduced token count. Triage cost uses gemini-2.5-flash at ~60
    output tokens per file.
    """
    if not files:
        return {"reduced_tokens": 0, "triage_cost": 0.0, "high": 0, "medium": 0, "low_skip": 0}

    n = len(files)
    sorted_files = sorted(files, key=lambda f: len(f.content))

    high_count = max(1, round(n * 0.15))
    medium_count = max(1, round(n * 0.30))
    low_skip_count = n - high_count - medium_count

    # Kept files: high + medium priority (smallest files in ascending sort)
    kept_files = sorted_files[: high_count + medium_count]
    reduced_tokens = sum(len(f.content) // 4 for f in kept_files)

    # Triage cost: run gemini-2.5-flash on all files (~60 output tokens per file)
    flash_pricing = MODEL_PRICING["gemini-2.5-flash"]
    triage_output = n * 60
    triage_cost = estimate_cost(total_tokens, triage_output, flash_pricing, use_batch=False)

    return {
        "reduced_tokens": reduced_tokens,
        "triage_cost": triage_cost,
        "high": high_count,
        "medium": medium_count,
        "low_skip": low_skip_count,
    }


def _fmt_tokens(n: int) -> str:
    """Format a token count as a human-readable string (287000 → '287K')."""
    if n >= 1_000_000:
        return f"{n / 1_000_000:.1f}M"
    if n >= 1_000:
        return f"{n // 1_000}K"
    return str(n)


def build_estimate_report(
    repo_name: str,
    focus_names: list[str],
    files: list,
    provider_name: str,
    model_key: str,
) -> str:
    """Assemble a human-readable cost estimate report."""
    lines: list[str] = [""]

    # Header
    focus_display = " + ".join(focus_names)
    header = f"  {repo_name} — {focus_display}"
    lines.append(header)
    lines.append("")

    # Token count (rough: 1 token per 4 chars)
    total_tokens = sum(len(f.content) // 4 for f in files)
    lines.append(f"  Files:     {len(files)} files, {_fmt_tokens(total_tokens)} tokens")

    pricing = MODEL_PRICING[model_key]
    use_batch = pricing.batch_discount > 0
    lines.append(f"  Provider:  {provider_name} ({model_key})")
    lines.append("")

    output_tokens = estimate_output_tokens(total_tokens, len(focus_names))
    cost = estimate_cost(total_tokens, output_tokens, pricing, use_batch=use_batch)
    tiered = bool(pricing.tier_threshold and total_tokens > pricing.tier_threshold)

    if tiered:
        lines.append(f"  ! Cost estimate: ~${cost:.2f}")
        lines.append(
            f"    Your token count ({_fmt_tokens(total_tokens)}) exceeds "
            f"Anthropic's 200K standard tier."
        )
        lines.append(
            f"    Tiered pricing applies: ${pricing.input_per_million_high:.2f}/M input "
            f"(2x standard rate)."
        )
        if use_batch:
            lines.append("    Batch API 50% discount applied.")
    else:
        lines.append(f"  Cost estimate: ~${cost:.2f}")
        if use_batch:
            lines.append("    Batch API 50% discount applied.")
    lines.append("")

    # Alternatives (cheaper models)
    alternatives: list[tuple[str, str, float, int]] = []
    for alt_key, alt_pricing in MODEL_PRICING.items():
        if alt_key == model_key:
            continue
        alt_provider = _MODEL_PROVIDER[alt_key]
        alt_output = estimate_output_tokens(total_tokens, len(focus_names))
        alt_use_batch = alt_pricing.batch_discount > 0
        alt_cost = estimate_cost(total_tokens, alt_output, alt_pricing, use_batch=alt_use_batch)
        if cost > 0 and alt_cost < cost:
            savings_pct = int((1.0 - alt_cost / cost) * 100)
            alternatives.append((alt_key, alt_provider, alt_cost, savings_pct))

    alternatives.sort(key=lambda x: x[2])

    # Pre-pass alternative (only when Anthropic exceeds tier threshold)
    prepass_info: tuple | None = None
    if tiered and provider_name == "anthropic":
        prepass = estimate_prepass_reduction(files, total_tokens)
        reduced_tokens = prepass["reduced_tokens"]
        reduced_output = estimate_output_tokens(reduced_tokens, len(focus_names))
        no_tier_cost = estimate_cost(reduced_tokens, reduced_output, pricing, use_batch=use_batch)
        total_prepass_cost = no_tier_cost + prepass["triage_cost"]
        if cost > 0 and total_prepass_cost < cost:
            pp_savings_pct = int((1.0 - total_prepass_cost / cost) * 100)
            prepass_info = (prepass, reduced_tokens, total_prepass_cost, pp_savings_pct)

    if alternatives or prepass_info:
        lines.append("  Alternatives:")
        for alt_key, alt_provider, alt_cost, savings_pct in alternatives:
            desc = f"{alt_provider} ({alt_key})"
            note = ""
            if alt_provider == "gemini" and savings_pct >= 90:
                note = " — best value"
            elif "sonnet" in alt_key:
                note = " — similar quality, lower cost"
            lines.append(f"    {desc:<38} ~${alt_cost:.2f}   {savings_pct}% cheaper{note}")
        if prepass_info:
            _, _, total_prepass_cost, pp_savings_pct = prepass_info
            desc = f"{provider_name} + pre-pass"
            note = f" — Flash triage keeps {model_key} under 200K tier"
            lines.append(
                f"    {desc:<38} ~${total_prepass_cost:.2f}   {pp_savings_pct}% cheaper{note}"
            )
        lines.append("")

    # Pre-pass detail section
    if prepass_info:
        prepass, reduced_tokens, total_prepass_cost, _ = prepass_info
        lines.append("  Pre-pass estimate:")
        lines.append(f"    Gemini Flash triage: ~${prepass['triage_cost']:.2f}")
        lines.append(
            f"    Expected reduction: {_fmt_tokens(total_tokens)} → "
            f"~{_fmt_tokens(reduced_tokens)} tokens "
            f"(high: {prepass['high']} files, medium: {prepass['medium']}, "
            f"low/skip: {prepass['low_skip']})"
        )
        lines.append(
            f"    With pre-pass, {model_key} stays in standard pricing tier "
            f"(${pricing.input_per_million:.2f}/M vs "
            f"${pricing.input_per_million_high:.2f}/M)"
        )
        lines.append("")

    # Monthly projection (if run daily)
    monthly_runs = 30
    monthly_cost = cost * monthly_runs
    lines.append(f"  Monthly estimate: ~${monthly_cost:.2f} (if run daily)")

    if alternatives:
        cheapest_key, cheapest_provider, cheapest_cost, _ = alternatives[0]
        monthly_cheapest = cheapest_cost * monthly_runs
        lines.append(f"  Monthly with {cheapest_key}: ~${monthly_cheapest:.2f}")

    lines.append("")

    # Recommendation for expensive Anthropic with cheap gemini alternative
    if tiered and alternatives:
        cheapest_key, cheapest_provider, cheapest_cost, _ = alternatives[0]
        if cheapest_provider == "gemini":
            lines.append(
                f"  Recommendation: {cheapest_key} offers significantly lower cost (~${cheapest_cost:.2f}/run)."
            )
            lines.append("  Reserve anthropic for deep dives where finding depth matters most.")
            lines.append("")

    return "\n".join(lines)
