const std = @import("std");

/// RunManifest represents a run manifest - identical schema between local and server
/// Schema compatibility is a hard requirement enforced here
pub const RunManifest = struct {
    run_id: []const u8,
    experiment: []const u8,
    command: []const u8,
    args: [][]const u8,
    commit_id: ?[]const u8,
    started_at: []const u8,
    ended_at: ?[]const u8,
    status: []const u8, // RUNNING, FINISHED, FAILED, CANCELLED
    exit_code: ?i32,
    params: std.StringHashMap([]const u8),
    metrics_summary: ?std.StringHashMap(f64),
    artifact_path: []const u8,
    synced: bool,

    pub fn init(allocator: std.mem.Allocator) RunManifest {
        return .{
            .run_id = "",
            .experiment = "",
            .command = "",
            .args = &[_][]const u8{},
            .commit_id = null,
            .started_at = "",
            .ended_at = null,
            .status = "RUNNING",
            .exit_code = null,
            .params = std.StringHashMap([]const u8).init(allocator),
            .metrics_summary = null,
            .artifact_path = "",
            .synced = false,
        };
    }

    pub fn deinit(self: *RunManifest, allocator: std.mem.Allocator) void {
        var params_iter = self.params.iterator();
        while (params_iter.next()) |entry| {
            allocator.free(entry.key_ptr.*);
            allocator.free(entry.value_ptr.*);
        }
        self.params.deinit();

        if (self.metrics_summary) |*summary| {
            var summary_iter = summary.iterator();
            while (summary_iter.next()) |entry| {
                allocator.free(entry.key_ptr.*);
            }
            summary.deinit();
        }

        for (self.args) |arg| {
            allocator.free(arg);
        }
        allocator.free(self.args);
    }
};

