fetch_ml/cli/src/commands/jupyter/lifecycle.zig

const std = @import("std");
const Config = @import("../../config.zig").Config;
const ws = @import("../../net/ws/client.zig");
const crypto = @import("../../utils/crypto.zig");
const protocol = @import("../../net/protocol.zig");
const validation = @import("validation.zig");

/// Context holding connection resources for cleanup
const ConnectionCtx = struct {
    config: Config,
    client: ws.Client,
    api_key_hash: []const u8,
    ws_url: []const u8,
    allocator: std.mem.Allocator,

    fn init(allocator: std.mem.Allocator) !ConnectionCtx {
        const config = try Config.load(allocator);
        errdefer {
            var mut = config;
            mut.deinit(allocator);
        }

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

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

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

        return ConnectionCtx{
            .config = config,
            .client = client,
            .api_key_hash = api_key_hash,
            .ws_url = ws_url,
            .allocator = allocator,
        };
    }

    fn deinit(self: *ConnectionCtx) void {
        self.allocator.free(self.api_key_hash);
        self.allocator.free(self.ws_url);
        self.client.close();
        var mut = self.config;
        mut.deinit(self.allocator);
    }
};

/// Create a new Jupyter workspace and start it
pub fn createJupyter(allocator: std.mem.Allocator, args: []const []const u8) !void {
    if (args.len < 1) {
        std.debug.print("Usage: ml jupyter create <name> [--path <path>] [--password <password>]\n", .{});
        return;
    }

    const name = args[0];
    var workspace_path_owned: ?[]u8 = null;
    defer if (workspace_path_owned) |p| allocator.free(p);
    var workspace_path: []const u8 = "";
    var password: []const u8 = "";

    var i: usize = 1;
    while (i < args.len) : (i += 1) {
        if (std.mem.eql(u8, args[i], "--path") and i + 1 < args.len) {
            workspace_path = args[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, args[i], "--password") and i + 1 < args.len) {
            password = args[i + 1];
            i += 1;
        }
    }

    if (workspace_path.len == 0) {
        const p = try defaultWorkspacePath(allocator, name);
        workspace_path_owned = p;
        workspace_path = p;
    }

    if (!validation.validateWorkspacePath(workspace_path)) {
        std.debug.print("Invalid workspace path\n", .{});
        return error.InvalidArgs;
    }

    std.fs.cwd().makePath(workspace_path) catch |err| {
        std.debug.print("Failed to create workspace directory: {}\n", .{err});
        return;
    };

    var start_args = std.ArrayList([]const u8).initCapacity(allocator, 8) catch |err| {
        std.debug.print("Failed to allocate args: {}\n", .{err});
        return;
    };
    defer start_args.deinit(allocator);

    try start_args.append(allocator, "--name");
    try start_args.append(allocator, name);
    try start_args.append(allocator, "--workspace");
    try start_args.append(allocator, workspace_path);
    if (password.len > 0) {
        try start_args.append(allocator, "--password");
        try start_args.append(allocator, password);
    }

    try startJupyter(allocator, start_args.items);
}

