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

pub fn run(allocator: std.mem.Allocator, args: []const []const u8) !void {
    if (args.len == 0) {
        try printUsage();
        return error.InvalidArgs;
    }

    if (std.mem.eql(u8, args[0], "--help") or std.mem.eql(u8, args[0], "-h")) {
        try printUsage();
        return;
    }

    const target = args[0];

    var author: []const u8 = "";
    var note: ?[]const u8 = null;
    var base_override: ?[]const u8 = null;
    var json_mode: bool = false;

    var i: usize = 1;
    while (i < args.len) : (i += 1) {
        const a = args[i];
        if (std.mem.eql(u8, a, "--author")) {
            if (i + 1 >= args.len) {
                colors.printError("Missing value for --author\n", .{});
                return error.InvalidArgs;
            }
            author = args[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, a, "--note")) {
            if (i + 1 >= args.len) {
                colors.printError("Missing value for --note\n", .{});
                return error.InvalidArgs;
            }
            note = args[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, a, "--base")) {
            if (i + 1 >= args.len) {
                colors.printError("Missing value for --base\n", .{});
                return error.InvalidArgs;
            }
            base_override = args[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, a, "--json")) {
            json_mode = true;
        } else if (std.mem.eql(u8, a, "--help") or std.mem.eql(u8, a, "-h")) {
            try printUsage();
            return;
        } else if (std.mem.startsWith(u8, a, "--")) {
            colors.printError("Unknown option: {s}\n", .{a});
            return error.InvalidArgs;
        } else {
            colors.printError("Unexpected argument: {s}\n", .{a});
            return error.InvalidArgs;
        }
    }

    if (note == null or std.mem.trim(u8, note.?, " \t\r\n").len == 0) {
        colors.printError("--note is required\n", .{});
        try printUsage();
        return error.InvalidArgs;
    }

    const cfg = try Config.load(allocator);
    defer {
        var mut_cfg = cfg;
        mut_cfg.deinit(allocator);
    }

    const resolved_base = base_override orelse cfg.worker_base;

    const manifest_path = manifest.resolvePathWithBase(allocator, target, resolved_base) catch |err| {
        if (err == error.FileNotFound) {
            colors.printError(
                "Could not locate run_manifest.json for '{s}'. Provide a path, or use --base <path> to scan finished/failed/running/pending.\n",
                .{target},
            );
        }
        return err;
    };
    defer allocator.free(manifest_path);

    const job_name = try manifest.readJobNameFromManifest(allocator, manifest_path);
    defer allocator.free(job_name);

    const api_key_hash = try crypto.hashApiKey(allocator, cfg.api_key);
    defer allocator.free(api_key_hash);

    const ws_url = try cfg.getWebSocketUrl(allocator);
    defer allocator.free(ws_url);

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

    try client.sendAnnotateRun(job_name, author, note.?, api_key_hash);

    if (json_mode) {
        const msg = try client.receiveMessage(allocator);
        defer allocator.free(msg);

        const packet = protocol.ResponsePacket.deserialize(msg, allocator) catch {
            var out = io.stdoutWriter();
            try out.print("{s}\n", .{msg});
            return error.InvalidPacket;
        };
        defer packet.deinit(allocator);

        const Result = struct {
            ok: bool,
            job_name: []const u8,
            message: []const u8,
            error_code: ?u8 = null,
            error_message: ?[]const u8 = null,
            details: ?[]const u8 = null,
        };

        var out = io.stdoutWriter();
        if (packet.packet_type == .error_packet) {
            const res = Result{
                .ok = false,
                .job_name = job_name,
                .message = "",
                .error_code = @intFromEnum(packet.error_code.?),
                .error_message = packet.error_message orelse "",
                .details = packet.error_details orelse "",
            };
            try out.print("{f}\n", .{std.json.fmt(res, .{})});
            return error.CommandFailed;
        }

        const res = Result{
            .ok = true,
            .job_name = job_name,
            .message = packet.success_message orelse "",
        };
        try out.print("{f}\n", .{std.json.fmt(res, .{})});
        return;
    }

    try client.receiveAndHandleResponse(allocator, "Annotate");

    colors.printSuccess("Annotation added\n", .{});
    colors.printInfo("Job: {s}\n", .{job_name});
}

fn printUsage() !void {
    colors.printInfo("Usage: ml annotate <path|run_id|task_id> --note <text> [--author <name>] [--base <path>] [--json]\n", .{});
    colors.printInfo("\nExamples:\n", .{});
    colors.printInfo("  ml annotate 8b3f... --note \"Try lr=3e-4 next\"\n", .{});
    colors.printInfo("  ml annotate ./finished/job-123 --note \"Baseline looks stable\" --author alice\n", .{});
}