const std = @import("std");
const Config = @import("../config.zig").Config;
const ws = @import("../net/ws/client.zig");
const crypto = @import("../utils/crypto.zig");
const core = @import("../core.zig");
const native_hash = @import("../native/hash.zig");

const DatasetOptions = struct {
    dry_run: bool = false,
    validate: bool = false,
    json: bool = false,
    csv: bool = false,
};

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

    var flags = core.flags.CommonFlags{};
    var dry_run = false;
    var validate = false;
    var csv = false;

    // Parse global flags: --dry-run, --validate, --json
    var positional = std.ArrayList([]const u8).initCapacity(allocator, args.len) catch |err| {
        return err;
    };
    defer positional.deinit(allocator);

    for (args) |arg| {
        if (std.mem.eql(u8, arg, "--help") or std.mem.eql(u8, arg, "-h")) {
            return printUsage();
        } else if (std.mem.eql(u8, arg, "--dry-run")) {
            dry_run = true;
        } else if (std.mem.eql(u8, arg, "--validate")) {
            validate = true;
        } else if (std.mem.eql(u8, arg, "--json")) {
            flags.json = true;
        } else if (std.mem.eql(u8, arg, "--csv")) {
            csv = true;
        } else if (std.mem.startsWith(u8, arg, "--")) {
            core.output.err("Unknown option");
            return printUsage();
        } else {
            try positional.append(allocator, arg);
        }
    }

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

    const action = positional.items[0];

    switch (positional.items.len) {
        0 => {
            return printUsage();
        },
        1 => {
            if (std.mem.eql(u8, action, "list")) {
                const options = DatasetOptions{ .json = flags.json, .csv = csv };
                try listDatasets(allocator, &options);
                return;
            }
        },
        2 => {
            if (std.mem.eql(u8, action, "info")) {
                const options = DatasetOptions{ .json = flags.json, .csv = csv };
                try showDatasetInfo(allocator, positional.items[1], &options);
                return;
            } else if (std.mem.eql(u8, action, "search")) {
                const options = DatasetOptions{ .json = flags.json, .csv = csv };
                try searchDatasets(allocator, positional.items[1], &options);
                return;
            } else if (std.mem.eql(u8, action, "verify")) {
                const options = DatasetOptions{ .json = flags.json, .validate = validate };
                try verifyDataset(allocator, positional.items[1], &options);
                return;
            }
        },
        3 => {
            if (std.mem.eql(u8, action, "register")) {
                const options = DatasetOptions{ .json = flags.json, .dry_run = dry_run };
                try registerDataset(allocator, positional.items[1], positional.items[2], &options);
                return;
            }
        },
        else => {
            core.output.err("Too many arguments");
            return error.InvalidArgs;
        },
    }

    return printUsage();
}

fn printUsage() void {
    std.debug.print("Usage: ml dataset <action> [options]\n\n", .{});
    std.debug.print("Actions:\n", .{});
    std.debug.print("\tlist\t\t\tList registered datasets\n", .{});
    std.debug.print("\tregister <name> <url>\tRegister a dataset with URL\n", .{});
    std.debug.print("\tinfo <name>\t\tShow dataset information\n", .{});
    std.debug.print("\tsearch <term>\t\tSearch datasets by name/description\n", .{});
    std.debug.print("\tverify <path|id>\tVerify dataset integrity (auto-hashes)\n", .{});
    std.debug.print("\nOptions:\n", .{});
    std.debug.print("\t--dry-run\t\tShow what would be requested\n", .{});
    std.debug.print("\t--validate\t\tValidate inputs only (no request)\n", .{});
    std.debug.print("\t--json\t\t\tOutput machine-readable JSON\n", .{});
    std.debug.print("\t--help, -h\t\tShow this help message\n", .{});
}

