const std = @import("std");
const config = @import("../config.zig");
const db = @import("../db.zig");
const ws = @import("../net/ws/client.zig");
const crypto = @import("../utils/crypto.zig");
const core = @import("../core.zig");
const mode = @import("../mode.zig");
const manifest_lib = @import("../manifest.zig");

pub fn run(allocator: std.mem.Allocator, args: []const []const u8) !void {
    var flags = core.flags.CommonFlags{};
    var force = false;
    var targets: std.ArrayList([]const u8) = .empty;
    defer targets.deinit(allocator);

    // Parse arguments
    var i: usize = 0;
    while (i < args.len) : (i += 1) {
        const arg = args[i];
        if (std.mem.eql(u8, arg, "--force")) {
            force = true;
        } 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 if (std.mem.startsWith(u8, arg, "--")) {
            core.output.err("Unknown option");
            return error.InvalidArgs;
        } else {
            try targets.append(allocator, arg);
        }
    }

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

    if (targets.items.len == 0) {
        core.output.err("No run_id specified");
        return error.InvalidArgs;
    }

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

    // Detect mode
    const mode_result = try mode.detect(allocator, cfg);
    if (mode_result.warning) |w| {
        std.log.warn("{s}", .{w});
    }

    var success_count: usize = 0;
    var failed_count: usize = 0;

    for (targets.items) |target| {
        if (mode.isOffline(mode_result.mode)) {
            // Local mode: kill by PID
            cancelLocal(allocator, target, force, flags.json) catch |err| {
                if (!flags.json) {
                    std.debug.print("Failed to cancel '{s}': {}\n", .{ target, err });
                }
                failed_count += 1;
                continue;
            };
        } else {
            // Online mode: cancel on server
            cancelServer(allocator, target, force, flags.json, cfg) catch |err| {
                if (!flags.json) {
                    std.debug.print("Failed to cancel '{s}': {}\n", .{ target, err });
                }
                failed_count += 1;
                continue;
            };
        }
        success_count += 1;
    }

    if (flags.json) {
        std.debug.print("{{\"success\":true,\"canceled\":{d},\"failed\":{d}}}\n", .{ success_count, failed_count });
    } else {
        std.debug.print("Canceled {d} run(s)\n", .{success_count});
        if (failed_count > 0) {
            std.debug.print("Failed to cancel {d} run(s)\n", .{failed_count});
        }
    }
}

/// Cancel local run by PID
fn cancelLocal(allocator: std.mem.Allocator, run_id: []const u8, force: bool, json: bool) !void {
    const cfg = try config.Config.load(allocator);
    defer {
        var mut_cfg = cfg;
        mut_cfg.deinit(allocator);
    }

    // Get DB path
    const db_path = try cfg.getDBPath(allocator);
    defer allocator.free(db_path);

    var database = try db.DB.init(allocator, db_path);
    defer database.close();

    // Look up PID
    const sql = "SELECT pid FROM ml_runs WHERE run_id = ? AND status = 'RUNNING';";
    const stmt = try database.prepare(sql);
    defer db.DB.finalize(stmt);
    try db.DB.bindText(stmt, 1, run_id);

    const has_row = try db.DB.step(stmt);
    if (!has_row) {
        return error.RunNotFoundOrNotRunning;
    }

    const pid = db.DB.columnInt64(stmt, 0);
    if (pid == 0) {
        return error.NoPIDAvailable;
    }

    // Send SIGTERM first
    std.posix.kill(@intCast(pid), std.posix.SIG.TERM) catch |err| {
        if (err == error.ProcessNotFound) {
            // Process already gone
        } else {
            return err;
        }
    };

    if (!force) {
        // Wait 5 seconds for graceful termination
        std.Thread.sleep(5 * std.time.ns_per_s);
    }

    // Check if still running, send SIGKILL if needed
    if (force or isProcessRunning(@intCast(pid))) {
        std.posix.kill(@intCast(pid), std.posix.SIG.KILL) catch |err| {
            if (err != error.ProcessNotFound) {
                return err;
            }
        };
    }

    // Update run status
    const update_sql = "UPDATE ml_runs SET status = 'CANCELLED', pid = NULL WHERE run_id = ?;";
    const update_stmt = try database.prepare(update_sql);
    defer db.DB.finalize(update_stmt);
    try db.DB.bindText(update_stmt, 1, run_id);
    _ = try db.DB.step(update_stmt);

    // Update manifest
    const artifact_path = try std.fs.path.join(allocator, &[_][]const u8{
        cfg.artifact_path,
        if (cfg.experiment) |exp| exp.name else "default",
        run_id,
        "run_manifest.json",
    });
    defer allocator.free(artifact_path);

    manifest_lib.updateManifestStatus(artifact_path, "CANCELLED", null, allocator) catch {};

    // Checkpoint
    database.checkpointOnExit();

    if (!json) {
        std.debug.print("Canceled run {s}\n", .{run_id[0..8]});
    }
}

/// Check if process is still running
fn isProcessRunning(pid: i32) bool {
    const result = std.posix.kill(pid, 0);
    return if (result) |_| true else |err| err == error.PermissionDenied;
}

/// Cancel server job
fn cancelServer(allocator: std.mem.Allocator, job_name: []const u8, force: bool, json: bool, cfg: config.Config) !void {
    _ = force;
    _ = json;

    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.sendCancelJob(job_name, api_key_hash);

    // Wait for acknowledgment
    const message = try client.receiveMessage(allocator);
    defer allocator.free(message);

    // Parse response (simplified)
    if (std.mem.indexOf(u8, message, "error") != null) {
        return error.ServerCancelFailed;
    }
}

fn printUsage() !void {
    std.debug.print("Usage: ml cancel [options] <run-id> [<run-id> ...]\n", .{});
    std.debug.print("\nCancel a local run (kill process) or server job.\n\n", .{});
    std.debug.print("Options:\n", .{});
    std.debug.print("\t--force\t\tForce cancel (SIGKILL immediately)\n", .{});
    std.debug.print("\t--json\t\tOutput structured JSON\n", .{});
    std.debug.print("\t--help, -h\tShow this help message\n", .{});
    std.debug.print("\nExamples:\n", .{});
    std.debug.print("\tml cancel abc123\t# Cancel local run by run_id\n", .{});
    std.debug.print("\tml cancel --force abc123\t# Force cancel\n", .{});
}
