//! Per-(provider, model) token pricing and cost calculation.
//!
//! Prices are stored as integers — specifically, *micro-cents per token*
//! (1/1,000,000 of a cent), equivalently picodollars (10^-12 USD) per
//! token. Reasons:
//!
//!   - All commonly-quoted "USD per million tokens" prices land on round
//!     integers under the conversion: $3.00 / 1M tokens = 300
//!     micro-cents per token.
//!   - Per-message token counts (10^2 to 10^5) multiplied by per-token
//!     prices (~10^2) stay well inside `u64` for entire long-running
//!     sessions — a session worth ~$1 (10^14 micro-cents) is nowhere
//!     near `u64.max` (~1.8 × 10^19).
//!   - Summing per-turn costs across a session is exact: no
//!     floating-point drift.
//!
//! The TOML on-disk format lets users write `input = 3.0` for $3/Mtok,
//! which is the natural unit for humans. The loader multiplies by 100
//! and rounds to the nearest integer; the rounding handles parse-time
//! float-precision wobble.
//!
//! `costMicroCents` does the integer arithmetic. Each `Pricing` field
//! is `?u64`: `null` means "we don't know the price for this token
//! category," which is distinct from a price of zero (e.g. OpenAI does
//! charge nothing for cache writes — that's a known 0, not unknown).
//! `costMicroCents` returns `?u64`: if any usage category with nonzero
//! count maps to a `null` price, the whole turn cost goes to `null`
//! ("unknown"), and any session-level sum that includes that turn must
//! likewise degenerate to `null`. This prevents silently treating
//! unknown costs as free.
//!
//! The display layer lives in the CLI; libpanto only computes.

const std = @import("std");

const session_mod = @import("session.zig");
pub const Usage = session_mod.Usage;

// =============================================================================
// Pricing struct
// =============================================================================

/// Per-token pricing for a single (provider, model) pair.
///
/// Units: micro-cents per token (1/1,000,000 of a cent per token), aka
/// picodollars (10^-12 USD) per token.
///
/// Conversion from "USD per million tokens":
///
///     micro_cents_per_token = round(USD_per_Mtok * 100)
///
/// So $3.00/Mtok → 300 micro-cents/token.
pub const Pricing = struct {
    /// Per fresh (uncached, non-cache-written) input token. `null` =
    /// unknown (e.g. field omitted from `models.toml`).
    input: ?u64 = null,
    /// Per output token. `null` = unknown.
    output: ?u64 = null,
    /// Per cache-read input token. Typically a fraction of `input`
    /// (0.1× on Anthropic, 0.5× on OpenAI for cached prompt tokens).
    /// `null` = unknown.
    cache_read: ?u64 = null,
    /// Per cache-write input token. Anthropic charges a premium
    /// (1.25× `input`); OpenAI doesn't bill a cache-write rate (a
    /// known 0, which should be written `cache_write = 0` rather than
    /// omitted). `null` = unknown.
    cache_write: ?u64 = null,

    /// Convert a USD-per-million-tokens float (the human-friendly unit)
    /// to the internal integer representation. Rounds to nearest.
    ///
    /// `dollars_per_mtok * 1_000_000 cents/dollar / 1_000_000 tokens` =
    /// cents-per-token, then * 1_000_000 micro-cents/cent =
    /// micro-cents-per-token. The 1_000_000s cancel, leaving
    /// `dollars_per_mtok * 100`.
    ///
    /// Non-finite inputs round to 0. Negative inputs are clamped to 0
    /// (treated as a known free price, not unknown).
    pub fn fromDollarsPerMtok(dollars_per_mtok: f64) u64 {
        if (!std.math.isFinite(dollars_per_mtok) or dollars_per_mtok <= 0) return 0;
        const scaled = dollars_per_mtok * 100.0;
        const r = @round(scaled);
        if (r >= @as(f64, @floatFromInt(std.math.maxInt(u64)))) return std.math.maxInt(u64);
        return @intFromFloat(r);
    }
};

// =============================================================================
// Cost calculation
// =============================================================================

