//! Pure provider usage query core.
//!
//! Driven by the host: no sockets, no clock, no filesystem.
//! Implements ABI 4 (limit id, windowSeconds and scope).
//! The core owns no cross-call state; the host owns it and passes it back.
const std = @import("std");
pub const iso8601 = @import("iso8601.zig");

pub const abi_version: u32 = 4;

pub const ProviderInfo = struct {
    provider: []const u8,
    providerName: []const u8,
    credentialKinds: []const []const u8,
};

pub const supported_providers = [_]ProviderInfo{
    .{
        .provider = "alibaba",
        .providerName = "Alibaba Model Studio",
        .credentialKinds = &[_][]const u8{ "cli", "session" },
    },
    .{
        .provider = "anthropic",
        .providerName = "Anthropic",
        .credentialKinds = &[_][]const u8{"oauth"},
    },
    .{
        .provider = "deepseek",
        .providerName = "DeepSeek",
        .credentialKinds = &[_][]const u8{"apiKey"},
    },
    .{
        .provider = "google-antigravity",
        .providerName = "Google Antigravity",
        .credentialKinds = &[_][]const u8{"oauth"},
    },
    .{
        .provider = "hyper",
        .providerName = "Charm Hyper",
        .credentialKinds = &[_][]const u8{ "oauth", "apiKey" },
    },
    .{
        .provider = "openai-codex",
        .providerName = "ChatGPT Codex",
        .credentialKinds = &[_][]const u8{"oauth"},
    },
    .{
        .provider = "opencode-go",
        .providerName = "OpenCode Zen",
        .credentialKinds = &[_][]const u8{ "oauth", "apiKey" },
    },
    .{
        .provider = "zai",
        .providerName = "Z.ai",
        .credentialKinds = &[_][]const u8{"apiKey"},
    },
};

/// Cross-round state owned by the host and passed back on each advance call.
/// The core owns no cross-call state; the host owns it and passes it back.
pub const State = struct {
    hyper_team_id: ?[]u8 = null,
    hyper_team_name: ?[]u8 = null,
    alloc: std.mem.Allocator,

    pub fn init(alloc: std.mem.Allocator) State {
        return .{ .alloc = alloc };
    }

    pub fn deinit(self: *State) void {
        if (self.hyper_team_id) |id| self.alloc.free(id);
        if (self.hyper_team_name) |name| self.alloc.free(name);
        self.hyper_team_id = null;
        self.hyper_team_name = null;
    }

    pub fn setHyperTeam(self: *State, team_id: []const u8, team_name: ?[]const u8) !void {
        if (self.hyper_team_id) |id| self.alloc.free(id);
        if (self.hyper_team_name) |name| self.alloc.free(name);
        self.hyper_team_id = try self.alloc.dupe(u8, team_id);
        self.hyper_team_name = if (team_name) |tn| try self.alloc.dupe(u8, tn) else null;
    }
};

pub const UsageLimit = struct {
    /// Stable machine key for this limit within its provider; `label` is display text.
    id: []const u8,
    label: []const u8,
    usedPercent: ?f64 = null,
    windowLabel: ?[]const u8 = null,
    /// Window length when the provider states it or the plan fixes it; null when unknown or
    /// calendar-shaped (a month).
    windowSeconds: ?i64 = null,
    resetsAt: ?i64 = null,
    /// The provider's own name for the model family or feature this limit is restricted to;
    /// null means it counts all usage on the account.
    scope: ?[]const u8 = null,
    status: []const u8,
    detail: ?[]const u8 = null,
};

pub const UsageRequestKind = enum {
    http,
    command,
};

pub const UsageRequest = struct {
    kind: UsageRequestKind = .http,
    // http
    method: []const u8 = "GET",
    url: []const u8 = "",
    headers: []const [2][]const u8 = &[_][2][]const u8{},
    body: ?[]const u8 = null,
    // command
    command: []const u8 = "",
    args: []const []const u8 = &[_][]const u8{},
    stdin: ?[]const u8 = null,
    timeoutMs: ?i64 = null,
};