/// Start a Jupyter service
pub fn startJupyter(allocator: std.mem.Allocator, args: []const []const u8) !void {
    var name: []const u8 = "default";
    var workspace: []const u8 = "./workspace";
    var password: []const u8 = "";

    var i: usize = 0;
    while (i < args.len) : (i += 1) {
        if (std.mem.eql(u8, args[i], "--name") and i + 1 < args.len) {
            name = args[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, args[i], "--workspace") and i + 1 < args.len) {
            workspace = args[i + 1];
            i += 1;
        } else if (std.mem.eql(u8, args[i], "--password") and i + 1 < args.len) {
            password = args[i + 1];
            i += 1;
        }
    }

    var ctx = try ConnectionCtx.init(allocator);
    defer ctx.deinit();

    std.debug.print("Starting Jupyter service '{s}'...\n", .{name});

    ctx.client.sendStartJupyter(name, workspace, password, ctx.api_key_hash) catch |err| {
        std.debug.print("Failed to send start command: {}\n", .{err});
        return;
    };

    const response = ctx.client.receiveMessage(allocator) catch |err| {
        std.debug.print("Failed to receive response: {}\n", .{err});
        return;
    };
    defer allocator.free(response);

    const packet = protocol.ResponsePacket.deserialize(response, allocator) catch |err| {
        std.debug.print("Failed to parse response: {}\n", .{err});
        return;
    };
    defer packet.deinit(allocator);

    switch (packet.packet_type) {
        .success => {
            std.debug.print("Jupyter service started!\n", .{});
            if (packet.success_message) |msg| {
                std.debug.print("{s}\n", .{msg});
            }
        },
        .error_packet => {
            const error_msg = protocol.ResponsePacket.getErrorMessage(packet.error_code.?);
            std.debug.print("Failed to start service: {s}\n", .{error_msg});
            if (packet.error_details) |details| {
                std.debug.print("Details: {s}\n", .{details});
            } else if (packet.error_message) |msg| {
                std.debug.print("Details: {s}\n", .{msg});
            }
        },
        else => {
            std.debug.print("Unexpected response type\n", .{});
        },
    }
}

/// Stop a Jupyter service
pub fn stopJupyter(allocator: std.mem.Allocator, args: []const []const u8) !void {
    if (args.len < 1) {
        std.debug.print("Usage: ml jupyter stop <service_id>\n", .{});
        return;
    }
    const service_id = args[0];

    var ctx = try ConnectionCtx.init(allocator);
    defer ctx.deinit();

    std.debug.print("Stopping service {s}...\n", .{service_id});

    ctx.client.sendStopJupyter(service_id, ctx.api_key_hash) catch |err| {
        std.debug.print("Failed to send stop command: {}\n", .{err});
        return;
    };

    const response = ctx.client.receiveMessage(allocator) catch |err| {
        std.debug.print("Failed to receive response: {}\n", .{err});
        return;
    };
    defer allocator.free(response);

    const packet = protocol.ResponsePacket.deserialize(response, allocator) catch |err| {
        std.debug.print("Failed to parse response: {}\n", .{err});
        return;
    };
    defer packet.deinit(allocator);

    switch (packet.packet_type) {
        .success => {
            std.debug.print("Service stopped.\n", .{});
        },
        .error_packet => {
            const error_msg = protocol.ResponsePacket.getErrorMessage(packet.error_code.?);
            std.debug.print("Failed to stop service: {s}\n", .{error_msg});
            if (packet.error_details) |details| {
                std.debug.print("Details: {s}\n", .{details});
            } else if (packet.error_message) |msg| {
                std.debug.print("Details: {s}\n", .{msg});
            }
        },
        else => {
            std.debug.print("Unexpected response type\n", .{});
        },
    }
}

