const std = @import("std");
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 core = @import("../core.zig");

pub const CompareOptions = struct {
    json: bool = false,
    csv: bool = false,
    all_fields: bool = false,
    fields: ?[]const u8 = null,
};

pub fn run(allocator: std.mem.Allocator, argv: []const []const u8) !void {
    if (argv.len < 2) {
        core.output.usage("compare", "Expected two run IDs");
        return error.InvalidArgs;
    }

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

    const run_a = argv[0];
    const run_b = argv[1];

    var flags = core.flags.CommonFlags{};
    var csv: bool = false;
    var all_fields: bool = false;
    var fields: ?[]const u8 = null;

    var i: usize = 2;
    while (i < argv.len) : (i += 1) {
        const arg = argv[i];
        if (std.mem.eql(u8, arg, "--json")) {
            flags.json = true;
        } else if (std.mem.eql(u8, arg, "--csv")) {
            csv = true;
        } else if (std.mem.eql(u8, arg, "--all")) {
            all_fields = true;
        } else if (std.mem.eql(u8, arg, "--fields") and i + 1 < argv.len) {
            fields = argv[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, arg, "--help") or std.mem.eql(u8, arg, "-h")) {
            return printUsage();
        } else {
            core.output.err("Unknown option");
            return error.InvalidArgs;
        }
    }

    core.output.setMode(if (flags.json) .json else .text);

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

    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);

    // Fetch both runs
    std.debug.print("Fetching run {s}...\n", .{run_a});
    var client_a = try ws.Client.connect(allocator, ws_url, cfg.api_key);
    defer client_a.close();

    // Try to get experiment info for run A
    try client_a.sendGetExperiment(run_a, api_key_hash);
    const msg_a = try client_a.receiveMessage(allocator);
    defer allocator.free(msg_a);

    std.debug.print("Fetching run {s}...\n", .{run_b});
    var client_b = try ws.Client.connect(allocator, ws_url, cfg.api_key);
    defer client_b.close();

    try client_b.sendGetExperiment(run_b, api_key_hash);
    const msg_b = try client_b.receiveMessage(allocator);
    defer allocator.free(msg_b);

    // Parse responses
    const parsed_a = std.json.parseFromSlice(std.json.Value, allocator, msg_a, .{}) catch {
        std.debug.print("Failed to parse response for {s}\n", .{run_a});
        return error.InvalidResponse;
    };
    defer parsed_a.deinit();

    const parsed_b = std.json.parseFromSlice(std.json.Value, allocator, msg_b, .{}) catch {
        std.debug.print("Failed to parse response for {s}\n", .{run_b});
        return error.InvalidResponse;
    };
    defer parsed_b.deinit();

    const root_a = parsed_a.value.object;
    const root_b = parsed_b.value.object;

    // Check for errors
    if (root_a.get("error")) |err_a| {
        std.debug.print("Error fetching {s}: {s}\n", .{ run_a, err_a.string });
        return error.ServerError;
    }
    if (root_b.get("error")) |err_b| {
        std.debug.print("Error fetching {s}: {s}\n", .{ run_b, err_b.string });
        return error.ServerError;
    }

    if (flags.json) {
        try outputJsonComparison(allocator, root_a, root_b, run_a, run_b);
    } else {
        try outputHumanComparison(root_a, root_b, run_a, run_b, all_fields);
    }
}

