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

pub const ExportOptions = struct {
    anonymize: bool = false,
    anonymize_level: []const u8 = "metadata-only", // metadata-only, full
    bundle: ?[]const u8 = null,
    base_override: ?[]const u8 = null,
    json: bool = false,
};

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

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

    const target = argv[0];

    var flags = core.flags.CommonFlags{};
    var anonymize = false;
    var anonymize_level: []const u8 = "metadata-only";
    var bundle: ?[]const u8 = null;
    var base_override: ?[]const u8 = null;

    var i: usize = 1;
    while (i < argv.len) : (i += 1) {
        const arg = argv[i];
        if (std.mem.eql(u8, arg, "--anonymize")) {
            anonymize = true;
        } else if (std.mem.eql(u8, arg, "--anonymize-level") and i + 1 < argv.len) {
            anonymize_level = argv[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, arg, "--bundle") and i + 1 < argv.len) {
            bundle = argv[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, arg, "--base") and i + 1 < argv.len) {
            base_override = argv[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, arg, "--json")) {
            flags.json = true;
        } 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);

    // Validate anonymize level
    if (!std.mem.eql(u8, anonymize_level, "metadata-only") and
        !std.mem.eql(u8, anonymize_level, "full"))
    {
        core.output.err("Invalid anonymize level");
        return error.InvalidArgs;
    }

    if (flags.json) {
        var stdout_writer = io.stdoutWriter();
        try stdout_writer.print("{{\"success\":true,\"anonymize_level\":\"{s}\"}}\n", .{anonymize_level});
    } else {
        std.debug.print("Anonymization level: {s}\n", .{anonymize_level});
    }

    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) {
            std.debug.print(
                "Could not locate run_manifest.json for '{s}'.\n",
                .{target},
            );
        }
        return err;
    };
    defer allocator.free(manifest_path);

    // Read the manifest
    const manifest_content = manifest.readFileAlloc(allocator, manifest_path) catch |err| {
        std.debug.print("Failed to read manifest: {}\n", .{err});
        return err;
    };
    defer allocator.free(manifest_content);

    // Parse the manifest
    const parsed = std.json.parseFromSlice(std.json.Value, allocator, manifest_content, .{}) catch |err| {
        std.debug.print("Failed to parse manifest: {}\n", .{err});
        return err;
    };
    defer parsed.deinit();

    // Anonymize if requested
    var final_content: []u8 = undefined;
    var final_content_owned = false;

    if (anonymize) {
        final_content = try anonymizeManifest(allocator, parsed.value, anonymize_level);
        final_content_owned = true;
    } else {
        final_content = manifest_content;
    }
    defer if (final_content_owned) allocator.free(final_content);

    // Output or bundle
    if (bundle) |bundle_path| {
        // Create a simple tar-like bundle (just the manifest for now)
        // In production, this would include code, configs, etc.
        var bundle_file = try std.fs.cwd().createFile(bundle_path, .{});
        defer bundle_file.close();

        try bundle_file.writeAll(final_content);

        if (flags.json) {
            var stdout_writer = io.stdoutWriter();
            try stdout_writer.print("{{\"success\":true,\"bundle\":\"{s}\",\"anonymized\":{}}}\n", .{
                bundle_path,
                anonymize,
            });
        } else {
            std.debug.print("Exported to {s}\n", .{bundle_path});
            if (anonymize) {
                std.debug.print("\tAnonymization level: {s}\n", .{anonymize_level});
                std.debug.print("\tPaths redacted, IPs removed, usernames anonymized\n", .{});
            }
        }
    } else {
        // Output to stdout
        var stdout_writer = io.stdoutWriter();
        try stdout_writer.print("{s}\n", .{final_content});
    }
}

fn anonymizeManifest(
    allocator: std.mem.Allocator,
    root: std.json.Value,
    level: []const u8,
) ![]u8 {
    // Clone the value by stringifying and re-parsing so we can modify it
    var buf = std.ArrayList(u8).empty;
    defer buf.deinit(allocator);
    try writeJSONValue(buf.writer(allocator), root);
    const json_str = try buf.toOwnedSlice(allocator);
    defer allocator.free(json_str);
    var parsed_clone = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{});
    defer parsed_clone.deinit();
    var cloned = parsed_clone.value;

    if (cloned != .object) {
        // For non-objects, just re-serialize and return
        var out_buf = std.ArrayList(u8).empty;
        defer out_buf.deinit(allocator);
        try writeJSONValue(out_buf.writer(allocator), cloned);
        return out_buf.toOwnedSlice(allocator);
    }

    const obj = &cloned.object;

    // Anonymize metadata fields
    if (obj.get("metadata")) |meta| {
        if (meta == .object) {
            var meta_obj = meta.object;

            // Path anonymization: /nas/private/user/data → /datasets/data
            if (meta_obj.get("dataset_path")) |dp| {
                if (dp == .string) {
                    const anon_path = try anonymizePath(allocator, dp.string);
                    defer allocator.free(anon_path);
                    try meta_obj.put("dataset_path", std.json.Value{ .string = anon_path });
                }
            }

            // Anonymize other paths if full level
            if (std.mem.eql(u8, level, "full")) {
                const path_fields = [_][]const u8{ "code_path", "output_path", "checkpoint_path" };
                for (path_fields) |field| {
                    if (meta_obj.get(field)) |p| {
                        if (p == .string) {
                            const anon_path = try anonymizePath(allocator, p.string);
                            defer allocator.free(anon_path);
                            try meta_obj.put(field, std.json.Value{ .string = anon_path });
                        }
                    }
                }
            }
        }
    }

    // Anonymize system info
    if (obj.get("system")) |sys| {
        if (sys == .object) {
            var sys_obj = sys.object;

            // Hostname: gpu-server-01.internal → worker-A
            if (sys_obj.get("hostname")) |h| {
                if (h == .string) {
                    try sys_obj.put("hostname", std.json.Value{ .string = "worker-A" });
                }
            }

            // IP addresses → [REDACTED]
            if (sys_obj.get("ip_address")) |_| {
                try sys_obj.put("ip_address", std.json.Value{ .string = "[REDACTED]" });
            }

            // Username: user@lab.edu → researcher-N
            if (sys_obj.get("username")) |_| {
                try sys_obj.put("username", std.json.Value{ .string = "[REDACTED]" });
            }
        }
    }

    // Anonymize logs reference (logs may contain PII)
    if (std.mem.eql(u8, level, "full")) {
        _ = obj.swapRemove("logs");
        _ = obj.swapRemove("log_path");
        _ = obj.swapRemove("annotations");
    }

    // Serialize back to JSON
    var out_buf = std.ArrayList(u8).empty;
    defer out_buf.deinit(allocator);
    try writeJSONValue(out_buf.writer(allocator), cloned);
    return out_buf.toOwnedSlice(allocator);
}