/// Compute the cost of a single turn's `usage` under the given `pricing`,
/// in micro-cents. Returns `null` if any usage category with a nonzero
/// token count maps to a `null` price — i.e. "we used some cache reads
/// but we don't know the cache-read price" poisons the whole turn cost
/// to "unknown". Categories with zero tokens are ignored regardless of
/// whether their price is known, so e.g. a model with `cache_write =
/// null` still produces a known cost on turns that never write to
/// cache.
///
/// `reasoning` does NOT contribute separately — it's already counted
/// inside `output`.
pub fn costMicroCents(usage: Usage, pricing: Pricing) ?u64 {
    var total: u64 = 0;
    total +%= component(usage.input, pricing.input) orelse return null;
    total +%= component(usage.output, pricing.output) orelse return null;
    total +%= component(usage.cache_read, pricing.cache_read) orelse return null;
    total +%= component(usage.cache_write, pricing.cache_write) orelse return null;
    return total;
}

/// Cost contribution (in micro-cents) from a single (tokens, price)
/// pair. Returns `0` when `tokens == 0` regardless of whether `price`
/// is known — so unknown prices don't poison turns that never used
/// that category. Returns `null` when tokens are nonzero but the
/// price is unknown; callers convert that to a `null` total.
fn component(tokens: u64, price: ?u64) ?u64 {
    if (tokens == 0) return 0;
    const p = price orelse return null;
    return tokens *% p;
}

// =============================================================================
// Registry
// =============================================================================

/// In-memory registry of `(provider, model) -> Pricing`. Lookups are by
/// exact match of both fields.
///
/// The registry owns the (provider, model) key strings; entries are
/// pushed via `set` (which duplicates the inputs).
pub const Registry = struct {
    allocator: std.mem.Allocator,
    entries: std.ArrayList(Entry),

    pub const Entry = struct {
        provider: []u8,
        model: []u8,
        pricing: Pricing,
    };

    pub fn init(allocator: std.mem.Allocator) Registry {
        return .{ .allocator = allocator, .entries = .empty };
    }

    pub fn deinit(self: *Registry) void {
        for (self.entries.items) |e| {
            self.allocator.free(e.provider);
            self.allocator.free(e.model);
        }
        self.entries.deinit(self.allocator);
    }

    /// Look up pricing for (provider, model). Returns null if no entry
    /// matches — distinct from "price is zero," which is a valid entry.
    pub fn get(self: *const Registry, provider: []const u8, model: []const u8) ?Pricing {
        for (self.entries.items) |e| {
            if (std.mem.eql(u8, e.provider, provider) and std.mem.eql(u8, e.model, model)) {
                return e.pricing;
            }
        }
        return null;
    }

    /// Insert or replace pricing for (provider, model). Duplicates the
    /// key strings.
    pub fn set(
        self: *Registry,
        provider: []const u8,
        model: []const u8,
        pricing: Pricing,
    ) !void {
        for (self.entries.items) |*e| {
            if (std.mem.eql(u8, e.provider, provider) and std.mem.eql(u8, e.model, model)) {
                e.pricing = pricing;
                return;
            }
        }
        const provider_copy = try self.allocator.dupe(u8, provider);
        errdefer self.allocator.free(provider_copy);
        const model_copy = try self.allocator.dupe(u8, model);
        errdefer self.allocator.free(model_copy);
        try self.entries.append(self.allocator, .{
            .provider = provider_copy,
            .model = model_copy,
            .pricing = pricing,
        });
    }

    pub fn count(self: *const Registry) usize {
        return self.entries.items.len;
    }
};

// =============================================================================
// Tests
// =============================================================================

const testing = std.testing;

test "Pricing.fromDollarsPerMtok: $3.00/Mtok -> 300 micro-cents/token" {
    try testing.expectEqual(@as(u64, 300), Pricing.fromDollarsPerMtok(3.0));
    try testing.expectEqual(@as(u64, 1500), Pricing.fromDollarsPerMtok(15.0));
    try testing.expectEqual(@as(u64, 30), Pricing.fromDollarsPerMtok(0.3));
    try testing.expectEqual(@as(u64, 0), Pricing.fromDollarsPerMtok(0));
    try testing.expectEqual(@as(u64, 0), Pricing.fromDollarsPerMtok(-1));
}

test "Pricing.fromDollarsPerMtok: rounds float noise to clean integer" {
    // 0.1 * 3 = 0.30000000000000004 in IEEE 754. Verify rounding
    // recovers the clean integer.
    const v = 0.1 * 3.0;
    try testing.expectEqual(@as(u64, 30), Pricing.fromDollarsPerMtok(v));
}