/// Write manifest to JSON file
pub fn writeManifest(manifest: RunManifest, path: []const u8, allocator: std.mem.Allocator) !void {
    var file = try std.fs.cwd().createFile(path, .{});
    defer file.close();

    // Write JSON manually to avoid std.json complexity with hash maps
    try file.writeAll("{\n");

    const line1 = try std.fmt.allocPrint(allocator, "  \"run_id\": \"{s}\",\n", .{manifest.run_id});
    defer allocator.free(line1);
    try file.writeAll(line1);

    const line2 = try std.fmt.allocPrint(allocator, "  \"experiment\": \"{s}\",\n", .{manifest.experiment});
    defer allocator.free(line2);
    try file.writeAll(line2);

    const line3 = try std.fmt.allocPrint(allocator, "  \"command\": \"{s}\",\n", .{manifest.command});
    defer allocator.free(line3);
    try file.writeAll(line3);

    // Args array
    try file.writeAll("  \"args\": [");
    for (manifest.args, 0..) |arg, i| {
        if (i > 0) try file.writeAll(", ");
        const arg_str = try std.fmt.allocPrint(allocator, "\"{s}\"", .{arg});
        defer allocator.free(arg_str);
        try file.writeAll(arg_str);
    }
    try file.writeAll("],\n");

    // Commit ID (optional)
    if (manifest.commit_id) |cid| {
        const cid_str = try std.fmt.allocPrint(allocator, "  \"commit_id\": \"{s}\",\n", .{cid});
        defer allocator.free(cid_str);
        try file.writeAll(cid_str);
    } else {
        try file.writeAll("  \"commit_id\": null,\n");
    }

    const started_str = try std.fmt.allocPrint(allocator, "  \"started_at\": \"{s}\",\n", .{manifest.started_at});
    defer allocator.free(started_str);
    try file.writeAll(started_str);

    // Ended at (optional)
    if (manifest.ended_at) |ended| {
        const ended_str = try std.fmt.allocPrint(allocator, "  \"ended_at\": \"{s}\",\n", .{ended});
        defer allocator.free(ended_str);
        try file.writeAll(ended_str);
    } else {
        try file.writeAll("  \"ended_at\": null,\n");
    }

    const status_str = try std.fmt.allocPrint(allocator, "  \"status\": \"{s}\",\n", .{manifest.status});
    defer allocator.free(status_str);
    try file.writeAll(status_str);

    // Exit code (optional)
    if (manifest.exit_code) |code| {
        const exit_str = try std.fmt.allocPrint(allocator, "  \"exit_code\": {d},\n", .{code});
        defer allocator.free(exit_str);
        try file.writeAll(exit_str);
    } else {
        try file.writeAll("  \"exit_code\": null,\n");
    }

    // Params object
    try file.writeAll("  \"params\": {");
    var params_first = true;
    var params_iter = manifest.params.iterator();
    while (params_iter.next()) |entry| {
        if (!params_first) try file.writeAll(", ");
        params_first = false;
        const param_str = try std.fmt.allocPrint(allocator, "\"{s}\": \"{s}\"", .{ entry.key_ptr.*, entry.value_ptr.* });
        defer allocator.free(param_str);
        try file.writeAll(param_str);
    }
    try file.writeAll("},\n");

    // Metrics summary (optional)
    if (manifest.metrics_summary) |summary| {
        try file.writeAll("  \"metrics_summary\": {");
        var summary_first = true;
        var summary_iter = summary.iterator();
        while (summary_iter.next()) |entry| {
            if (!summary_first) try file.writeAll(", ");
            summary_first = false;
            const metric_str = try std.fmt.allocPrint(allocator, "\"{s}\": {d:.4}", .{ entry.key_ptr.*, entry.value_ptr.* });
            defer allocator.free(metric_str);
            try file.writeAll(metric_str);
        }
        try file.writeAll("},\n");
    } else {
        try file.writeAll("  \"metrics_summary\": null,\n");
    }

    const artifact_str = try std.fmt.allocPrint(allocator, "  \"artifact_path\": \"{s}\",\n", .{manifest.artifact_path});
    defer allocator.free(artifact_str);
    try file.writeAll(artifact_str);

    const synced_str = try std.fmt.allocPrint(allocator, "  \"synced\": {}", .{manifest.synced});
    defer allocator.free(synced_str);
    try file.writeAll(synced_str);

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

/// Read manifest from JSON file
pub fn readManifest(path: []const u8, allocator: std.mem.Allocator) !RunManifest {
    var file = try std.fs.cwd().openFile(path, .{});
    defer file.close();

    const content = try file.readToEndAlloc(allocator, 1024 * 1024);
    defer allocator.free(content);

    const parsed = try std.json.parseFromSlice(std.json.Value, allocator, content, .{});
    defer parsed.deinit();

    if (parsed.value != .object) {
        return error.InvalidManifest;
    }

    const root = parsed.value.object;
    var manifest = RunManifest.init(allocator);

    // Required fields
    manifest.run_id = try getStringField(allocator, root, "run_id") orelse return error.MissingRunId;
    manifest.experiment = try getStringField(allocator, root, "experiment") orelse return error.MissingExperiment;
    manifest.command = try getStringField(allocator, root, "command") orelse return error.MissingCommand;
    manifest.status = try getStringField(allocator, root, "status") orelse "RUNNING";
    manifest.started_at = try getStringField(allocator, root, "started_at") orelse "";

    // Optional fields
    manifest.ended_at = try getStringField(allocator, root, "ended_at");
    manifest.commit_id = try getStringField(allocator, root, "commit_id");
    manifest.artifact_path = try getStringField(allocator, root, "artifact_path") orelse "";

    // Synced boolean
    if (root.get("synced")) |synced_val| {
        if (synced_val == .bool) {
            manifest.synced = synced_val.bool;
        }
    }

    // Exit code
    if (root.get("exit_code")) |exit_val| {
        if (exit_val == .integer) {
            manifest.exit_code = @intCast(exit_val.integer);
        }
    }

    // Args array
    if (root.get("args")) |args_val| {
        if (args_val == .array) {
            const args = try allocator.alloc([]const u8, args_val.array.items.len);
            for (args_val.array.items, 0..) |arg, i| {
                if (arg == .string) {
                    args[i] = try allocator.dupe(u8, arg.string);
                }
            }
            manifest.args = args;
        }
    }

    // Params object
    if (root.get("params")) |params_val| {
        if (params_val == .object) {
            var params_iter = params_val.object.iterator();
            while (params_iter.next()) |entry| {
                if (entry.value_ptr.* == .string) {
                    const key = try allocator.dupe(u8, entry.key_ptr.*);
                    const value = try allocator.dupe(u8, entry.value_ptr.*.string);
                    try manifest.params.put(key, value);
                }
            }
        }
    }

    // Metrics summary
    if (root.get("metrics_summary")) |metrics_val| {
        if (metrics_val == .object) {
            var summary = std.StringHashMap(f64).init(allocator);
            var metrics_iter = metrics_val.object.iterator();
            while (metrics_iter.next()) |entry| {
                const val = entry.value_ptr.*;
                if (val == .float) {
                    const key = try allocator.dupe(u8, entry.key_ptr.*);
                    try summary.put(key, val.float);
                } else if (val == .integer) {
                    const key = try allocator.dupe(u8, entry.key_ptr.*);
                    try summary.put(key, @floatFromInt(val.integer));
                }
            }
            manifest.metrics_summary = summary;
        }
    }

    return manifest;
}

/// Get string field from JSON object, duplicating the string
fn getStringField(allocator: std.mem.Allocator, obj: std.json.ObjectMap, field: []const u8) !?[]const u8 {
    const val = obj.get(field) orelse return null;
    if (val != .string) return null;
    return try allocator.dupe(u8, val.string);
}

/// Update manifest status and ended_at on run completion
pub fn updateManifestStatus(path: []const u8, status: []const u8, exit_code: ?i32, allocator: std.mem.Allocator) !void {
    var manifest = try readManifest(path, allocator);
    defer manifest.deinit(allocator);

    manifest.status = status;
    manifest.exit_code = exit_code;

    // Set ended_at to current timestamp
    const now = std.time.timestamp();
    const epoch_seconds = std.time.epoch.EpochSeconds{ .secs = @intCast(now) };
    const epoch_day = epoch_seconds.getEpochDay();
    const year_day = epoch_day.calculateYearDay();
    const month_day = year_day.calculateMonthDay();
    const day_seconds = epoch_seconds.getDaySeconds();

    var buf: [30]u8 = undefined;
    const timestamp = std.fmt.bufPrint(&buf, "{d:0>4}-{d:0>2}-{d:0>2}T{d:0>2}:{d:0>2}:{d:0>2}Z", .{
        year_day.year,
        month_day.month.numeric(),
        month_day.day_index + 1,
        day_seconds.getHoursIntoDay(),
        day_seconds.getMinutesIntoHour(),
        day_seconds.getSecondsIntoMinute(),
    }) catch unreachable;

    manifest.ended_at = try allocator.dupe(u8, timestamp);

    try writeManifest(manifest, path, allocator);
}

/// Mark manifest as synced
pub fn markManifestSynced(path: []const u8, allocator: std.mem.Allocator) !void {
    var manifest = try readManifest(path, allocator);
    defer manifest.deinit(allocator);

    manifest.synced = true;
    try writeManifest(manifest, path, allocator);
}

/// Build manifest path from experiment and run_id
pub fn buildManifestPath(artifact_path: []const u8, experiment: []const u8, run_id: []const u8, allocator: std.mem.Allocator) ![]const u8 {
    return std.fs.path.join(allocator, &[_][]const u8{
        artifact_path,
        experiment,
        run_id,
        "run_manifest.json",
    });
}

/// Resolve manifest path from input (path, run_id, or task_id)
pub fn resolveManifestPath(input: []const u8, base_path: ?[]const u8, allocator: std.mem.Allocator) ![]const u8 {
    // If input is a valid file path, use it directly
    if (std.fs.path.isAbsolute(input)) {
        if (std.fs.cwd().access(input, .{})) {
            // It's a file or directory
            const stat = std.fs.cwd().statFile(input) catch {
                // It's a directory, append manifest name
                return std.fs.path.join(allocator, &[_][]const u8{ input, "run_manifest.json" });
            };
            _ = stat;
            // It's a file, use as-is
            return try allocator.dupe(u8, input);
        } else |_| {}
    }

    // Try relative path
    if (std.fs.cwd().access(input, .{})) {
        const stat = std.fs.cwd().statFile(input) catch {
            return std.fs.path.join(allocator, &[_][]const u8{ input, "run_manifest.json" });
        };
        _ = stat;
        return try allocator.dupe(u8, input);
    } else |_| {}

    // Search by run_id in base_path
    if (base_path) |bp| {
        return try findManifestById(bp, input, allocator);
    }

    return error.ManifestNotFound;
}

/// Find manifest by run_id in base path
fn findManifestById(base_path: []const u8, id: []const u8, allocator: std.mem.Allocator) ![]const u8 {
    // Look in experiments/ subdirectories
    var experiments_dir = std.fs.cwd().openDir(base_path, .{ .iterate = true }) catch {
        return error.ManifestNotFound;
    };
    defer experiments_dir.close();

    var iter = experiments_dir.iterate();
    while (try iter.next()) |entry| {
        if (entry.kind != .directory) continue;

        // Check if this experiment has a subdirectory matching the run_id
        const run_dir_path = try std.fs.path.join(allocator, &[_][]const u8{
            base_path,
            entry.name,
            id,
        });
        defer allocator.free(run_dir_path);

        const manifest_path = try std.fs.path.join(allocator, &[_][]const u8{
            run_dir_path,
            "run_manifest.json",
        });

        if (std.fs.cwd().access(manifest_path, .{})) {
            return manifest_path;
        } else |_| {
            allocator.free(manifest_path);
        }
    }

    return error.ManifestNotFound;
}