/// Remove a Jupyter service
pub fn removeJupyter(allocator: std.mem.Allocator, args: []const []const u8) !void {
    if (args.len < 1) {
        std.debug.print("Usage: ml jupyter remove <service_id> [--purge] [--force]\n", .{});
        return;
    }

    const service_id = args[0];
    var purge: bool = false;
    var force: bool = false;

    var i: usize = 1;
    while (i < args.len) : (i += 1) {
        if (std.mem.eql(u8, args[i], "--purge")) {
            purge = true;
        } else if (std.mem.eql(u8, args[i], "--force")) {
            force = true;
        } else {
            std.debug.print("Unknown option: {s}\n", .{args[i]});
            std.debug.print("Usage: ml jupyter remove <service_id> [--purge] [--force]\n", .{});
            return error.InvalidArgs;
        }
    }

    // Trash-first by default: no confirmation.
    // Permanent deletion requires explicit --purge and a strong confirmation unless --force.
    if (purge and !force) {
        std.debug.print("PERMANENT deletion requested for '{s}'.\n", .{service_id});
        std.debug.print("This cannot be undone.\n", .{});
        std.debug.print("Type the service name to confirm: ", .{});

        const stdin = std.fs.File{ .handle = @intCast(0) };
        var buffer: [256]u8 = undefined;
        const bytes_read = stdin.read(&buffer) catch |err| {
            std.debug.print("Failed to read input: {}\n", .{err});
            return;
        };
        const line = buffer[0..bytes_read];
        const typed = std.mem.trim(u8, line, "\n\r ");
        if (!std.mem.eql(u8, typed, service_id)) {
            std.debug.print("Operation cancelled.\n", .{});
            return;
        }
    }

    var ctx = try ConnectionCtx.init(allocator);
    defer ctx.deinit();

    if (purge) {
        std.debug.print("Permanently deleting service {s}...\n", .{service_id});
    } else {
        std.debug.print("Removing service {s} (move to trash)...\n", .{service_id});
    }

    ctx.client.sendRemoveJupyter(service_id, ctx.api_key_hash, purge) catch |err| {
        std.debug.print("Failed to send remove command: {}\n", .{err});
        return;
    };

    const response = ctx.client.receiveMessage(allocator) catch |err| {
        std.debug.print("Failed to receive response: {}\n", .{err});
        return;
    };
    defer allocator.free(response);

    const packet = protocol.ResponsePacket.deserialize(response, allocator) catch |err| {
        std.debug.print("Failed to parse response: {}\n", .{err});
        return;
    };
    defer packet.deinit(allocator);

    switch (packet.packet_type) {
        .success => {
            std.debug.print("Service removed successfully.\n", .{});
        },
        .error_packet => {
            const error_msg = protocol.ResponsePacket.getErrorMessage(packet.error_code.?);
            std.debug.print("Failed to remove service: {s}\n", .{error_msg});
            if (packet.error_details) |details| {
                std.debug.print("Details: {s}\n", .{details});
            } else if (packet.error_message) |msg| {
                std.debug.print("Details: {s}\n", .{msg});
            }
        },
        else => {
            std.debug.print("Unexpected response type\n", .{});
        },
    }
}

/// Restore a Jupyter workspace
pub fn restoreJupyter(allocator: std.mem.Allocator, args: []const []const u8, json: bool) !void {
    _ = json;
    if (args.len < 1) {
        std.debug.print("Usage: ml jupyter restore <name>\n", .{});
        return;
    }
    const name = args[0];

    var ctx = try ConnectionCtx.init(allocator);
    defer ctx.deinit();

    std.debug.print("Restoring workspace {s}...", .{name});

    ctx.client.sendRestoreJupyter(name, ctx.api_key_hash) catch {
        std.debug.print("Failed to send restore command\n", .{});
        return;
    };

    const response = ctx.client.receiveMessage(allocator) catch |err| {
        std.debug.print("Failed to receive response: {}\n", .{err});
        return;
    };
    defer allocator.free(response);

    const packet = protocol.ResponsePacket.deserialize(response, allocator) catch |err| {
        std.debug.print("Failed to parse response: {}\n", .{err});
        return;
    };
    defer packet.deinit(allocator);

    switch (packet.packet_type) {
        .success => {
            if (packet.success_message) |msg| {
                std.debug.print("{s}", .{msg});
            } else {
                std.debug.print("Workspace restored.", .{});
            }
        },
        .error_packet => {
            const error_msg = protocol.ResponsePacket.getErrorMessage(packet.error_code.?);
            std.debug.print("Error: {s}\n", .{error_msg});
        },
        else => {
            std.debug.print("Unexpected response type\n", .{});
        },
    }
}

/// Default workspace path generator
fn defaultWorkspacePath(allocator: std.mem.Allocator, name: []const u8) ![]u8 {
    return std.fmt.allocPrint(allocator, "./{s}", .{name});
}