fn listDatasets(allocator: std.mem.Allocator, options: *const DatasetOptions) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

    // Connect to WebSocket and request dataset list
    const api_key_hash = try crypto.hashApiKey(allocator, config.api_key);
    defer allocator.free(api_key_hash);

    if (options.validate) {
        if (options.json) {
            const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
            var buffer: [4096]u8 = undefined;
            const formatted = std.fmt.bufPrint(&buffer, "{{\"ok\":true,\"action\":\"list\",\"validated\":true}}\n", .{}) catch unreachable;
            try stdout_file.writeAll(formatted);
        } else {
            std.debug.print("Validation OK\n", .{});
        }
        return;
    }

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

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

    if (options.dry_run) {
        if (options.json) {
            const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
            var buffer: [4096]u8 = undefined;
            const formatted = std.fmt.bufPrint(&buffer, "{{\"dry_run\":true,\"action\":\"list\"}}\n", .{}) catch unreachable;
            try stdout_file.writeAll(formatted);
        } else {
            std.debug.print("Dry run: would request dataset list\n", .{});
        }
        return;
    }

    try client.sendDatasetList(api_key_hash);

    // Receive and display dataset list
    const response = try client.receiveAndHandleDatasetResponse(allocator);
    defer allocator.free(response);

    if (options.json) {
        const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
        var buffer: [4096]u8 = undefined;
        const formatted = std.fmt.bufPrint(&buffer, "{s}\n", .{response}) catch unreachable;
        try stdout_file.writeAll(formatted);
        return;
    }

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

    // Parse and display datasets (simplified for now)
    if (std.mem.eql(u8, response, "[]")) {
        std.debug.print("No datasets registered.\n", .{});
    } else {
        std.debug.print("{s}\n", .{response});
    }
}

fn registerDataset(allocator: std.mem.Allocator, name: []const u8, url: []const u8, options: *const DatasetOptions) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

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

    if (options.validate) {
        if (name.len == 0 or name.len > 255) return error.InvalidArgs;
        if (url.len == 0 or url.len > 1023) return error.InvalidURL;

        // Validate URL format
        if (!std.mem.startsWith(u8, url, "http://") and !std.mem.startsWith(u8, url, "https://") and
            !std.mem.startsWith(u8, url, "s3://") and !std.mem.startsWith(u8, url, "gs://"))
        {
            return error.InvalidURL;
        }

        if (options.json) {
            const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
            var buffer: [4096]u8 = undefined;
            const formatted = std.fmt.bufPrint(&buffer, "{{\"ok\":true,\"action\":\"register\",\"validated\":true,\"name\":\"{s}\",\"url\":\"{s}\"}}\n", .{ name, url }) catch unreachable;
            try stdout_file.writeAll(formatted);
        } else {
            std.debug.print("Validation OK\n", .{});
        }
        return;
    }

    // Validate URL format
    if (!std.mem.startsWith(u8, url, "http://") and !std.mem.startsWith(u8, url, "https://") and
        !std.mem.startsWith(u8, url, "s3://") and !std.mem.startsWith(u8, url, "gs://"))
    {
        std.debug.print("Invalid URL format. Supported: http://, https://, s3://, gs://\n", .{});
        return error.InvalidURL;
    }

    // Connect to WebSocket and register dataset

    if (options.dry_run) {
        if (options.json) {
            const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
            var buffer: [4096]u8 = undefined;
            const formatted = std.fmt.bufPrint(&buffer, "{{\"dry_run\":true,\"action\":\"register\",\"name\":\"{s}\",\"url\":\"{s}\"}}\n", .{ name, url }) catch unreachable;
            try stdout_file.writeAll(formatted);
        } else {
            std.debug.print("Dry run: would register dataset '{s}' -> {s}\n", .{ name, url });
        }
        return;
    }

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

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

    try client.sendDatasetRegister(name, url, api_key_hash);

    // Receive response
    const response = try client.receiveAndHandleDatasetResponse(allocator);
    defer allocator.free(response);

    if (options.json) {
        const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
        var buffer: [4096]u8 = undefined;
        const formatted = std.fmt.bufPrint(&buffer, "{{\"ok\":true,\"action\":\"register\",\"message\":\"{s}\"}}\n", .{response}) catch unreachable;
        try stdout_file.writeAll(formatted);
        return;
    }

    if (std.mem.startsWith(u8, response, "ERROR")) {
        std.debug.print("Failed to register dataset: {s}\n", .{response});
    } else {
        std.debug.print("Dataset '{s}' registered\n", .{name});
    }
}