fn anonymizePath(allocator: std.mem.Allocator, path: []const u8) ![]const u8 {
    // Simple path anonymization: replace leading path components with generic names
    // /home/user/project/data → /workspace/data
    // /nas/private/lab/experiments → /datasets/experiments

    // Find the last component
    const last_sep = std.mem.lastIndexOf(u8, path, "/");
    if (last_sep == null) return allocator.dupe(u8, path);

    const filename = path[last_sep.? + 1 ..];

    // Determine prefix based on context
    const prefix = if (std.mem.indexOf(u8, path, "data") != null)
        "/datasets"
    else if (std.mem.indexOf(u8, path, "model") != null or std.mem.indexOf(u8, path, "checkpoint") != null)
        "/models"
    else if (std.mem.indexOf(u8, path, "code") != null or std.mem.indexOf(u8, path, "src") != null)
        "/code"
    else
        "/workspace";

    return std.fs.path.join(allocator, &[_][]const u8{ prefix, filename });
}

fn printUsage() !void {
    std.debug.print("Usage: ml export <run-id|path> [options]\n", .{});
    std.debug.print("\nExport experiment for sharing or archiving:\n", .{});
    std.debug.print("\t--bundle <path>\t\tCreate tarball at path\n", .{});
    std.debug.print("\t--anonymize\t\tEnable anonymization\n", .{});
    std.debug.print("\t--anonymize-level <lvl>\t'metadata-only' or 'full'\n", .{});
    std.debug.print("\t--base <path>\t\tBase path to find run\n", .{});
    std.debug.print("\t--json\t\t\tOutput JSON response\n", .{});
    std.debug.print("\nAnonymization rules:\n", .{});
    std.debug.print("\t- Paths: /nas/private/... → /datasets/...\n", .{});
    std.debug.print("\t- Hostnames: gpu-server-01 → worker-A\n", .{});
    std.debug.print("\t- IPs: 192.168.1.100 → [REDACTED]\n", .{});
    std.debug.print("\t- Usernames: user@lab.edu → [REDACTED]\n", .{});
    std.debug.print("\t- Full level: Also removes logs and annotations\n", .{});
    std.debug.print("\nExamples:\n", .{});
    std.debug.print("\tml export run_abc --bundle run_abc.tar.gz\n", .{});
    std.debug.print("\tml export run_abc --bundle run_abc.tar.gz --anonymize\n", .{});
    std.debug.print("\tml export run_abc --anonymize --anonymize-level full\n", .{});
}

fn writeJSONValue(writer: anytype, v: std.json.Value) !void {
    switch (v) {
        .null => try writer.writeAll("null"),
        .bool => |b| try writer.print("{}", .{b}),
        .integer => |i| try writer.print("{d}", .{i}),
        .float => |f| try writer.print("{d}", .{f}),
        .string => |s| try writeJSONString(writer, s),
        .array => |arr| {
            try writer.writeAll("[");
            for (arr.items, 0..) |item, idx| {
                if (idx > 0) try writer.writeAll(",");
                try writeJSONValue(writer, item);
            }
            try writer.writeAll("]");
        },
        .object => |obj| {
            try writer.writeAll("{");
            var first = true;
            var it = obj.iterator();
            while (it.next()) |entry| {
                if (!first) try writer.writeAll(",");
                first = false;
                try writer.print("\"{s}\":", .{entry.key_ptr.*});
                try writeJSONValue(writer, entry.value_ptr.*);
            }
            try writer.writeAll("}");
        },
        .number_string => |s| try writer.print("{s}", .{s}),
    }
}

fn writeJSONString(writer: anytype, s: []const u8) !void {
    try writer.writeAll("\"");
    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.writeAll(&[_]u8{c});
                }
            },
        }
    }
    try writer.writeAll("\"");
}

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