pub const Outcome = union(enum) {
    requests: []const UsageRequest,
    report: ReportOutcome,
    err: []const u8,
};

pub const ReportOutcome = struct {
    accountId: ?[]const u8 = null,
    email: ?[]const u8 = null,
    planType: ?[]const u8 = null,
    limits: []const UsageLimit = &[_]UsageLimit{},
    balance: ?[]const u8 = null,
    notes: []const []const u8 = &[_][]const u8{},
    error_message: ?[]const u8 = null,
};

/// Append-only JSON writer over an arena.
pub const Writer = struct {
    buf: std.ArrayList(u8) = .empty,
    alloc: std.mem.Allocator,

    pub fn init(alloc: std.mem.Allocator) Writer {
        return .{ .alloc = alloc };
    }

    pub fn raw(self: *Writer, text: []const u8) !void {
        try self.buf.appendSlice(self.alloc, text);
    }

    pub fn byte(self: *Writer, c: u8) !void {
        try self.buf.append(self.alloc, c);
    }

    pub fn print(self: *Writer, comptime fmt: []const u8, args: anytype) !void {
        const text = try std.fmt.allocPrint(self.alloc, fmt, args);
        try self.buf.appendSlice(self.alloc, text);
    }

    pub fn string(self: *Writer, value: []const u8) !void {
        try writeJsonString(self, value);
    }

    pub fn items(self: *Writer) []const u8 {
        return self.buf.items;
    }
};

pub fn writeJsonString(w: *Writer, value: []const u8) !void {
    try w.byte('"');
    for (value) |c| {
        switch (c) {
            '"' => try w.raw("\\\""),
            '\\' => try w.raw("\\\\"),
            '\n' => try w.raw("\\n"),
            '\r' => try w.raw("\\r"),
            '\t' => try w.raw("\\t"),
            else => {
                if (c < 0x20) {
                    try w.print("\\u{x:0>4}", .{c});
                } else {
                    try w.byte(c);
                }
            },
        }
    }
    try w.byte('"');
}

pub fn asObject(value: ?std.json.Value) ?std.json.ObjectMap {
    const v = value orelse return null;
    return switch (v) {
        .object => |o| o,
        else => null,
    };
}

pub fn asArray(value: ?std.json.Value) ?std.json.Array {
    const v = value orelse return null;
    return switch (v) {
        .array => |a| a,
        else => null,
    };
}

pub fn asNumber(value: ?std.json.Value) ?f64 {
    const v = value orelse return null;
    const f: ?f64 = switch (v) {
        .integer => |i| @floatFromInt(i),
        .float => |fl| fl,
        .number_string => |s| std.fmt.parseFloat(f64, s) catch null,
        .string => |s| std.fmt.parseFloat(f64, s) catch null,
        else => null,
    };
    if (f) |num| {
        if (std.math.isFinite(num)) return num;
    }
    return null;
}

pub fn asString(value: ?std.json.Value) ?[]const u8 {
    const v = value orelse return null;
    return switch (v) {
        .string => |s| s,
        else => null,
    };
}

pub fn stringOf(value: ?std.json.Value) ?[]const u8 {
    const s = asString(value) orelse return null;
    const trimmed = std.mem.trim(u8, s, " \t\r\n");
    if (trimmed.len == 0) return null;
    return trimmed;
}

pub fn toInt(comptime T: type, val: ?f64) ?T {
    const v = val orelse return null;
    if (!std.math.isFinite(v)) return null;
    const min_f = @as(f64, @floatFromInt(std.math.minInt(T)));
    const max_f = @as(f64, @floatFromInt(std.math.maxInt(T)));
    if (v < min_f or v > max_f) return null;
    return @as(T, @intFromFloat(v));
}