fn outputHumanComparison(
    root_a: std.json.ObjectMap,
    root_b: std.json.ObjectMap,
    run_a: []const u8,
    run_b: []const u8,
    all_fields: bool,
) !void {
    std.debug.print("\n=== Comparison: {s} vs {s} ===\n\n", .{ run_a, run_b });

    // Common fields
    const job_name_a = jsonGetString(root_a, "job_name") orelse "unknown";
    const job_name_b = jsonGetString(root_b, "job_name") orelse "unknown";

    if (!std.mem.eql(u8, job_name_a, job_name_b)) {
        std.debug.print("Job names differ:\n", .{});
        std.debug.print("\t{s}: {s}\n", .{ run_a, job_name_a });
        std.debug.print("\t{s}: {s}\n", .{ run_b, job_name_b });
    } else {
        std.debug.print("Job Name: {s}\n", .{job_name_a});
    }

    // Experiment group
    const group_a = jsonGetString(root_a, "experiment_group") orelse "";
    const group_b = jsonGetString(root_b, "experiment_group") orelse "";
    if (group_a.len > 0 or group_b.len > 0) {
        std.debug.print("\nExperiment Group:\n", .{});
        if (std.mem.eql(u8, group_a, group_b)) {
            std.debug.print("\tBoth: {s}\n", .{group_a});
        } else {
            std.debug.print("\t{s}: {s}\n", .{ run_a, group_a });
            std.debug.print("\t{s}: {s}\n", .{ run_b, group_b });
        }
    }

    // Narrative fields
    const narrative_a = root_a.get("narrative");
    const narrative_b = root_b.get("narrative");

    if (narrative_a != null or narrative_b != null) {
        std.debug.print("\n--- Narrative ---\n", .{});

        if (narrative_a) |na| {
            if (narrative_b) |nb| {
                if (na == .object and nb == .object) {
                    try compareNarrativeFields(na.object, nb.object, run_a, run_b);
                }
            } else {
                std.debug.print("\t{s} has narrative, {s} does not\n", .{ run_a, run_b });
            }
        } else if (narrative_b) |_| {
            std.debug.print("\t{s} has narrative, {s} does not\n", .{ run_b, run_a });
        }
    }

    // Metadata differences
    const meta_a = root_a.get("metadata");
    const meta_b = root_b.get("metadata");

    if (meta_a) |ma| {
        if (meta_b) |mb| {
            if (ma == .object and mb == .object) {
                std.debug.print("\n--- Metadata Differences ---\n", .{});
                try compareMetadata(ma.object, mb.object, run_a, run_b, all_fields);
            }
        }
    }

    // Metrics (if available)
    const metrics_a = root_a.get("metrics");
    const metrics_b = root_b.get("metrics");

    if (metrics_a) |ma| {
        if (metrics_b) |mb| {
            if (ma == .object and mb == .object) {
                std.debug.print("\n--- Metrics ---\n", .{});
                try compareMetrics(ma.object, mb.object, run_a, run_b);
            }
        }
    }

    // Outcome
    const outcome_a = jsonGetString(root_a, "outcome") orelse "";
    const outcome_b = jsonGetString(root_b, "outcome") orelse "";
    if (outcome_a.len > 0 or outcome_b.len > 0) {
        std.debug.print("\n--- Outcome ---\n", .{});
        if (std.mem.eql(u8, outcome_a, outcome_b)) {
            std.debug.print("\tBoth: {s}\n", .{outcome_a});
        } else {
            std.debug.print("\t{s}: {s}\n", .{ run_a, outcome_a });
            std.debug.print("\t{s}: {s}\n", .{ run_b, outcome_b });
        }
    }

    std.debug.print("\n", .{});
}

fn outputJsonComparison(
    allocator: std.mem.Allocator,
    root_a: std.json.ObjectMap,
    root_b: std.json.ObjectMap,
    run_a: []const u8,
    run_b: []const u8,
) !void {
    var buf = std.ArrayList(u8).empty;
    defer buf.deinit(allocator);

    const writer = buf.writer(allocator);

    try writer.writeAll("{\"run_a\":\"");
    try writer.writeAll(run_a);
    try writer.writeAll("\",\"run_b\":\"");
    try writer.writeAll(run_b);
    try writer.writeAll("\",\"differences\":");
    try writer.writeAll("{");

    var first = true;

    // Job names
    const job_name_a = jsonGetString(root_a, "job_name") orelse "";
    const job_name_b = jsonGetString(root_b, "job_name") orelse "";
    if (!std.mem.eql(u8, job_name_a, job_name_b)) {
        if (!first) try writer.writeAll(",");
        first = false;
        try writer.writeAll("\"job_name\":{\"a\":\"");
        try writer.writeAll(job_name_a);
        try writer.writeAll("\",\"b\":\"");
        try writer.writeAll(job_name_b);
        try writer.writeAll("\"}");
    }

    // Experiment group
    const group_a = jsonGetString(root_a, "experiment_group") orelse "";
    const group_b = jsonGetString(root_b, "experiment_group") orelse "";
    if (!std.mem.eql(u8, group_a, group_b)) {
        if (!first) try writer.writeAll(",");
        first = false;
        try writer.writeAll("\"experiment_group\":{\"a\":\"");
        try writer.writeAll(group_a);
        try writer.writeAll("\",\"b\":\"");
        try writer.writeAll(group_b);
        try writer.writeAll("\"}");
    }

    // Outcomes
    const outcome_a = jsonGetString(root_a, "outcome") orelse "";
    const outcome_b = jsonGetString(root_b, "outcome") orelse "";
    if (!std.mem.eql(u8, outcome_a, outcome_b)) {
        if (!first) try writer.writeAll(",");
        first = false;
        try writer.writeAll("\"outcome\":{\"a\":\"");
        try writer.writeAll(outcome_a);
        try writer.writeAll("\",\"b\":\"");
        try writer.writeAll(outcome_b);
        try writer.writeAll("\"}");
    }

    try writer.writeAll("}}");
    try writer.writeAll("}\n");

    const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
    try stdout_file.writeAll(buf.items);
}