fn showDatasetInfo(allocator: std.mem.Allocator, name: []const u8, options: *const DatasetOptions) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

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

    if (options.validate) {
        if (name.len == 0 or name.len > 255) return error.InvalidArgs;
        if (options.json) {
            const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
            var buffer: [4096]u8 = undefined;
            const formatted = std.fmt.bufPrint(&buffer, "{{\"ok\":true,\"action\":\"info\",\"validated\":true,\"name\":\"{s}\"}}\n", .{name}) catch unreachable;
            try stdout_file.writeAll(formatted);
        } else {
            std.debug.print("Validation OK\n", .{});
        }
        return;
    }

    // Connect to WebSocket and get dataset info

    if (options.dry_run) {
        if (options.json) {
            const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
            var buffer: [4096]u8 = undefined;
            const formatted = std.fmt.bufPrint(&buffer, "{{\"dry_run\":true,\"action\":\"info\",\"name\":\"{s}\"}}\n", .{name}) catch unreachable;
            try stdout_file.writeAll(formatted);
        } else {
            std.debug.print("Dry run: would request dataset info for '{s}'\n", .{name});
        }
        return;
    }

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

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

    try client.sendDatasetInfo(name, api_key_hash);

    // Receive response
    const response = try client.receiveAndHandleDatasetResponse(allocator);
    defer allocator.free(response);

    if (options.json) {
        const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
        var buffer: [4096]u8 = undefined;
        const formatted = std.fmt.bufPrint(&buffer, "{s}\n", .{response}) catch unreachable;
        try stdout_file.writeAll(formatted);
        return;
    }

    if (std.mem.startsWith(u8, response, "ERROR") or std.mem.startsWith(u8, response, "NOT_FOUND")) {
        std.debug.print("Dataset '{s}' not found.\n", .{name});
    } else {
        std.debug.print("{s}\n", .{response});
    }
}

fn searchDatasets(allocator: std.mem.Allocator, term: []const u8, options: *const DatasetOptions) !void {
    const config = try Config.load(allocator);
    defer {
        var mut_config = config;
        mut_config.deinit(allocator);
    }

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

    if (options.validate) {
        if (term.len == 0 or term.len > 255) return error.InvalidArgs;
        if (options.json) {
            const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
            var buffer: [4096]u8 = undefined;
            const formatted = std.fmt.bufPrint(&buffer, "{{\"ok\":true,\"action\":\"search\",\"validated\":true,\"term\":\"{s}\"}}\n", .{term}) catch unreachable;
            try stdout_file.writeAll(formatted);
        } else {
            std.debug.print("Validation OK\n", .{});
        }
        return;
    }

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

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

    try client.sendDatasetSearch(term, api_key_hash);

    // Receive response
    const response = try client.receiveAndHandleDatasetResponse(allocator);
    defer allocator.free(response);

    if (options.json) {
        const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
        var buffer: [4096]u8 = undefined;
        const formatted = std.fmt.bufPrint(&buffer, "{s}\n", .{response}) catch unreachable;
        try stdout_file.writeAll(formatted);
        return;
    }

    if (std.mem.eql(u8, response, "[]")) {
        std.debug.print("No datasets found matching '{s}'.\n", .{term});
    } else {
        std.debug.print("{s}\n", .{response});
    }
}