pub fn roundToInt(comptime T: type, val: ?f64) ?T {
    const v = val orelse return null;
    if (!std.math.isFinite(v)) return null;
    return toInt(T, @round(v));
}

pub fn statusFor(used_percent: ?f64) []const u8 {
    const used = used_percent orelse return "unknown";
    if (!std.math.isFinite(used)) return "unknown";
    if (used >= 100) return "exhausted";
    if (used >= 90) return "warning";
    return "ok";
}

pub fn titleCase(alloc: std.mem.Allocator, value: []const u8) ![]const u8 {
    var parts = std.mem.splitAny(u8, value, "_-");
    var res = std.ArrayList(u8).empty;
    var first = true;
    while (parts.next()) |part| {
        if (part.len == 0) continue;
        if (!first) try res.append(alloc, ' ');
        first = false;
        try res.append(alloc, std.ascii.toUpper(part[0]));
        for (part[1..]) |c| {
            try res.append(alloc, std.ascii.toLower(c));
        }
    }
    return res.items;
}

pub fn formatWindowLabel(alloc: std.mem.Allocator, seconds: ?f64) !?[]const u8 {
    const s = seconds orelse return null;
    if (!std.math.isFinite(s) or !(s > 0)) return null;
    if (s >= 86_400) {
        const days = roundToInt(i64, s / 86_400.0) orelse return null;
        return try std.fmt.allocPrint(alloc, "{d}d window", .{days});
    }
    const hours = roundToInt(i64, s / 3600.0) orelse return null;
    return try std.fmt.allocPrint(alloc, "{d}h window", .{if (hours < 1) 1 else hours});
}

pub fn epochToWindowLabel(alloc: std.mem.Allocator, resets_at_ms: ?i64, now_ms: i64) !?[]const u8 {
    const ra = resets_at_ms orelse return null;
    // A reset already in the past floors to the 1h label rather than disappearing: the reference
    // (prime-agent's driver.ts epochToWindowLabel) does `Math.max(1, …)`, and a host that shows
    // nothing would look like a provider with no window at all.
    const diff = ra - now_ms;
    const diff_f = @as(f64, @floatFromInt(diff)) / 1000.0;
    const seconds = @max(1, roundToInt(i64, diff_f) orelse return null);
    if (seconds >= 86_400) {
        const days = roundToInt(i64, @as(f64, @floatFromInt(seconds)) / 86_400.0) orelse return null;
        return try std.fmt.allocPrint(alloc, "{d}d window", .{days});
    }
    const hours = @max(1, roundToInt(i64, @as(f64, @floatFromInt(seconds)) / 3600.0) orelse return null);
    return try std.fmt.allocPrint(alloc, "{d}h window", .{hours});
}

pub fn httpStatusError(alloc: std.mem.Allocator, response_obj: ?std.json.ObjectMap) ![]const u8 {
    const ro = response_obj orelse return "HTTP no response";
    const status_num = asNumber(ro.get("status"));
    const status = toInt(i64, status_num) orelse return "HTTP no response";
    const body_str = asString(ro.get("body"));
    if (body_str) |bs| {
        if (std.json.parseFromSlice(std.json.Value, alloc, bs, .{})) |parsed| {
            if (asObject(parsed.value)) |root| {
                if (root.get("error")) |err_val| {
                    if (asString(err_val)) |err_str| {
                        const trimmed = std.mem.trim(u8, err_str, " \t\r\n");
                        if (trimmed.len > 0) {
                            return try std.fmt.allocPrint(alloc, "HTTP {d}: {s}", .{ status, trimmed });
                        }
                    } else if (asObject(err_val)) |err_obj| {
                        if (asString(err_obj.get("message"))) |msg| {
                            const trimmed = std.mem.trim(u8, msg, " \t\r\n");
                            if (trimmed.len > 0) {
                                return try std.fmt.allocPrint(alloc, "HTTP {d}: {s}", .{ status, trimmed });
                            }
                        }
                    }
                }
            }
        } else |_| {}
    }
    return try std.fmt.allocPrint(alloc, "HTTP {d}", .{status});
}

