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

// Optimized command dispatch
const Command = enum {
    jupyter,
    init,
    sync,
    requeue,
    queue,
    status,
    monitor,
    cancel,
    prune,
    watch,
    dataset,
    experiment,
    validate,
    info,
    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,
            'i' => if (std.mem.eql(u8, str, "info")) return .info,
            '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,
            'v' => if (std.mem.eql(u8, str, "validate")) return .validate,
            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];

    // Track if we found a valid command
    var command_found = false;

    // Fast dispatch using switch on first character
    switch (command[0]) {
        'j' => if (std.mem.eql(u8, command, "jupyter")) {
            command_found = true;
            try @import("commands/jupyter.zig").run(allocator, args[2..]);
        },
        'i' => if (std.mem.eql(u8, command, "init")) {
            command_found = true;
            colors.printInfo("Setup configuration interactively\n", .{});
        } else if (std.mem.eql(u8, command, "info")) {
            command_found = true;
            try @import("commands/info.zig").run(allocator, args[2..]);
        },
        'a' => if (std.mem.eql(u8, command, "annotate")) {
            command_found = true;
            try @import("commands/annotate.zig").run(allocator, args[2..]);
        },
        'n' => if (std.mem.eql(u8, command, "narrative")) {
            command_found = true;
            try @import("commands/narrative.zig").run(allocator, args[2..]);
        },
        's' => if (std.mem.eql(u8, command, "sync")) {
            command_found = true;
            if (args.len < 3) {
                colors.printError("Usage: ml sync <path>\n", .{});
                std.process.exit(1);
            }
            colors.printInfo("Sync project to server: {s}\n", .{args[2]});
        } else if (std.mem.eql(u8, command, "status")) {
            command_found = true;
            try @import("commands/status.zig").run(allocator, args[2..]);
        },
        'r' => if (std.mem.eql(u8, command, "requeue")) {
            command_found = true;
            try @import("commands/requeue.zig").run(allocator, args[2..]);
        },
        'q' => if (std.mem.eql(u8, command, "queue")) {
            command_found = true;
            try @import("commands/queue.zig").run(allocator, args[2..]);
        },
        'd' => if (std.mem.eql(u8, command, "dataset")) {
            command_found = true;
            try @import("commands/dataset.zig").run(allocator, args[2..]);
        },
        'e' => if (std.mem.eql(u8, command, "experiment")) {
            command_found = true;
            try @import("commands/experiment.zig").execute(allocator, args[2..]);
        },
        'c' => if (std.mem.eql(u8, command, "cancel")) {
            command_found = true;
            try @import("commands/cancel.zig").run(allocator, args[2..]);
        },
        'v' => if (std.mem.eql(u8, command, "validate")) {
            command_found = true;
            try @import("commands/validate.zig").run(allocator, args[2..]);
        },
        else => {},
    }

    // If no command was found, show error and exit
    if (!command_found) {
        colors.printError("Unknown command: {s}\n", .{args[1]});
        printUsage();
        std.process.exit(1);
    }
}

// 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("  annotate <path|id> Add an annotation to run_manifest.json (--note \"...\")\n", .{});
    std.debug.print("  narrative set <path|id> Set run narrative fields (hypothesis/context/...)\n", .{});
    std.debug.print("  info <path|id> Show run info from run_manifest.json (optionally --base <path>)\n", .{});
    std.debug.print("  sync <path>    Sync project to server\n", .{});
    std.debug.print("  requeue <id>   Re-submit from run_id/task_id/path (supports -- <args>)\n", .{});
    std.debug.print("  queue (q) <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 and metrics\n", .{});
    std.debug.print("  validate       Validate provenance and integrity for a commit/task\n", .{});
    std.debug.print("\nUse 'ml <command> --help' for detailed help.\n", .{});
}

test {
    _ = @import("commands/info.zig");
    _ = @import("commands/requeue.zig");
    _ = @import("commands/annotate.zig");
    _ = @import("commands/narrative.zig");
}