fetch_ml/cli/src/commands/dataset.zig

const std = @import("std");
const Config = @import("../config.zig").Config;
const ws = @import("../net/ws.zig");
const crypto = @import("../utils/crypto.zig");
const colors = @import("../utils/colors.zig");
const logging = @import("../utils/logging.zig");

pub fn run(allocator: std.mem.Allocator, args: []const []const u8) !void {
    if (args.len == 0) {
        colors.printError("Usage: ml dataset <action> [options]\n", .{});
        colors.printInfo("Actions:\n", .{});
        colors.printInfo("  list                    List registered datasets\n", .{});
        colors.printInfo("  register <name> <url>   Register a dataset with URL\n", .{});
        colors.printInfo("  info <name>            Show dataset information\n", .{});
        colors.printInfo("  search <term>          Search datasets by name/description\n", .{});
        return error.InvalidArgs;
    }

    const action = args[0];

    if (std.mem.eql(u8, action, "list")) {
        try listDatasets(allocator);
    } else if (std.mem.eql(u8, action, "register")) {
        if (args.len < 3) {
            colors.printError("Usage: ml dataset register <name> <url>\n", .{});
            return error.InvalidArgs;
        }
        try registerDataset(allocator, args[1], args[2]);
    } else if (std.mem.eql(u8, action, "info")) {
        if (args.len < 2) {
            colors.printError("Usage: ml dataset info <name>\n", .{});
            return error.InvalidArgs;
        }
        try showDatasetInfo(allocator, args[1]);
    } else if (std.mem.eql(u8, action, "search")) {
        if (args.len < 2) {
            colors.printError("Usage: ml dataset search <term>\n", .{});
            return error.InvalidArgs;
        }
        try searchDatasets(allocator, args[1]);
    } else {
        colors.printError("Unknown action: {s}\n", .{action});
        return error.InvalidArgs;
    }
}

fn listDatasets(allocator: std.mem.Allocator) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

    // Authenticate with server to get user context
    var user_context = try authenticateUser(allocator, config);
    defer user_context.deinit();

    // Connect to WebSocket and request dataset list
    const api_key_plain = config.api_key;
    const api_key_hash = try crypto.hashString(allocator, api_key_plain);
    defer allocator.free(api_key_hash);

    const ws_url = try std.fmt.allocPrint(allocator, "ws://{s}:9101/ws", .{config.worker_host});
    defer allocator.free(ws_url);

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

    try client.sendDatasetList(api_key_hash);

    // Receive and display dataset list
    const response = try client.receiveAndHandleDatasetResponse(allocator);
    defer allocator.free(response);

    colors.printInfo("Registered Datasets:\n", .{});
    colors.printInfo("=====================\n\n", .{});

    // Parse and display datasets (simplified for now)
    if (std.mem.eql(u8, response, "[]")) {
        colors.printWarning("No datasets registered.\n", .{});
        colors.printInfo("Use 'ml dataset register <name> <url>' to add a dataset.\n", .{});
    } else {
        colors.printSuccess("{s}\n", .{response});
    }
}

fn registerDataset(allocator: std.mem.Allocator, name: []const u8, url: []const u8) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

    // Validate URL format
    if (!std.mem.startsWith(u8, url, "http://") and !std.mem.startsWith(u8, url, "https://") and
        !std.mem.startsWith(u8, url, "s3://") and !std.mem.startsWith(u8, url, "gs://"))
    {
        colors.printError("Invalid URL format. Supported: http://, https://, s3://, gs://\n", .{});
        return error.InvalidURL;
    }

    // Authenticate with server
    var user_context = try authenticateUser(allocator, config);
    defer user_context.deinit();

    // Connect to WebSocket and register dataset
    const api_key_plain = config.api_key;
    const api_key_hash = try crypto.hashString(allocator, api_key_plain);
    defer allocator.free(api_key_hash);

    const ws_url = try std.fmt.allocPrint(allocator, "ws://{s}:9101/ws", .{config.worker_host});
    defer allocator.free(ws_url);

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

    try client.sendDatasetRegister(name, url, api_key_hash);

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

    if (std.mem.startsWith(u8, response, "ERROR")) {
        colors.printError("Failed to register dataset: {s}\n", .{response});
    } else {
        colors.printSuccess("Dataset '{s}' registered successfully!\n", .{name});
        colors.printInfo("URL: {s}\n", .{url});
    }
}

fn showDatasetInfo(allocator: std.mem.Allocator, name: []const u8) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

    // Authenticate with server
    var user_context = try authenticateUser(allocator, config);
    defer user_context.deinit();

    // Connect to WebSocket and get dataset info
    const api_key_plain = config.api_key;
    const api_key_hash = try crypto.hashString(allocator, api_key_plain);
    defer allocator.free(api_key_hash);

    const ws_url = try std.fmt.allocPrint(allocator, "ws://{s}:9101/ws", .{config.worker_host});
    defer allocator.free(ws_url);

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

    try client.sendDatasetInfo(name, api_key_hash);

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

    if (std.mem.startsWith(u8, response, "ERROR") or std.mem.startsWith(u8, response, "NOT_FOUND")) {
        colors.printError("Dataset '{s}' not found.\n", .{name});
    } else {
        colors.printInfo("Dataset Information:\n", .{});
        colors.printInfo("===================\n", .{});
        colors.printSuccess("Name: {s}\n", .{name});
        colors.printSuccess("Details: {s}\n", .{response});
    }
}

fn searchDatasets(allocator: std.mem.Allocator, term: []const u8) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

    // Authenticate with server
    var user_context = try authenticateUser(allocator, config);
    defer user_context.deinit();

    // Connect to WebSocket and search datasets
    const api_key_plain = config.api_key;
    const api_key_hash = try crypto.hashString(allocator, api_key_plain);
    defer allocator.free(api_key_hash);

    const ws_url = try std.fmt.allocPrint(allocator, "ws://{s}:9101/ws", .{config.worker_host});
    defer allocator.free(ws_url);

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

    try client.sendDatasetSearch(term, api_key_hash);

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

    colors.printInfo("Search Results for '{s}':\n", .{term});
    colors.printInfo("========================\n\n", .{});

    if (std.mem.eql(u8, response, "[]")) {
        colors.printWarning("No datasets found matching '{s}'.\n", .{term});
    } else {
        colors.printSuccess("{s}\n", .{response});
    }
}

// Reuse authenticateUser from other commands
fn authenticateUser(allocator: std.mem.Allocator, config: Config) !UserContext {
    const ws_url = try std.fmt.allocPrint(allocator, "ws://{s}:9101/ws", .{config.worker_host});
    defer allocator.free(ws_url);

    // Try to connect with the API key to validate it
    var client = ws.Client.connect(allocator, ws_url, config.api_key) catch |err| {
        switch (err) {
            error.ConnectionRefused => return error.ConnectionFailed,
            error.NetworkUnreachable => return error.ServerUnreachable,
            error.InvalidURL => return error.ConfigInvalid,
            else => return error.AuthenticationFailed,
        }
    };
    defer client.close();

    // For now, create a user context after successful authentication
    const user_name = try allocator.dupe(u8, "authenticated_user");
    return UserContext{
        .name = user_name,
        .admin = false,
        .allocator = allocator,
    };
}

const UserContext = struct {
    name: []const u8,
    admin: bool,
    allocator: std.mem.Allocator,

    pub fn deinit(self: *UserContext) void {
        self.allocator.free(self.name);
    }
};