//! Loader for `models.toml` — model aliases and per-model pricing.
//!
//! A model entry is keyed by `<provider_name>.<alias>`, where
//! `<provider_name>` matches a provider defined in `config.toml`'s
//! `[providers.<name>]` table, and `<alias>` is the short name the user
//! references via `<provider_name>:<alias>` (e.g. `anthropic:sonnet`).
//!
//! Schema:
//!
//!     [<provider>.<alias>]
//!     model       = <string>  # wire model id; defaults to <alias> if omitted
//!     reasoning   = <string>  # default | off | minimal | low | medium | high
//!     max_tokens  = <int>     # per-request output token cap
//!     api_version = <string>  # anthropic_messages only
//!     # pricing (all optional; USD per million tokens):
//!     input       = <float>
//!     output      = <float>
//!     cache_read  = <float>
//!     cache_write = <float>
//!
//! Pricing semantics are unchanged from the original models.toml: any
//! omitted price field is "unknown" (null), not zero. Write `= 0`
//! explicitly to declare a known-zero rate.
//!
//! Example:
//!
//!     [anthropic.sonnet]
//!     model       = "claude-sonnet-4-20250514"
//!     reasoning   = "high"
//!     max_tokens  = 8192
//!     input       = 3.0
//!     output      = 15.0
//!     cache_read  = 0.3
//!     cache_write = 3.75
//!
//!     [openai.gpt]
//!     model  = "gpt-4o"
//!     input  = 2.5
//!     output = 10.0
//!
//! A referenced alias with no entry here is not an error: the alias is
//! used verbatim as the wire model name (zero-config convenience), with
//! default knobs and unknown pricing.

const std = @import("std");
const Allocator = std.mem.Allocator;
const Io = std.Io;

const toml = @import("toml");
const panto = @import("panto");

pub const Pricing = panto.Pricing;
pub const Registry = panto.PricingRegistry;
pub const ReasoningEffort = panto.ReasoningEffort;

/// A resolved model definition: the wire model id plus per-model knobs.
/// Strings are owned by the `ModelRegistry`.
pub const ModelDef = struct {
    /// `<provider_name>` — matches a `[providers.<name>]` in config.toml.
    provider: []const u8,
    /// `<alias>` — the short reference name.
    alias: []const u8,
    /// Wire model id sent to the provider API. Equals `alias` when the
    /// TOML omitted `model`.
    model: []const u8,
    reasoning: ReasoningEffort,
    /// Per-request output token cap; null = use the provider/library default.
    max_tokens: ?u32,
    /// anthropic_messages only; null = use the library default.
    api_version: ?[]const u8,
};

/// In-memory map of `(provider, alias) -> ModelDef`. Owns all strings.
pub const ModelRegistry = struct {
    allocator: Allocator,
    entries: std.ArrayList(ModelDef),

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

    pub fn deinit(self: *ModelRegistry) void {
        for (self.entries.items) |e| {
            self.allocator.free(e.provider);
            self.allocator.free(e.alias);
            self.allocator.free(e.model);
            if (e.api_version) |v| self.allocator.free(v);
        }
        self.entries.deinit(self.allocator);
    }

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

    /// Look up a model definition by provider name + alias.
    pub fn get(self: *const ModelRegistry, provider: []const u8, alias: []const u8) ?ModelDef {
        for (self.entries.items) |e| {
            if (std.mem.eql(u8, e.provider, provider) and std.mem.eql(u8, e.alias, alias)) {
                return e;
            }
        }
        return null;
    }
};

/// Everything `models.toml` yields: model definitions and pricing. The
/// two registries are independent (pricing keys on the *wire* model id;
/// model defs key on the alias) but parsed in one pass.
pub const Models = struct {
    defs: ModelRegistry,
    pricing: Registry,

    pub fn deinit(self: *Models) void {
        self.defs.deinit();
        self.pricing.deinit();
    }
};

/// Resolve the absolute path to `models.toml`. Honors `XDG_CONFIG_HOME`,
/// falling back to `$HOME/.config`. Caller owns the returned slice.
pub fn configPath(allocator: Allocator, environ_map: *const std.process.Environ.Map) ![]u8 {
    if (environ_map.get("XDG_CONFIG_HOME")) |xdg| {
        return try std.fs.path.join(allocator, &.{ xdg, "panto", "models.toml" });
    }
    if (environ_map.get("HOME")) |home| {
        return try std.fs.path.join(allocator, &.{ home, ".config", "panto", "models.toml" });
    }
    return error.NoHomeDirectory;
}

