//! OpenAI Codex usage feed implementation.
//!
//! Mirrors prime-agent:packages/ai/src/utils/usage/providers/codex.ts field for field.
const std = @import("std");
const usage = @import("../usage.zig");

pub const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";

fn parseWindow(alloc: std.mem.Allocator, id: []const u8, window_val: ?std.json.Value, now_ms: i64, scope: ?[]const u8) !?usage.UsageLimit {
    const value = usage.asObject(window_val) orelse return null;
    const used_percent = usage.asNumber(value.get("used_percent"));
    const reset_at = usage.asNumber(value.get("reset_at"));
    const reset_after = usage.asNumber(value.get("reset_after_seconds"));
    const limit_window = usage.asNumber(value.get("limit_window_seconds"));

    const reset_ms: ?i64 = if (reset_at) |ra|
        (if (ra > 1_000_000_000_000.0) usage.toInt(i64, ra) else usage.toInt(i64, ra * 1000.0))
    else if (reset_after) |after| blk: {
        const ms = usage.toInt(i64, after * 1000.0) orelse break :blk null;
        break :blk now_ms + ms;
    } else
        null;

    const window_lbl = try usage.formatWindowLabel(alloc, limit_window);
    const window_seconds: ?i64 = if (limit_window) |lw| (if (lw > 0) usage.roundToInt(i64, lw) else null) else null;

    return usage.UsageLimit{
        .id = id,
        .label = "Usage",
        .usedPercent = used_percent,
        .windowLabel = window_lbl,
        .windowSeconds = window_seconds,
        .resetsAt = reset_ms,
        .scope = scope,
        .status = usage.statusFor(used_percent),
    };
}

pub const ParseResult = struct {
    planType: ?[]const u8 = null,
    limits: []const usage.UsageLimit = &[_]usage.UsageLimit{},
};

pub fn parseCodexUsage(alloc: std.mem.Allocator, payload: ?std.json.Value, now_ms: i64) !ParseResult {
    const root = usage.asObject(payload);
    const rate_limit = if (root) |r| usage.asObject(r.get("rate_limit")) else null;
    var limits = std.ArrayList(usage.UsageLimit).empty;

    const primary_val = if (rate_limit) |rl| rl.get("primary_window") else null;
    if (try parseWindow(alloc, "primary", primary_val, now_ms, null)) |pw| {
        var w = pw;
        w.label = "Primary";
        try limits.append(alloc, w);
    }

    const secondary_val = if (rate_limit) |rl| rl.get("secondary_window") else null;
    if (try parseWindow(alloc, "secondary", secondary_val, now_ms, null)) |sw| {
        var w = sw;
        w.label = "Secondary";
        try limits.append(alloc, w);
    }

    const additional = if (root) |r| usage.asArray(r.get("additional_rate_limits")) else null;
    if (additional) |arr| {
        for (arr.items) |entry| {
            const item = usage.asObject(entry) orelse continue;
            const extra = usage.asObject(item.get("rate_limit"));
            const limit_name = usage.asString(item.get("limit_name"));
            const feature = usage.asString(item.get("metered_feature"));
            const name = limit_name orelse feature orelse "Additional";
            // The metered feature is the machine key; the limit name is what it restricts.
            const key = feature orelse limit_name orelse "additional";
            const scope = limit_name orelse feature;

            const extra_primary = if (extra) |e| e.get("primary_window") else null;
            if (try parseWindow(alloc, try std.fmt.allocPrint(alloc, "{s}.primary", .{key}), extra_primary, now_ms, scope)) |ep| {
                var w = ep;
                w.label = try std.fmt.allocPrint(alloc, "{s} · Primary", .{name});
                try limits.append(alloc, w);
            }

            const extra_secondary = if (extra) |e| e.get("secondary_window") else null;
            if (try parseWindow(alloc, try std.fmt.allocPrint(alloc, "{s}.secondary", .{key}), extra_secondary, now_ms, scope)) |es| {
                var w = es;
                w.label = try std.fmt.allocPrint(alloc, "{s} · Secondary", .{name});
                try limits.append(alloc, w);
            }
        }
    }

    const plan_type = if (root) |r| usage.asString(r.get("plan_type")) else null;
    return ParseResult{
        .planType = plan_type,
        .limits = limits.items,
    };
}