pub const FirstResponse = struct {
    obj: std.json.ObjectMap,
    parsed: std.json.Parsed(std.json.Value),
};

pub const FirstResponseOutcome = union(enum) {
    ok: FirstResponse,
    err: []const u8,
};

pub fn parseFirstResponse(
    alloc: std.mem.Allocator,
    resp_array: std.json.Array,
    not_json_msg: []const u8,
) !FirstResponseOutcome {
    return parseFirstResponseEx(alloc, resp_array, not_json_msg, false);
}

pub fn parseFirstResponseEx(
    alloc: std.mem.Allocator,
    resp_array: std.json.Array,
    not_json_msg: []const u8,
    unwrap_http_error: bool,
) !FirstResponseOutcome {
    if (resp_array.items.len == 0) return .{ .err = "HTTP no response" };
    const resp0_obj = asObject(resp_array.items[0]) orelse return .{ .err = "HTTP no response" };
    const status_num = asNumber(resp0_obj.get("status"));
    const status = toInt(i64, status_num);
    if (status == null or status.? != 200) {
        if (unwrap_http_error) {
            return .{ .err = try httpStatusError(alloc, resp0_obj) };
        }
        const err_text = if (status) |st|
            try std.fmt.allocPrint(alloc, "HTTP {d}", .{st})
        else
            "HTTP no response";
        return .{ .err = err_text };
    }
    const body_str = asString(resp0_obj.get("body")) orelse return .{ .err = not_json_msg };
    const parsed = std.json.parseFromSlice(std.json.Value, alloc, body_str, .{}) catch {
        return .{ .err = not_json_msg };
    };
    return .{ .ok = .{ .obj = resp0_obj, .parsed = parsed } };
}

pub fn urlEncode(alloc: std.mem.Allocator, input: []const u8) ![]const u8 {
    var out = std.ArrayList(u8).empty;
    const hex = "0123456789ABCDEF";
    for (input) |c| {
        switch (c) {
            'a'...'z', 'A'...'Z', '0'...'9', '-', '_', '.', '*' => try out.append(alloc, c),
            ' ' => try out.append(alloc, '+'),
            else => {
                try out.append(alloc, '%');
                try out.append(alloc, hex[(c >> 4) & 0x0F]);
                try out.append(alloc, hex[c & 0x0F]);
            },
        }
    }
    return out.items;
}

