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

/// Logs command - fetch and display job logs via WebSocket API
pub fn run(allocator: std.mem.Allocator, argv: []const []const u8) !void {
    if (argv.len == 0) {
        try printUsage();
        return error.InvalidArgs;
    }
    if (std.mem.eql(u8, argv[0], "--help") or std.mem.eql(u8, argv[0], "-h")) {
        try printUsage();
        return;
    }

    const target = argv[0];

    // Parse optional flags
    var follow = false;
    var tail: ?usize = null;

    var i: usize = 1;
    while (i < argv.len) : (i += 1) {
        const a = argv[i];
        if (std.mem.eql(u8, a, "-f") or std.mem.eql(u8, a, "--follow")) {
            follow = true;
        } else if (std.mem.eql(u8, a, "-n") and i + 1 < argv.len) {
            tail = try std.fmt.parseInt(usize, argv[i + 1], 10);
            i += 1;
        } else if (std.mem.eql(u8, a, "--tail") and i + 1 < argv.len) {
            tail = try std.fmt.parseInt(usize, argv[i + 1], 10);
            i += 1;
        } else {
            colors.printError("Unknown option: {s}\n", .{a});
            return error.InvalidArgs;
        }
    }

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

    colors.printInfo("Fetching logs for: {s}\n", .{target});

    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();

    // Send appropriate request based on follow flag
    if (follow) {
        try client.sendStreamLogs(target, api_key_hash);
    } else {
        try client.sendGetLogs(target, api_key_hash);
    }

    // Receive and display response
    const message = try client.receiveMessage(allocator);
    defer allocator.free(message);

    const packet = protocol.ResponsePacket.deserialize(message, allocator) catch {
        // Fallback: treat as plain text response
        std.debug.print("{s}\n", .{message});
        return;
    };
    defer {
        if (packet.success_message) |m| allocator.free(m);
        if (packet.error_message) |m| allocator.free(m);
        if (packet.error_details) |m| allocator.free(m);
        if (packet.data_payload) |m| allocator.free(m);
        if (packet.data_type) |m| allocator.free(m);
    }

    switch (packet.packet_type) {
        .data => {
            if (packet.data_payload) |payload| {
                // Parse JSON response
                const parsed = std.json.parseFromSlice(std.json.Value, allocator, payload, .{}) catch {
                    std.debug.print("{s}\n", .{payload});
                    return;
                };
                defer parsed.deinit();

                const root = parsed.value.object;

                // Display logs
                if (root.get("logs")) |logs| {
                    if (logs == .string) {
                        std.debug.print("{s}\n", .{logs.string});
                    }
                } else if (root.get("message")) |msg| {
                    if (msg == .string) {
                        colors.printInfo("{s}\n", .{msg.string});
                    }
                }

                // Show truncation warning if applicable
                if (root.get("truncated")) |truncated| {
                    if (truncated == .bool and truncated.bool) {
                        if (root.get("total_lines")) |total| {
                            if (total == .integer) {
                                colors.printWarning("\n[Output truncated. Total lines: {d}]\n", .{total.integer});
                            }
                        }
                    }
                }
            }
        },
        .error_packet => {
            const err_msg = packet.error_message orelse "Unknown error";
            colors.printError("Error: {s}\n", .{err_msg});
            return error.ServerError;
        },
        else => {
            if (packet.success_message) |msg| {
                colors.printSuccess("{s}\n", .{msg});
            } else {
                colors.printInfo("Logs retrieved successfully\n", .{});
            }
        },
    }
}

fn printUsage() !void {
    colors.printInfo("Usage:\n", .{});
    colors.printInfo("  ml logs <task_id|run_id|experiment_id> [-f|--follow] [-n <count>|--tail <count>]\n", .{});
    colors.printInfo("\nExamples:\n", .{});
    colors.printInfo("  ml logs abc123              # Show full logs\n", .{});
    colors.printInfo("  ml logs abc123 -f             # Follow logs in real-time\n", .{});
    colors.printInfo("  ml logs abc123 -n 100         # Show last 100 lines\n", .{});
}