fn compareNarrativeFields(
    na: std.json.ObjectMap,
    nb: std.json.ObjectMap,
    run_a: []const u8,
    run_b: []const u8,
) !void {
    const fields = [_][]const u8{ "hypothesis", "context", "intent", "expected_outcome" };

    for (fields) |field| {
        const val_a = jsonGetString(na, field);
        const val_b = jsonGetString(nb, field);

        if (val_a != null and val_b != null) {
            if (!std.mem.eql(u8, val_a.?, val_b.?)) {
                std.debug.print("\t{s}:\n", .{field});
                std.debug.print("\t\t{s}: {s}\n", .{ run_a, val_a.? });
                std.debug.print("\t\t{s}: {s}\n", .{ run_b, val_b.? });
            }
        } else if (val_a != null) {
            std.debug.print("\t{s}: only in {s}\n", .{ field, run_a });
        } else if (val_b != null) {
            std.debug.print("\t{s}: only in {s}\n", .{ field, run_b });
        }
    }
}

fn compareMetadata(
    ma: std.json.ObjectMap,
    mb: std.json.ObjectMap,
    run_a: []const u8,
    run_b: []const u8,
    show_all: bool,
) !void {
    var has_differences = false;

    // Compare key metadata fields
    const keys = [_][]const u8{ "batch_size", "learning_rate", "epochs", "model", "dataset" };

    for (keys) |key| {
        if (ma.get(key)) |va| {
            if (mb.get(key)) |vb| {
                const str_a = jsonValueToString(va);
                const str_b = jsonValueToString(vb);

                if (!std.mem.eql(u8, str_a, str_b)) {
                    has_differences = true;
                    std.debug.print("\t{s}: {s} ~> {s}\n", .{ key, str_a, str_b });
                } else if (show_all) {
                    std.debug.print("\t{s}: {s} (same)\n", .{ key, str_a });
                }
            } else if (show_all) {
                std.debug.print("\t{s}: only in {s}\n", .{ key, run_a });
            }
        } else if (mb.get(key)) |_| {
            if (show_all) {
                std.debug.print("\t{s}: only in {s}\n", .{ key, run_b });
            }
        }
    }

    if (!has_differences and !show_all) {
        std.debug.print("\t(no significant differences in common metadata)\n", .{});
    }
}

fn compareMetrics(
    ma: std.json.ObjectMap,
    mb: std.json.ObjectMap,
    run_a: []const u8,
    run_b: []const u8,
) !void {
    _ = run_a;
    _ = run_b;

    // Common metrics to compare
    const metrics = [_][]const u8{ "accuracy", "loss", "f1_score", "precision", "recall", "training_time", "validation_loss" };

    for (metrics) |metric| {
        if (ma.get(metric)) |va| {
            if (mb.get(metric)) |vb| {
                const val_a = jsonValueToFloat(va);
                const val_b = jsonValueToFloat(vb);

                const diff = val_b - val_a;
                const percent = if (val_a != 0) (diff / val_a) * 100 else 0;

                const arrow = if (diff > 0) "+" else if (diff < 0) "-" else "=";

                std.debug.print("  {s}: {d:.4} ~> {d:.4} ({s}{d:.4}, {d:.1}%)\n", .{
                    metric, val_a, val_b, arrow, @abs(diff), percent,
                });
            }
        }
    }
}

