const std = @import("std");

// Handle unknown command - prints error and exits
fn handleUnknownCommand(cmd: []const u8) noreturn {
    std.debug.print("Error: Unknown command: {s}\n", .{cmd});
    printUsage();
    std.process.exit(1);
}

pub fn main() !void {
    // Use c_allocator for better performance on Linux
    const allocator = std.heap.c_allocator;

    const args = std.process.argsAlloc(allocator) catch |err| {
        std.debug.print("Failed to allocate args: {}\n", .{err});
        return;
    };
    defer std.process.argsFree(allocator, args);

    if (args.len < 2) {
        printUsage();
        return;
    }

    const command = args[1];

    // Fast dispatch using switch on first character
    switch (command[0]) {
        'j' => if (std.mem.eql(u8, command, "jupyter")) {
            try @import("commands/jupyter.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'i' => if (std.mem.eql(u8, command, "init")) {
            try @import("commands/init.zig").run(allocator, args[2..]);
        } else if (std.mem.eql(u8, command, "info")) {
            try @import("commands/info.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'a' => if (std.mem.eql(u8, command, "annotate")) {
            try @import("commands/annotate.zig").execute(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'e' => if (std.mem.eql(u8, command, "experiment")) {
            try @import("commands/experiment.zig").execute(allocator, args[2..]);
        } else if (std.mem.eql(u8, command, "export")) {
            try @import("commands/export_cmd.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        's' => if (std.mem.eql(u8, command, "sync")) {
            try @import("commands/sync.zig").run(allocator, args[2..]);
        } else if (std.mem.eql(u8, command, "status")) {
            try @import("commands/status.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'r' => if (std.mem.eql(u8, command, "run")) {
            try @import("commands/run.zig").execute(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'q' => if (std.mem.eql(u8, command, "queue")) {
            try @import("commands/queue.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'd' => if (std.mem.eql(u8, command, "dataset")) {
            try @import("commands/dataset.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'x' => if (std.mem.eql(u8, command, "export")) {
            try @import("commands/export_cmd.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'c' => if (std.mem.eql(u8, command, "cancel")) {
            try @import("commands/cancel.zig").run(allocator, args[2..]);
        } else if (std.mem.eql(u8, command, "compare")) {
            try @import("commands/compare.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'f' => if (std.mem.eql(u8, command, "find")) {
            try @import("commands/find.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'v' => if (std.mem.eql(u8, command, "validate")) {
            try @import("commands/validate.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'l' => if (std.mem.eql(u8, command, "logs")) {
            try @import("commands/log.zig").run(allocator, args[2..]);
        } else if (std.mem.eql(u8, command, "log")) {
            try @import("commands/log.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        'w' => if (std.mem.eql(u8, command, "watch")) {
            try @import("commands/watch.zig").run(allocator, args[2..]);
        } else handleUnknownCommand(command),
        else => {
            std.debug.print("Error: Unknown command: {s}\n", .{args[1]});
            printUsage();
            return error.InvalidCommand;
        },
    }
}

// Optimized usage printer
fn printUsage() void {
    std.debug.print("ML Experiment Manager\n\n", .{});
    std.debug.print("Usage: ml <command> [options]\n\n", .{});
    std.debug.print("Commands:\n", .{});
    std.debug.print("  init           Initialize project with config\n", .{});
    std.debug.print("  run [args]     Execute a run locally\n", .{});
    std.debug.print("  queue <job>    Queue job on server\n", .{});
    std.debug.print("  annotate <id>  Add metadata annotations\n", .{});
    std.debug.print("  experiment     Manage experiments (create, list, show)\n", .{});
    std.debug.print("  logs <id>      Fetch or stream run logs\n", .{});
    std.debug.print("  sync [id]      Push local runs to server\n", .{});
    std.debug.print("  cancel <id>    Cancel local run\n", .{});
    std.debug.print("  watch [--sync] Watch directory with optional auto-sync\n", .{});
    std.debug.print("  status         Get system status\n", .{});
    std.debug.print("  dataset        Manage datasets\n", .{});
    std.debug.print("  export <id>    Export experiment bundle\n", .{});
    std.debug.print("  validate       Validate provenance and integrity\n", .{});
    std.debug.print("  compare <a> <b>  Compare two runs\n", .{});
    std.debug.print("  find [query]   Search experiments\n", .{});
    std.debug.print("  jupyter        Jupyter workspace management\n", .{});
    std.debug.print("  info <id>      Show run info\n", .{});
    std.debug.print("\nUse 'ml <command> --help' for detailed help.\n", .{});
}

test {
    _ = @import("commands/info.zig");
    _ = @import("commands/compare.zig");
    _ = @import("commands/find.zig");
    _ = @import("commands/export_cmd.zig");
    _ = @import("commands/log.zig");
    _ = @import("commands/annotate.zig");
    _ = @import("commands/experiment.zig");
}