test "costMicroCents: standard mixed input/output" {
    const pricing: Pricing = .{
        .input = 300, // $3/Mtok
        .output = 1500, // $15/Mtok
    };
    const usage: Usage = .{ .input = 1000, .output = 200 };
    // 1000*300 + 200*1500 = 300_000 + 300_000 = 600_000 micro-cents.
    // = 6 cents = $0.06.
    try testing.expectEqual(@as(?u64, 600_000), costMicroCents(usage, pricing));
}

test "costMicroCents: cache_read and cache_write discounts are honored" {
    const pricing: Pricing = .{
        .input = 300,
        .output = 1500,
        .cache_read = 30, // 0.1x input
        .cache_write = 375, // 1.25x input
    };
    const usage: Usage = .{
        .input = 1000,
        .output = 200,
        .cache_read = 5000,
        .cache_write = 500,
    };
    // 1000*300 + 200*1500 + 5000*30 + 500*375
    //   = 300_000 + 300_000 + 150_000 + 187_500 = 937_500.
    try testing.expectEqual(@as(?u64, 937_500), costMicroCents(usage, pricing));
}

test "costMicroCents: reasoning tokens do not double-count" {
    const pricing: Pricing = .{ .input = 0, .output = 1500 };
    const usage: Usage = .{ .output = 100, .reasoning = 60 };
    // Cost is from `output` alone; reasoning is a subset of output.
    try testing.expectEqual(@as(?u64, 150_000), costMicroCents(usage, pricing));
}

test "costMicroCents: unknown price + nonzero usage poisons to null" {
    // gpt-4o written with only input/output set; cache fields default
    // to null (unknown). A turn that uses any cache reads should
    // surface as unknown cost, not silently free.
    const pricing: Pricing = .{
        .input = 250,
        .output = 1000,
        // cache_read, cache_write left null.
    };
    const usage: Usage = .{ .input = 1000, .output = 200, .cache_read = 500 };
    try testing.expectEqual(@as(?u64, null), costMicroCents(usage, pricing));
}

test "costMicroCents: unknown price + zero usage stays known" {
    // Same partially-specified pricing, but the turn never touched
    // cache. Cost should remain known.
    const pricing: Pricing = .{ .input = 250, .output = 1000 };
    const usage: Usage = .{ .input = 1000, .output = 200 };
    try testing.expectEqual(@as(?u64, 450_000), costMicroCents(usage, pricing));
}

test "Registry: set, get, replace" {
    var reg = Registry.init(testing.allocator);
    defer reg.deinit();

    try testing.expect(reg.get("anthropic", "claude-sonnet-4") == null);

    try reg.set("anthropic", "claude-sonnet-4", .{ .input = 300, .output = 1500 });
    try testing.expectEqual(@as(usize, 1), reg.count());

    const p = reg.get("anthropic", "claude-sonnet-4").?;
    try testing.expectEqual(@as(?u64, 300), p.input);
    try testing.expectEqual(@as(?u64, 1500), p.output);

    // Replace existing.
    try reg.set("anthropic", "claude-sonnet-4", .{ .input = 400, .output = 1600 });
    try testing.expectEqual(@as(usize, 1), reg.count());
    const p2 = reg.get("anthropic", "claude-sonnet-4").?;
    try testing.expectEqual(@as(?u64, 400), p2.input);
}

test "Registry: distinct (provider, model) pairs" {
    var reg = Registry.init(testing.allocator);
    defer reg.deinit();
    try reg.set("openai", "gpt-4o", .{ .input = 250 });
    try reg.set("openai", "gpt-4o-mini", .{ .input = 15 });
    try reg.set("anthropic", "claude-sonnet-4", .{ .input = 300 });
    try testing.expectEqual(@as(usize, 3), reg.count());
    try testing.expectEqual(@as(?u64, 250), reg.get("openai", "gpt-4o").?.input);
    try testing.expectEqual(@as(?u64, 15), reg.get("openai", "gpt-4o-mini").?.input);
    try testing.expectEqual(@as(?u64, 300), reg.get("anthropic", "claude-sonnet-4").?.input);
}