fn outputCsvComparison(
    allocator: std.mem.Allocator,
    root_a: std.json.ObjectMap,
    root_b: std.json.ObjectMap,
    run_a: []const u8,
    run_b: []const u8,
) !void {
    var buf = std.ArrayList(u8).empty;
    defer buf.deinit(allocator);

    const writer = buf.writer(allocator);

    // Header with actual run IDs as column names
    try writer.print("field,{s},{s},delta,notes\n", .{ run_a, run_b });

    // Job names
    const job_name_a = jsonGetString(root_a, "job_name") orelse "";
    const job_name_b = jsonGetString(root_b, "job_name") orelse "";
    const job_same = std.mem.eql(u8, job_name_a, job_name_b);
    try writer.print("job_name,\"{s}\",\"{s}\",{s},\"{s}\"\n", .{
        job_name_a,                          job_name_b,
        if (job_same) "same" else "changed", if (job_same) "" else "different job names",
    });

    // Outcomes
    const outcome_a = jsonGetString(root_a, "outcome") orelse "";
    const outcome_b = jsonGetString(root_b, "outcome") orelse "";
    const outcome_same = std.mem.eql(u8, outcome_a, outcome_b);
    try writer.print("outcome,{s},{s},{s},\"{s}\"\n", .{
        outcome_a,                               outcome_b,
        if (outcome_same) "same" else "changed", if (outcome_same) "" else "different outcomes",
    });

    // Experiment group
    const group_a = jsonGetString(root_a, "experiment_group") orelse "";
    const group_b = jsonGetString(root_b, "experiment_group") orelse "";
    const group_same = std.mem.eql(u8, group_a, group_b);
    try writer.print("experiment_group,\"{s}\",\"{s}\",{s},\"{s}\"\n", .{
        group_a,                               group_b,
        if (group_same) "same" else "changed", if (group_same) "" else "different groups",
    });

    // Metadata fields with delta calculation for numeric values
    const keys = [_][]const u8{ "batch_size", "learning_rate", "epochs", "model", "dataset" };
    for (keys) |key| {
        if (root_a.get(key)) |va| {
            if (root_b.get(key)) |vb| {
                const str_a = jsonValueToString(va);
                const str_b = jsonValueToString(vb);
                const same = std.mem.eql(u8, str_a, str_b);

                // Try to calculate delta for numeric values
                const delta = if (!same) blk: {
                    const f_a = jsonValueToFloat(va);
                    const f_b = jsonValueToFloat(vb);
                    if (f_a != 0 or f_b != 0) {
                        break :blk try std.fmt.allocPrint(allocator, "{d:.4}", .{f_b - f_a});
                    }
                    break :blk "changed";
                } else "0";
                defer if (!same and (jsonValueToFloat(va) != 0 or jsonValueToFloat(vb) != 0)) allocator.free(delta);

                try writer.print("{s},{s},{s},{s},\"{s}\"\n", .{
                    key,                             str_a, str_b, delta,
                    if (same) "same" else "changed",
                });
            }
        }
    }

    // Metrics with delta calculation
    const metrics = [_][]const u8{ "accuracy", "loss", "f1_score", "precision", "recall", "training_time" };
    for (metrics) |metric| {
        if (root_a.get(metric)) |va| {
            if (root_b.get(metric)) |vb| {
                const val_a = jsonValueToFloat(va);
                const val_b = jsonValueToFloat(vb);
                const diff = val_b - val_a;
                const percent = if (val_a != 0) (diff / val_a) * 100 else 0;

                const notes = if (std.mem.eql(u8, metric, "loss") or std.mem.eql(u8, metric, "training_time"))
                    if (diff < 0) "improved" else if (diff > 0) "degraded" else "same"
                else if (diff > 0) "improved" else if (diff < 0) "degraded" else "same";

                try writer.print("{s},{d:.4},{d:.4},{d:.4},\"{d:.1}% - {s}\"\n", .{
                    metric, val_a, val_b, diff, percent, notes,
                });
            }
        }
    }

    const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
    try stdout_file.writeAll(buf.items);
}

fn jsonGetString(obj: std.json.ObjectMap, key: []const u8) ?[]const u8 {
    const v_opt = obj.get(key);
    if (v_opt == null) return null;
    const v = v_opt.?;
    if (v != .string) return null;
    return v.string;
}

fn jsonValueToString(v: std.json.Value) []const u8 {
    return switch (v) {
        .string => |s| s,
        .integer => "number",
        .float => "number",
        .bool => |b| if (b) "true" else "false",
        else => "complex",
    };
}

fn jsonValueToFloat(v: std.json.Value) f64 {
    return switch (v) {
        .float => |f| f,
        .integer => |i| @as(f64, @floatFromInt(i)),
        else => 0,
    };
}

fn printUsage() !void {
    std.debug.print("Usage: ml compare <run-a> <run-b> [options]\n", .{});
    std.debug.print("\nCompare two runs and show differences in:\n", .{});
    std.debug.print("\t- Job metadata (batch_size, learning_rate, etc.)\n", .{});
    std.debug.print("\t- Narrative fields (hypothesis, context, intent)\n", .{});
    std.debug.print("\t- Metrics (accuracy, loss, training_time)\n", .{});
    std.debug.print("\t- Outcome status\n", .{});
    std.debug.print("\nOptions:\n", .{});
    std.debug.print("\t--json\t\tOutput as JSON\n", .{});
    std.debug.print("\t--all\t\tShow all fields (including unchanged)\n", .{});
    std.debug.print("\t--fields <csv>\tCompare only specific fields\n", .{});
    std.debug.print("\t--help, -h\tShow this help\n", .{});
    std.debug.print("\nExamples:\n", .{});
    std.debug.print("\tml compare run_abc run_def\n", .{});
    std.debug.print("\tml compare run_abc run_def --json\n", .{});
    std.debug.print("\tml compare run_abc run_def --all\n", .{});
}