fn verifyDataset(allocator: std.mem.Allocator, target: []const u8, options: *const DatasetOptions) !void {
    std.debug.print("Verifying dataset: {s}\n", .{target});

    const path = if (std.fs.path.isAbsolute(target))
        target
    else
        try std.fs.path.join(allocator, &[_][]const u8{ ".", target });
    defer if (!std.fs.path.isAbsolute(target)) allocator.free(path);

    var dir = std.fs.openDirAbsolute(path, .{ .iterate = true }) catch {
        std.debug.print("Dataset not found: {s}\n", .{target});
        return error.FileNotFound;
    };
    defer dir.close();

    var file_count: usize = 0;
    var total_size: u64 = 0;

    var walker = try dir.walk(allocator);
    defer walker.deinit();

    while (try walker.next()) |entry| {
        if (entry.kind != .file) continue;
        file_count += 1;

        const full_path = try std.fs.path.join(allocator, &[_][]const u8{ path, entry.path });
        defer allocator.free(full_path);

        const stat = std.fs.cwd().statFile(full_path) catch continue;
        total_size += stat.size;
    }

    // Compute native SHA256 hash
    const hash = blk: {
        break :blk native_hash.hashDirectory(allocator, path) catch |err| {
            std.debug.print("Hash computation failed: {s}\n", .{@errorName(err)});
            // Continue without hash - verification still succeeded
            break :blk null;
        };
    };
    defer if (hash) |h| allocator.free(h);

    if (options.json) {
        const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
        var buffer: [4096]u8 = undefined;
        const hash_str = if (hash) |h| h else "null";
        const formatted = std.fmt.bufPrint(&buffer, "{{\"path\":\"{s}\",\"files\":{d},\"size\":{d},\"hash\":\"{s}\",\"ok\":true}}\n", .{
            target, file_count, total_size, hash_str,
        }) catch unreachable;
        try stdout_file.writeAll(formatted);
    } else if (options.csv) {
        const stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
        try stdout_file.writeAll("metric,value\n");
        var buf: [256]u8 = undefined;
        const line1 = try std.fmt.bufPrint(&buf, "path,{s}\n", .{target});
        try stdout_file.writeAll(line1);
        const line2 = try std.fmt.bufPrint(&buf, "files,{d}\n", .{file_count});
        try stdout_file.writeAll(line2);
        const line3 = try std.fmt.bufPrint(&buf, "size_bytes,{d}\n", .{total_size});
        try stdout_file.writeAll(line3);
        const line4 = try std.fmt.bufPrint(&buf, "size_mb,{d:.2}\n", .{@as(f64, @floatFromInt(total_size)) / (1024 * 1024)});
        try stdout_file.writeAll(line4);
        if (hash) |h| {
            const line5 = try std.fmt.bufPrint(&buf, "sha256,{s}\n", .{h});
            try stdout_file.writeAll(line5);
        }
    } else {
        std.debug.print("Dataset verified\n", .{});
        std.debug.print("{s}\t{d}\t{d:.2} MB\n", .{ target, file_count, @as(f64, @floatFromInt(total_size)) / (1024 * 1024) });
        if (hash) |h| {
            std.debug.print("SHA256\t{s}\n", .{h});
        }
    }
}

fn hashDataset(allocator: std.mem.Allocator, path: []const u8) !void {
    std.debug.print("Computing native SHA256 hash for: {s}\n", .{path});

    // Check SIMD availability
    if (!native_hash.hasSimdSha256()) {
        std.debug.print("SIMD SHA256 not available, using generic implementation\n", .{});
    } else {
        const impl_name = native_hash.getSimdImplName();
        std.debug.print("Using {s} SHA256 implementation\n", .{impl_name});
    }

    // Compute hash using native library
    const hash = native_hash.hashDirectory(allocator, path) catch |err| {
        switch (err) {
            error.ContextInitFailed => {
                std.debug.print("Failed to initialize native hash context\n", .{});
            },
            error.HashFailed => {
                std.debug.print("Hash computation failed\n", .{});
            },
            error.InvalidPath => {
                std.debug.print("Invalid path: {s}\n", .{path});
            },
            error.OutOfMemory => {
                std.debug.print("Out of memory\n", .{});
            },
        }
        return err;
    };
    defer allocator.free(hash);

    // Print result
    std.debug.print("SHA256: {s}\n", .{hash});
}

fn writeJSONString(writer: anytype, s: []const u8) !void {
    try writer.writeByte('"');
    for (s) |c| {
        switch (c) {
            '"' => try writer.writeAll("\\\""),
            '\\' => try writer.writeAll("\\\\"),
            '\n' => try writer.writeAll("\\n"),
            '\r' => try writer.writeAll("\\r"),
            '\t' => try writer.writeAll("\\t"),
            else => {
                if (c < 0x20) {
                    var buf: [6]u8 = undefined;
                    buf[0] = '\\';
                    buf[1] = 'u';
                    buf[2] = '0';
                    buf[3] = '0';
                    buf[4] = hexDigit(@intCast((c >> 4) & 0x0F));
                    buf[5] = hexDigit(@intCast(c & 0x0F));
                    try writer.writeAll(&buf);
                } else {
                    try writer.writeByte(c);
                }
            },
        }
    }
    try writer.writeByte('"');
}

fn hexDigit(v: u8) u8 {
    return if (v < 10) ('0' + v) else ('a' + (v - 10));
}
