const std = @import("std");
const colors = @import("utils/colors.zig");

// Optimized command dispatch
const Command = enum {
    jupyter,
    init,
    sync,
    queue,
    status,
    monitor,
    cancel,
    prune,
    watch,
    dataset,
    experiment,
    unknown,

    fn fromString(str: []const u8) Command {
        if (str.len == 0) return .unknown;

        // Fast path for common commands
        switch (str[0]) {
            'j' => if (std.mem.eql(u8, str, "jupyter")) return .jupyter,
            'i' => if (std.mem.eql(u8, str, "init")) return .init,
            's' => if (std.mem.eql(u8, str, "sync")) return .sync else if (std.mem.eql(u8, str, "status")) return .status,
            'q' => if (std.mem.eql(u8, str, "queue")) return .queue,
            'm' => if (std.mem.eql(u8, str, "monitor")) return .monitor,
            'c' => if (std.mem.eql(u8, str, "cancel")) return .cancel,
            'p' => if (std.mem.eql(u8, str, "prune")) return .prune,
            'w' => if (std.mem.eql(u8, str, "watch")) return .watch,
            'd' => if (std.mem.eql(u8, str, "dataset")) return .dataset,
            'e' => if (std.mem.eql(u8, str, "experiment")) return .experiment,
            else => return .unknown,
        }
        return .unknown;
    }
};

pub fn main() !void {
    // Initialize colors based on environment
    colors.initColors();

    // Use ArenaAllocator for thread-safe memory management
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

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

    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..]);
        },
        'i' => if (std.mem.eql(u8, command, "init")) {
            colors.printInfo("Setup configuration interactively\n", .{});
        },
        's' => if (std.mem.eql(u8, command, "sync")) {
            if (args.len < 3) {
                colors.printError("Usage: ml sync <path>\n", .{});
                return;
            }
            colors.printInfo("Sync project to server: {s}\n", .{args[2]});
        } else if (std.mem.eql(u8, command, "status")) {
            colors.printInfo("Getting system status...\n", .{});
        },
        'q' => if (std.mem.eql(u8, command, "queue")) {
            if (args.len < 3) {
                colors.printError("Usage: ml queue <job>\n", .{});
                return;
            }
            colors.printInfo("Queue job for execution: {s}\n", .{args[2]});
        },
        'm' => if (std.mem.eql(u8, command, "monitor")) {
            colors.printInfo("Launching TUI via SSH...\n", .{});
        },
        'c' => if (std.mem.eql(u8, command, "cancel")) {
            if (args.len < 3) {
                colors.printError("Usage: ml cancel <job>\n", .{});
                return;
            }
            colors.printInfo("Canceling job: {s}\n", .{args[2]});
        },
        else => {
            colors.printError("Unknown command: {s}\n", .{args[1]});
            printUsage();
        },
    }
}

// Optimized usage printer
fn printUsage() void {
    colors.printInfo("ML Experiment Manager\n\n", .{});
    std.debug.print("Usage: ml <command> [options]\n\n", .{});
    std.debug.print("Commands:\n", .{});
    std.debug.print("  jupyter        Jupyter workspace management\n", .{});
    std.debug.print("  init           Setup configuration interactively\n", .{});
    std.debug.print("  sync <path>    Sync project to server\n", .{});
    std.debug.print("  queue <job>    Queue job for execution\n", .{});
    std.debug.print("  status         Get system status\n", .{});
    std.debug.print("  monitor        Launch TUI via SSH\n", .{});
    std.debug.print("  cancel <job>   Cancel running job\n", .{});
    std.debug.print("  prune          Remove old experiments\n", .{});
    std.debug.print("  watch <path>   Watch directory for auto-sync\n", .{});
    std.debug.print("  dataset        Manage datasets\n", .{});
    std.debug.print("  experiment     Manage experiments\n", .{});
    std.debug.print("\nUse 'ml <command> --help' for detailed help.\n", .{});
}