const std = @import("std");
const Config = @import("../config.zig").Config;
const ws = @import("../net/ws/client.zig");
const crypto = @import("../utils/crypto.zig");
const io = @import("../utils/io.zig");
const core = @import("../core.zig");

pub fn run(allocator: std.mem.Allocator, args: []const []const u8) !void {
    var flags = core.flags.CommonFlags{};
    var keep_count: ?u32 = null;
    var older_than_days: ?u32 = null;

    // Parse flags
    var i: usize = 0;
    while (i < args.len) : (i += 1) {
        if (std.mem.eql(u8, args[i], "--help") or std.mem.eql(u8, args[i], "-h")) {
            return printUsage();
        } else if (std.mem.eql(u8, args[i], "--json")) {
            flags.json = true;
        } else if (std.mem.eql(u8, args[i], "--keep") and i + 1 < args.len) {
            keep_count = try std.fmt.parseInt(u32, args[i + 1], 10);
            i += 1;
        } else if (std.mem.eql(u8, args[i], "--older-than") and i + 1 < args.len) {
            older_than_days = try std.fmt.parseInt(u32, args[i + 1], 10);
            i += 1;
        }
    }

    core.output.setMode(if (flags.flags.json) .flags.json else .text);

    if (keep_count == null and older_than_days == null) {
        core.output.usage("prune", "ml prune --keep <n> | --older-than <days>");
        return error.InvalidArgs;
    }

    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

    // Add confirmation prompt
    if (!flags.flags.json) {
        if (keep_count) |count| {
            if (!io.confirm("This will permanently delete all but the {d} most recent experiments. Continue?", .{count})) {
                io.info("Prune cancelled.\n", .{});
                return;
            }
        } else if (older_than_days) |days| {
            if (!io.confirm("This will permanently delete all experiments older than {d} days. Continue?", .{days})) {
                io.info("Prune cancelled.\n", .{});
                return;
            }
        }
    }

    io.info("Pruning experiments...\n", .{});

    // Use plain password for WebSocket authentication, hash for binary protocol
    const api_key_plain = config.api_key; // Plain password from config
    const api_key_hash = try crypto.hashApiKey(allocator, api_key_plain);
    defer allocator.free(api_key_hash);

    // Connect to WebSocket and send prune message
    const ws_url = try config.getWebSocketUrl(allocator);
    defer allocator.free(ws_url);

    var client = try ws.Client.connect(allocator, ws_url, api_key_plain);
    defer client.close();

    // Determine prune type and send message
    var prune_type: u8 = undefined;
    var value: u32 = undefined;

    if (keep_count) |count| {
        prune_type = 0; // keep N
        value = count;
        io.info("Keeping {d} most recent experiments\n", .{count});
    }
    if (older_than_days) |days| {
        prune_type = 1; // older than days
        value = days;
        io.info("Removing experiments older than {d} days\n", .{days});
    }

    try client.sendPrune(api_key_hash, prune_type, value);

    // Receive response
    const response = try client.receiveMessage(allocator);
    defer allocator.free(response);

    // Parse prune response (simplified - assumes success/failure byte)
    if (response.len > 0) {
        if (response[0] == 0x00) {
            if (flags.json) {
                std.debug.print("{\"ok\":true}\n", .{});
            } else {
                io.success("Prune operation completed successfully\n", .{});
            }
        } else {
            if (flags.json) {
                std.debug.print("{\"ok\":false,\"error_code\":{d}}\n", .{response[0]});
            } else {
                io.err("[FAIL] Prune operation failed: error code {d}\n", .{response[0]});
            }
            return error.PruneFailed;
        }
    } else {
        if (flags.json) {
            std.debug.print("{\"ok\":true,\"note\":\"no_response\"}\n", .{});
        } else {
            io.success("Prune request sent (no response received)\n", .{});
        }
    }
}

fn printUsage() void {
    io.info("Usage: ml prune [options]\n\n", .{});
    io.info("Options:\n", .{});
    io.info("\t--keep <N>\t\tKeep N most recent experiments\n", .{});
    io.info("\t--older-than <days>\tRemove experiments older than N days\n", .{});
    io.info("\t--json\t\t\tOutput machine-readable JSON\n", .{});
    io.info("\t--help, -h\t\tShow this help message\n", .{});
}