/// Load `models.toml` from `path`. A missing file yields empty registries
/// (no error). Parse errors return `error.InvalidModelsToml`.
pub fn loadFromPath(allocator: Allocator, io: Io, path: []const u8) !Models {
    var models: Models = .{
        .defs = ModelRegistry.init(allocator),
        .pricing = Registry.init(allocator),
    };
    errdefer models.deinit();

    const file = Io.Dir.cwd().openFile(io, path, .{ .mode = .read_only }) catch |err| switch (err) {
        error.FileNotFound => return models, // empty; perfectly fine.
        else => return err,
    };
    defer file.close(io);

    const file_len = try file.length(io);
    const bytes = try allocator.alloc(u8, @intCast(file_len));
    defer allocator.free(bytes);
    _ = try file.readPositionalAll(io, bytes, 0);

    try parseInto(&models, bytes);
    return models;
}

/// Parse a TOML string into the given registries. Useful for tests.
pub fn parseInto(models: *Models, source: []const u8) !void {
    const alloc = models.defs.allocator;
    const result = toml.parseWithError(alloc, source, .{});
    switch (result) {
        .err => |e| {
            if (!@import("builtin").is_test) {
                std.log.err(
                    "models.toml: parse error at line {d}, column {d}: {s}",
                    .{ e.line, e.column, e.message },
                );
            }
            return error.InvalidModelsToml;
        },
        .ok => |doc| {
            defer {
                var d = doc;
                d.deinit();
            }
            try ingestDocument(models, doc);
        },
    }
}

fn ingestDocument(models: *Models, doc: *toml.Document) !void {
    const root_val: *const toml.Value = doc.root;
    if (root_val.* != .table) return;

    var provider_it = toml.tableIterator(root_val);
    while (provider_it.next()) |provider_entry| {
        const provider = provider_entry.key;
        const provider_val: *const toml.Value = provider_entry.value;
        if (provider_val.* != .table) continue;

        var alias_it = toml.tableIterator(provider_val);
        while (alias_it.next()) |alias_entry| {
            const alias = alias_entry.key;
            const v: *const toml.Value = alias_entry.value;
            if (v.* != .table) continue;

            try ingestModel(models, provider, alias, v);
        }
    }
}

fn ingestModel(
    models: *Models,
    provider: []const u8,
    alias: []const u8,
    v: *const toml.Value,
) !void {
    const alloc = models.defs.allocator;

    // Wire model id: explicit `model`, else the alias itself.
    const wire_model: []const u8 = blk: {
        if (v.get("model")) |m| {
            if (m.asString()) |s| break :blk s;
        }
        break :blk alias;
    };

    const reasoning: ReasoningEffort = blk: {
        if (v.get("reasoning")) |r| {
            if (r.asString()) |s| {
                break :blk std.meta.stringToEnum(ReasoningEffort, s) orelse .default;
            }
        }
        break :blk .default;
    };

    const max_tokens: ?u32 = blk: {
        if (v.get("max_tokens")) |mt| {
            if (mt.asI64()) |i| {
                if (i > 0) break :blk @intCast(i);
            }
        }
        break :blk null;
    };

    const api_version_src: ?[]const u8 = blk: {
        if (v.get("api_version")) |av| {
            if (av.asString()) |s| break :blk s;
        }
        break :blk null;
    };

    // Own all the strings.
    const provider_copy = try alloc.dupe(u8, provider);
    errdefer alloc.free(provider_copy);
    const alias_copy = try alloc.dupe(u8, alias);
    errdefer alloc.free(alias_copy);
    const model_copy = try alloc.dupe(u8, wire_model);
    errdefer alloc.free(model_copy);
    const api_version_copy: ?[]const u8 = if (api_version_src) |s| try alloc.dupe(u8, s) else null;
    errdefer if (api_version_copy) |s| alloc.free(s);

    try models.defs.entries.append(alloc, .{
        .provider = provider_copy,
        .alias = alias_copy,
        .model = model_copy,
        .reasoning = reasoning,
        .max_tokens = max_tokens,
        .api_version = api_version_copy,
    });

    // Pricing keys on the *wire* model id so usage records (which carry
    // the wire model) resolve directly.
    const pricing = pricingFromValue(v);
    try models.pricing.set(provider, wire_model, pricing);
}

fn pricingFromValue(v: *const toml.Value) Pricing {
    return .{
        .input = readPriceField(v, "input"),
        .output = readPriceField(v, "output"),
        .cache_read = readPriceField(v, "cache_read"),
        .cache_write = readPriceField(v, "cache_write"),
    };
}

/// Returns `null` if the field is absent or has a type we can't
/// interpret as a price. An explicit `0` comes through as a known zero.
fn readPriceField(table: *const toml.Value, name: []const u8) ?u64 {
    const field = table.get(name) orelse return null;
    if (field.asF64()) |f| return Pricing.fromDollarsPerMtok(f);
    if (field.asI64()) |i| return Pricing.fromDollarsPerMtok(@floatFromInt(i));
    return null;
}

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

const testing = std.testing;

