#!/usr/bin/env python3
"""Token cost tracking and optimization levers.

Estimates cost per call, aggregates by endpoint/model, and demonstrates the
cost-cutting levers: prompt caching, model routing, context compression.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from collections import defaultdict


# --- Pricing table (USD per 1M tokens; illustrative) ---
PRICING: dict[str, tuple[float, float]] = {
    # model: (input_price, output_price) per 1M tokens
    "gpt-4o": (2.50, 10.00),
    "gpt-4o-mini": (0.15, 0.60),
    "claude-3.5-sonnet": (3.00, 15.00),
    "llama-3.1-70b": (0.60, 0.60),
}


@dataclass
class UsageRecord:
    model: str
    input_tokens: int
    output_tokens: int
    cached_input_tokens: int = 0
    endpoint: str = "default"


@dataclass
class CostTracker:
    records: list[UsageRecord] = field(default_factory=list)

    def add(self, rec: UsageRecord) -> None:
        self.records.append(rec)

    def cost_of(self, rec: UsageRecord) -> float:
        in_price, out_price = PRICING[rec.model]
        # Cached input tokens are typically ~90% cheaper
        cached_cost = rec.cached_input_tokens * in_price * 0.1 / 1_000_000
        fresh_input = (rec.input_tokens - rec.cached_input_tokens) * in_price / 1_000_000
        output_cost = rec.output_tokens * out_price / 1_000_000
        return cached_cost + fresh_input + output_cost

    def total(self) -> float:
        return sum(self.cost_of(r) for r in self.records)

    def by_endpoint(self) -> dict[str, float]:
        agg: dict[str, float] = defaultdict(float)
        for r in self.records:
            agg[r.endpoint] += self.cost_of(r)
        return dict(agg)

    def by_model(self) -> dict[str, float]:
        agg: dict[str, float] = defaultdict(float)
        for r in self.records:
            agg[r.model] += self.cost_of(r)
        return dict(agg)

    def report(self) -> None:
        print(f"\n=== Cost Report ===")
        print(f"Total cost: ${self.total():.4f}")
        print(f"\nBy endpoint:")
        for ep, cost in sorted(self.by_endpoint().items(), key=lambda x: -x[1]):
            print(f"  {ep:30s} ${cost:.4f}")
        print(f"\nBy model:")
        for m, cost in sorted(self.by_model().items(), key=lambda x: -x[1]):
            print(f"  {m:30s} ${cost:.4f}")


def demo_model_routing() -> None:
    """Show the cost benefit of routing easy calls to a cheap model."""
    print("\n=== Model Routing Demo ===")
    # Scenario: 100 classification calls per day
    classification_calls = 100
    # All on gpt-4o vs routed to gpt-4o-mini
    expensive = classification_calls * (500 * 2.50 + 50 * 10.00) / 1_000_000
    cheap = classification_calls * (500 * 0.15 + 50 * 0.60) / 1_000_000
    print(f"All on gpt-4o:       ${expensive:.4f}/day")
    print(f"Routed to 4o-mini:   ${cheap:.4f}/day")
    print(f"Savings:             ${expensive - cheap:.4f}/day "
          f"({(1 - cheap/expensive)*100:.0f}%)")


def demo_prompt_caching() -> None:
    """Show the benefit of caching stable prefixes."""
    print("\n=== Prompt Caching Demo ===")
    # Stable system prompt (500 tokens) + few-shot examples (1000 tokens)
    stable_prefix = 1500
    per_call_input = 200  # variable part
    calls = 1000

    no_cache = calls * (stable_prefix + per_call_input) * 2.50 / 1_000_000
    with_cache = calls * (
        stable_prefix * 2.50 * 0.1 + per_call_input * 2.50
    ) / 1_000_000

    print(f"Without caching: ${no_cache:.4f}")
    print(f"With caching:    ${with_cache:.4f}")
    print(f"Savings:         ${no_cache - with_cache:.4f} "
          f"({(1 - with_cache/no_cache)*100:.0f}%)")


def main() -> None:
    # Build a usage record set
    tracker = CostTracker()
    tracker.add(UsageRecord("gpt-4o", 5000, 1000, endpoint="rag_generation"))
    tracker.add(UsageRecord("gpt-4o", 5000, 1000, cached_input_tokens=4500,
                            endpoint="rag_generation"))
    tracker.add(UsageRecord("gpt-4o-mini", 2000, 100, endpoint="classification"))
    tracker.add(UsageRecord("claude-3.5-sonnet", 3000, 800, endpoint="summarization"))

    tracker.report()
    demo_model_routing()
    demo_prompt_caching()

    print("\n=== Key levers ===")
    print("1. Prompt caching: ~90% cheaper on cached prefix hits")
    print("2. Model routing: 10-50x cheaper on routed share")
    print("3. Context compression: retrieve less, truncate history")
    print("4. Batch: amortize fixed overhead across many calls")
    print("5. Stream + early-exit: don't pay for tokens you don't need")


if __name__ == "__main__":
    main()