pub fn writeStep(w: *Writer, outcome: Outcome) !void {
    switch (outcome) {
        .requests => |reqs| {
            try w.raw("{\"requests\":[");
            for (reqs, 0..) |req, i| {
                if (i > 0) try w.byte(',');
                switch (req.kind) {
                    .http => {
                        try w.raw("{\"kind\":\"http\",\"method\":");
                        try w.string(req.method);
                        try w.raw(",\"url\":");
                        try w.string(req.url);
                        try w.raw(",\"headers\":[");
                        for (req.headers, 0..) |h, hi| {
                            if (hi > 0) try w.byte(',');
                            try w.raw("[");
                            try w.string(h[0]);
                            try w.byte(',');
                            try w.string(h[1]);
                            try w.raw("]");
                        }
                        try w.raw("]");
                        if (req.body) |body| {
                            try w.raw(",\"body\":");
                            try w.string(body);
                        }
                        try w.byte('}');
                    },
                    .command => {
                        try w.raw("{\"kind\":\"command\",\"command\":");
                        try w.string(req.command);
                        try w.raw(",\"args\":[");
                        for (req.args, 0..) |arg, ai| {
                            if (ai > 0) try w.byte(',');
                            try w.string(arg);
                        }
                        try w.byte(']');
                        if (req.stdin) |st| {
                            try w.raw(",\"stdin\":");
                            try w.string(st);
                        }
                        if (req.timeoutMs) |to| {
                            try w.print(",\"timeoutMs\":{d}", .{to});
                        }
                        try w.byte('}');
                    },
                }
            }
            try w.raw("]}");
        },
        .report => |rep| {
            try w.raw("{\"report\":{");
            var wrote = false;
            if (rep.accountId) |aid| {
                try w.raw("\"accountId\":");
                try w.string(aid);
                wrote = true;
            }
            if (rep.email) |em| {
                if (wrote) try w.byte(',');
                try w.raw("\"email\":");
                try w.string(em);
                wrote = true;
            }
            if (rep.planType) |pt| {
                if (wrote) try w.byte(',');
                try w.raw("\"planType\":");
                try w.string(pt);
                wrote = true;
            }

            if (wrote) try w.byte(',');
            try w.raw("\"limits\":[");
            for (rep.limits, 0..) |lim, i| {
                if (i > 0) try w.byte(',');
                try w.raw("{\"id\":");
                try w.string(lim.id);
                try w.raw(",\"label\":");
                try w.string(lim.label);
                if (lim.usedPercent) |up| {
                    if (std.math.isFinite(up)) {
                        const clamped = @max(0.0, @min(100.0, up));
                        try w.print(",\"usedPercent\":{d}", .{clamped});
                    }
                }
                if (lim.windowLabel) |wl| {
                    try w.raw(",\"windowLabel\":");
                    try w.string(wl);
                }
                if (lim.windowSeconds) |ws| {
                    try w.print(",\"windowSeconds\":{d}", .{ws});
                }
                if (lim.resetsAt) |ra| {
                    try w.print(",\"resetsAt\":{d}", .{ra});
                }
                if (lim.scope) |sc| {
                    try w.raw(",\"scope\":");
                    try w.string(sc);
                }
                if (lim.detail) |dt| {
                    try w.raw(",\"detail\":");
                    try w.string(dt);
                }
                try w.raw(",\"status\":");
                try w.string(lim.status);
                try w.byte('}');
            }
            try w.byte(']');
            wrote = true;

            if (rep.balance) |bal| {
                try w.raw(",\"balance\":");
                try w.string(bal);
            }
            if (rep.notes.len > 0) {
                try w.raw(",\"notes\":[");
                for (rep.notes, 0..) |note, ni| {
                    if (ni > 0) try w.byte(',');
                    try w.string(note);
                }
                try w.byte(']');
            }

            try w.raw(",\"error\":");
            if (rep.error_message) |err_msg| {
                try w.string(err_msg);
            } else {
                try w.raw("null");
            }
            try w.raw("}}");
        },
        .err => |msg| {
            try w.raw("{\"error\":");
            try w.string(msg);
            try w.raw("}");
        },
    }
}

pub fn writeProviders(w: *Writer) !void {
    try w.raw("[");
    for (supported_providers, 0..) |p, i| {
        if (i > 0) try w.byte(',');
        try w.raw("{\"provider\":");
        try w.string(p.provider);
        try w.raw(",\"providerName\":");
        try w.string(p.providerName);
        try w.raw(",\"credentialKinds\":[");
        for (p.credentialKinds, 0..) |k, ki| {
            if (ki > 0) try w.byte(',');
            try w.string(k);
        }
        try w.raw("]}");
    }
    try w.raw("]");
}

// Re-export providers
pub const codex = @import("providers/codex.zig");
pub const anthropic = @import("providers/anthropic.zig");
pub const antigravity = @import("providers/antigravity.zig");
pub const hyper = @import("providers/hyper.zig");
pub const opencode_go = @import("providers/opencode_go.zig");
pub const zai = @import("providers/zai.zig");
pub const deepseek = @import("providers/deepseek.zig");
pub const alibaba = @import("providers/alibaba.zig");

const ProviderDispatch = struct {
    name: []const u8,
    advance: *const fn (std.mem.Allocator, ?*State, []const u8, []const u8, i64) anyerror!Outcome,
};