fn emptyModels(a: Allocator) Models {
    return .{ .defs = ModelRegistry.init(a), .pricing = Registry.init(a) };
}

test "parseInto: model defs carry wire name and knobs, keyed by provider.alias" {
    var models = emptyModels(testing.allocator);
    defer models.deinit();

    const src =
        \\[anthropic.sonnet]
        \\model = "claude-sonnet-4-20250514"
        \\reasoning = "high"
        \\max_tokens = 8192
        \\input = 3.0
        \\output = 15.0
        \\
        \\[openai.gpt]
        \\model = "gpt-4o"
        \\input = 2.5
        \\output = 10.0
    ;
    try parseInto(&models, src);

    const sonnet = models.defs.get("anthropic", "sonnet").?;
    try testing.expectEqualStrings("claude-sonnet-4-20250514", sonnet.model);
    try testing.expectEqual(ReasoningEffort.high, sonnet.reasoning);
    try testing.expectEqual(@as(?u32, 8192), sonnet.max_tokens);

    const gpt = models.defs.get("openai", "gpt").?;
    try testing.expectEqualStrings("gpt-4o", gpt.model);
    try testing.expectEqual(ReasoningEffort.default, gpt.reasoning);
    try testing.expectEqual(@as(?u32, null), gpt.max_tokens);

    // Pricing keyed on the wire model id.
    const sonnet_price = models.pricing.get("anthropic", "claude-sonnet-4-20250514").?;
    try testing.expectEqual(@as(?u64, 300), sonnet_price.input);
    try testing.expectEqual(@as(?u64, 1500), sonnet_price.output);
}

test "parseInto: omitted model defaults the wire name to the alias" {
    var models = emptyModels(testing.allocator);
    defer models.deinit();

    const src =
        \\[openai.gpt-4o]
        \\input = 2.5
        \\output = 10.0
    ;
    try parseInto(&models, src);

    const def = models.defs.get("openai", "gpt-4o").?;
    try testing.expectEqualStrings("gpt-4o", def.model);
}

test "parseInto: pricing omissions stay null (unknown, not zero)" {
    var models = emptyModels(testing.allocator);
    defer models.deinit();

    const src =
        \\[openai.gpt]
        \\model = "gpt-4o"
        \\input = 2.5
        \\output = 10.0
    ;
    try parseInto(&models, src);
    const p = models.pricing.get("openai", "gpt-4o").?;
    try testing.expectEqual(@as(?u64, null), p.cache_read);
    try testing.expectEqual(@as(?u64, null), p.cache_write);
}

test "parseInto: explicit cache_write = 0 is a known zero" {
    var models = emptyModels(testing.allocator);
    defer models.deinit();

    const src =
        \\[openai.gpt]
        \\model = "gpt-4o"
        \\input = 2.5
        \\output = 10.0
        \\cache_write = 0
    ;
    try parseInto(&models, src);
    const p = models.pricing.get("openai", "gpt-4o").?;
    try testing.expectEqual(@as(?u64, 0), p.cache_write);
}

test "parseInto: malformed TOML returns InvalidModelsToml" {
    var models = emptyModels(testing.allocator);
    defer models.deinit();
    try testing.expectError(error.InvalidModelsToml, parseInto(&models, "not valid = ="));
}

test "loadFromPath: missing file returns empty registries, no error" {
    const io = testing.io;
    var models = try loadFromPath(testing.allocator, io, "/nonexistent/models.toml");
    defer models.deinit();
    try testing.expectEqual(@as(usize, 0), models.defs.count());
    try testing.expectEqual(@as(usize, 0), models.pricing.count());
}

test "ModelRegistry.get: returns null for unknown alias" {
    var models = emptyModels(testing.allocator);
    defer models.deinit();
    const src =
        \\[anthropic.sonnet]
        \\model = "claude-sonnet-4-20250514"
    ;
    try parseInto(&models, src);
    try testing.expect(models.defs.get("anthropic", "opus") == null);
    try testing.expect(models.defs.get("openai", "sonnet") == null);
}

test "configPath: XDG_CONFIG_HOME wins" {
    const a = testing.allocator;
    var env: std.process.Environ.Map = .init(a);
    defer env.deinit();
    try env.put("XDG_CONFIG_HOME", "/custom/cfg");
    try env.put("HOME", "/ignored");
    const got = try configPath(a, &env);
    defer a.free(got);
    try testing.expectEqualStrings("/custom/cfg/panto/models.toml", got);
}

test "configPath: falls back to HOME/.config" {
    const a = testing.allocator;
    var env: std.process.Environ.Map = .init(a);
    defer env.deinit();
    try env.put("HOME", "/home/user");
    const got = try configPath(a, &env);
    defer a.free(got);
    try testing.expectEqualStrings("/home/user/.config/panto/models.toml", got);
}