fn jwtClaim(alloc: std.mem.Allocator, token: []const u8, namespace: []const u8, key: []const u8) ?[]const u8 {
    const payload = usage.jwtPayload(alloc, token) catch return null;
    const parsed = std.json.parseFromSlice(std.json.Value, alloc, payload, .{}) catch return null;
    const root = usage.asObject(parsed.value) orelse return null;
    const ns = usage.asObject(root.get(namespace)) orelse return null;
    return usage.asString(ns.get(key));
}

fn jwtClaimEmail(alloc: std.mem.Allocator, token: []const u8) ?[]const u8 {
    const email_raw = jwtClaim(alloc, token, "https://api.openai.com/profile", "email") orelse return null;
    const trimmed = std.mem.trim(u8, email_raw, " \t\r\n");
    if (trimmed.len == 0) return null;
    const lowered = alloc.alloc(u8, trimmed.len) catch return null;
    return std.ascii.lowerString(lowered, trimmed);
}

pub fn advance(
    alloc: std.mem.Allocator,
    credential_json: []const u8,
    responses_json: []const u8,
    now_ms: i64,
) !usage.Outcome {
    const cred_parsed = std.json.parseFromSlice(std.json.Value, alloc, credential_json, .{}) catch {
        return .{ .report = .{ .error_message = "invalid credential JSON" } };
    };
    const cred_root = usage.asObject(cred_parsed.value) orelse {
        return .{ .report = .{ .error_message = "invalid credential JSON" } };
    };
    const kind = usage.asString(cred_root.get("kind")) orelse "";
    if (!std.mem.eql(u8, kind, "oauth")) {
        return .{ .report = .{ .error_message = "openai-codex needs an OAuth account" } };
    }

    const access = usage.asString(cred_root.get("access")) orelse "";
    const cred_account_id = usage.asString(cred_root.get("accountId"));
    const account_id = if (cred_account_id != null and cred_account_id.?.len > 0)
        cred_account_id
    else
        jwtClaim(alloc, access, "https://api.openai.com/auth", "chatgpt_account_id");

    const cred_email = usage.asString(cred_root.get("email"));
    const email = if (cred_email != null and cred_email.?.len > 0)
        cred_email
    else
        jwtClaimEmail(alloc, access);

    const resp_parsed = std.json.parseFromSlice(std.json.Value, alloc, responses_json, .{}) catch {
        return .{ .report = .{ .error_message = "invalid responses JSON" } };
    };
    const resp_array = usage.asArray(resp_parsed.value) orelse {
        return .{ .report = .{ .error_message = "invalid responses JSON" } };
    };

    if (resp_array.items.len == 0) {
        var headers = std.ArrayList([2][]const u8).empty;
        const auth_hdr = try std.fmt.allocPrint(alloc, "Bearer {s}", .{access});
        try headers.append(alloc, .{ "Authorization", auth_hdr });
        try headers.append(alloc, .{ "User-Agent", "prime-agent" });
        if (account_id) |aid| {
            if (aid.len > 0) {
                try headers.append(alloc, .{ "ChatGPT-Account-Id", aid });
            }
        }
        const requests = try alloc.alloc(usage.UsageRequest, 1);
        requests[0] = .{
            .method = "GET",
            .url = CODEX_USAGE_URL,
            .headers = headers.items,
        };
        return .{ .requests = requests };
    }

    const first = switch (try usage.parseFirstResponse(alloc, resp_array, "usage body is not JSON")) {
        .ok => |f| f,
        .err => |msg| return .{ .report = .{ .error_message = msg } },
    };

    const parsed_usage = try parseCodexUsage(alloc, first.parsed.value, now_ms);

    return .{
        .report = .{
            .planType = parsed_usage.planType,
            .limits = parsed_usage.limits,
            .accountId = account_id,
            .email = email,
            .error_message = null,
        },
    };
}