fn wrapStateless(comptime func: anytype) *const fn (std.mem.Allocator, ?*State, []const u8, []const u8, i64) anyerror!Outcome {
    return struct {
        fn call(alloc: std.mem.Allocator, state: ?*State, cred: []const u8, resp: []const u8, now_ms: i64) anyerror!Outcome {
            _ = state;
            return func(alloc, cred, resp, now_ms);
        }
    }.call;
}

const provider_table = [_]ProviderDispatch{
    .{ .name = "alibaba", .advance = wrapStateless(alibaba.advance) },
    .{ .name = "anthropic", .advance = wrapStateless(anthropic.advance) },
    .{ .name = "deepseek", .advance = wrapStateless(deepseek.advance) },
    .{ .name = "google-antigravity", .advance = wrapStateless(antigravity.advance) },
    .{ .name = "hyper", .advance = hyper.advance },
    .{ .name = "openai-codex", .advance = wrapStateless(codex.advance) },
    .{ .name = "opencode-go", .advance = wrapStateless(opencode_go.advance) },
    .{ .name = "zai", .advance = wrapStateless(zai.advance) },
};

pub fn advance(
    alloc: std.mem.Allocator,
    state: ?*State,
    provider: []const u8,
    credential_json: []const u8,
    responses_json: []const u8,
    now_ms: i64,
) !Outcome {
    inline for (provider_table) |entry| {
        if (std.mem.eql(u8, provider, entry.name)) {
            return entry.advance(alloc, state, credential_json, responses_json, now_ms);
        }
    }
    return .{ .err = "unknown provider" };
}

pub const ClaimsOutcome = union(enum) { ok, err: []const u8 };

pub fn claimsFromAccessToken(alloc: std.mem.Allocator, jwt: []const u8, w: *Writer) !ClaimsOutcome {
    const payload = jwtPayload(alloc, jwt) catch return .{ .err = "access token is not a JWT" };
    defer alloc.free(payload);

    var parsed = std.json.parseFromSlice(std.json.Value, alloc, payload, .{}) catch {
        return .{ .err = "access token claims are not JSON" };
    };
    defer parsed.deinit();

    const root = asObject(parsed.value) orelse return .{ .err = "access token claims are not an object" };
    const auth = asObject(root.get("https://api.openai.com/auth"));
    const profile = asObject(root.get("https://api.openai.com/profile"));

    const account_id = if (auth) |a| asString(a.get("chatgpt_account_id")) else null;
    const email_claim = if (profile) |p| asString(p.get("email")) else null;
    const email: ?[]const u8 = if (email_claim) |raw| blk: {
        const trimmed = std.mem.trim(u8, raw, " \t\r\n");
        if (trimmed.len == 0) break :blk null;
        const lowered = try alloc.alloc(u8, trimmed.len);
        break :blk std.ascii.lowerString(lowered, trimmed);
    } else null;

    try w.raw("{");
    var wrote = false;
    if (account_id) |id| {
        try w.raw("\"accountId\":");
        try w.string(id);
        wrote = true;
    }
    if (email) |value| {
        if (wrote) try w.byte(',');
        try w.raw("\"email\":");
        try w.string(value);
    }
    try w.byte('}');
    return .ok;
}

pub fn jwtPayload(alloc: std.mem.Allocator, jwt: []const u8) ![]u8 {
    var parts = std.mem.splitScalar(u8, jwt, '.');
    _ = parts.next() orelse return error.NotAJwt;
    const encoded = parts.next() orelse return error.NotAJwt;
    _ = parts.next() orelse return error.NotAJwt;

    const decoder = std.base64.url_safe_no_pad.Decoder;
    const size = try decoder.calcSizeForSlice(encoded);
    const decoded = try alloc.alloc(u8, size);
    errdefer alloc.free(decoded);
    try decoder.decode(decoded, encoded);
    return decoded;
}