const std = @import("std");

/// Response packet types for structured server responses
pub const PacketType = enum(u8) {
    success = 0x00,
    error_packet = 0x01,
    progress = 0x02,
    status = 0x03,
    data = 0x04,
    log = 0x05,
};

/// Error codes for structured error responses
pub const ErrorCode = enum(u8) {
    // General errors (0x00-0x0F)
    unknown_error = 0x00,
    invalid_request = 0x01,
    authentication_failed = 0x02,
    permission_denied = 0x03,
    resource_not_found = 0x04,
    resource_already_exists = 0x05,

    // Server errors (0x10-0x1F)
    server_overloaded = 0x10,
    database_error = 0x11,
    network_error = 0x12,
    storage_error = 0x13,
    timeout = 0x14,

    // Job errors (0x20-0x2F)
    job_not_found = 0x20,
    job_already_running = 0x21,
    job_failed_to_start = 0x22,
    job_execution_failed = 0x23,
    job_cancelled = 0x24,

    // System errors (0x30-0x3F)
    out_of_memory = 0x30,
    disk_full = 0x31,
    invalid_configuration = 0x32,
    service_unavailable = 0x33,
};

/// Progress update types
pub const ProgressType = enum(u8) {
    percentage = 0x00,
    stage = 0x01,
    message = 0x02,
    bytes_transferred = 0x03,
};

/// Base response packet structure
pub const ResponsePacket = struct {
    packet_type: PacketType,
    timestamp: u64, // Unix timestamp

    // Success packet fields
    success_message: ?[]const u8 = null,

    // Error packet fields
    error_code: ?ErrorCode = null,
    error_message: ?[]const u8 = null,
    error_details: ?[]const u8 = null,

    // Progress packet fields
    progress_type: ?ProgressType = null,
    progress_value: ?u32 = null,
    progress_total: ?u32 = null,
    progress_message: ?[]const u8 = null,

    // Status packet fields
    status_data: ?[]const u8 = null,

    // Data packet fields
    data_type: ?[]const u8 = null,
    data_payload: ?[]const u8 = null,

    // Log packet fields
    log_level: ?u8 = null,
    log_message: ?[]const u8 = null,

    pub fn initSuccess(timestamp: u64, message: ?[]const u8) ResponsePacket {
        return ResponsePacket{
            .packet_type = .success,
            .timestamp = timestamp,
            .success_message = message,
        };
    }

    pub fn initError(timestamp: u64, code: ErrorCode, message: []const u8, details: ?[]const u8) ResponsePacket {
        return ResponsePacket{
            .packet_type = .error_packet,
            .timestamp = timestamp,
            .error_code = code,
            .error_message = message,
            .error_details = details,
        };
    }

    pub fn initProgress(timestamp: u64, ptype: ProgressType, value: u32, total: ?u32, message: ?[]const u8) ResponsePacket {
        return ResponsePacket{
            .packet_type = .progress,
            .timestamp = timestamp,
            .progress_type = ptype,
            .progress_value = value,
            .progress_total = total,
            .progress_message = message,
        };
    }

    pub fn initStatus(timestamp: u64, data: []const u8) ResponsePacket {
        return ResponsePacket{
            .packet_type = .status,
            .timestamp = timestamp,
            .status_data = data,
        };
    }

    pub fn initData(timestamp: u64, data_type: []const u8, payload: []const u8) ResponsePacket {
        return ResponsePacket{
            .packet_type = .data,
            .timestamp = timestamp,
            .data_type = data_type,
            .data_payload = payload,
        };
    }

    pub fn initLog(timestamp: u64, level: u8, message: []const u8) ResponsePacket {
        return ResponsePacket{
            .packet_type = .log,
            .timestamp = timestamp,
            .log_level = level,
            .log_message = message,
        };
    }

    /// Serialize packet to binary format
    pub fn serialize(self: ResponsePacket, allocator: std.mem.Allocator) ![]u8 {
        var buffer = try std.ArrayList(u8).initCapacity(allocator, 256);
        defer buffer.deinit(allocator);

        try buffer.append(allocator, @intFromEnum(self.packet_type));
        var ts_bytes: [8]u8 = undefined;
        std.mem.writeInt(u64, ts_bytes[0..8], self.timestamp, .big);
        try buffer.appendSlice(allocator, &ts_bytes);

        switch (self.packet_type) {
            .success => {
                if (self.success_message) |msg| {
                    try writeString(&buffer, allocator, msg);
                } else {
                    try writeString(&buffer, allocator, "");
                }
            },
            .error_packet => {
                try buffer.append(allocator, @intFromEnum(self.error_code.?));
                try writeString(&buffer, allocator, self.error_message.?);
                if (self.error_details) |details| {
                    try writeString(&buffer, allocator, details);
                } else {
                    try writeString(&buffer, allocator, "");
                }
            },
            .progress => {
                try buffer.append(allocator, @intFromEnum(self.progress_type.?));
                var pv_bytes: [4]u8 = undefined;
                std.mem.writeInt(u32, pv_bytes[0..4], self.progress_value.?, .big);
                try buffer.appendSlice(allocator, &pv_bytes);
                if (self.progress_total) |total| {
                    var pt_bytes: [4]u8 = undefined;
                    std.mem.writeInt(u32, pt_bytes[0..4], total, .big);
                    try buffer.appendSlice(allocator, &pt_bytes);
                } else {
                    try buffer.appendSlice(allocator, &[4]u8{ 0, 0, 0, 0 }); // 0 indicates no total
                }
                if (self.progress_message) |msg| {
                    try writeString(&buffer, allocator, msg);
                } else {
                    try writeString(&buffer, allocator, "");
                }
            },
            .status => {
                try writeString(&buffer, allocator, self.status_data.?);
            },
            .data => {
                try writeString(&buffer, allocator, self.data_type.?);
                try writeBytes(&buffer, allocator, self.data_payload.?);
            },
            .log => {
                try buffer.append(allocator, self.log_level.?);
                try writeString(&buffer, allocator, self.log_message.?);
            },
        }

        return buffer.toOwnedSlice(allocator);
    }

    /// Deserialize packet from binary format
    pub fn deserialize(data: []const u8, allocator: std.mem.Allocator) !ResponsePacket {
        if (data.len < 9) return error.InvalidPacket; // packet_type + timestamp

        var offset: usize = 0;
        const packet_type = @as(PacketType, @enumFromInt(data[offset]));
        offset += 1;

        const timestamp = std.mem.readInt(u64, data[offset .. offset + 8][0..8], .big);
        offset += 8;

        switch (packet_type) {
            .success => {
                const message = try readString(data, &offset, allocator);
                return ResponsePacket.initSuccess(timestamp, message);
            },
            .error_packet => {
                if (offset >= data.len) return error.InvalidPacket;
                const error_code = @as(ErrorCode, @enumFromInt(data[offset]));
                offset += 1;

                const error_message = try readString(data, &offset, allocator);
                const error_details = try readString(data, &offset, allocator);

                return ResponsePacket.initError(timestamp, error_code, error_message, if (error_details.len > 0) error_details else null);
            },
            .progress => {
                if (offset + 1 + 4 + 4 > data.len) return error.InvalidPacket;
                const progress_type = @as(ProgressType, @enumFromInt(data[offset]));
                offset += 1;

                const progress_value = std.mem.readInt(u32, data[offset .. offset + 4][0..4], .big);
                offset += 4;

                const progress_total = std.mem.readInt(u32, data[offset .. offset + 4][0..4], .big);
                offset += 4;

                const progress_message = try readString(data, &offset, allocator);

                return ResponsePacket.initProgress(timestamp, progress_type, progress_value, if (progress_total > 0) progress_total else null, if (progress_message.len > 0) progress_message else null);
            },
            .status => {
                const status_data = try readString(data, &offset, allocator);
                return ResponsePacket.initStatus(timestamp, status_data);
            },
            .data => {
                const data_type = try readString(data, &offset, allocator);
                const data_payload = try readBytes(data, &offset, allocator);
                return ResponsePacket.initData(timestamp, data_type, data_payload);
            },
            .log => {
                if (offset >= data.len) return error.InvalidPacket;
                const log_level = data[offset];
                offset += 1;

                const log_message = try readString(data, &offset, allocator);
                return ResponsePacket.initLog(timestamp, log_level, log_message);
            },
        }
    }

    /// Get human-readable error message for error code
    pub fn getErrorMessage(code: ErrorCode) []const u8 {
        return switch (code) {
            .unknown_error => "Unknown error occurred",
            .invalid_request => "Invalid request format",
            .authentication_failed => "Authentication failed",
            .permission_denied => "Permission denied",
            .resource_not_found => "Resource not found",
            .resource_already_exists => "Resource already exists",

            .server_overloaded => "Server is overloaded",
            .database_error => "Database error occurred",
            .network_error => "Network error occurred",
            .storage_error => "Storage error occurred",
            .timeout => "Operation timed out",

            .job_not_found => "Job not found",
            .job_already_running => "Job is already running",
            .job_failed_to_start => "Job failed to start",
            .job_execution_failed => "Job execution failed",
            .job_cancelled => "Job was cancelled",

            .out_of_memory => "Server out of memory",
            .disk_full => "Server disk full",
            .invalid_configuration => "Invalid server configuration",
            .service_unavailable => "Service temporarily unavailable",
        };
    }

    /// Get log level name
    pub fn getLogLevelName(level: u8) []const u8 {
        return switch (level) {
            0 => "DEBUG",
            1 => "INFO",
            2 => "WARN",
            3 => "ERROR",
            else => "UNKNOWN",
        };
    }

    /// Free all allocated memory in the packet
    pub fn deinit(self: ResponsePacket, allocator: std.mem.Allocator) void {
        if (self.success_message) |msg| allocator.free(msg);
        if (self.error_message) |msg| allocator.free(msg);
        if (self.error_details) |details| allocator.free(details);
        if (self.progress_message) |msg| allocator.free(msg);
        if (self.status_data) |data| allocator.free(data);
        if (self.data_type) |dtype| allocator.free(dtype);
        if (self.data_payload) |payload| allocator.free(payload);
        if (self.log_message) |msg| allocator.free(msg);
    }
};

/// Helper function to write string with length prefix
fn writeString(buffer: *std.ArrayList(u8), allocator: std.mem.Allocator, str: []const u8) !void {
    try writeUvarint(buffer, allocator, @as(u64, str.len));
    try buffer.appendSlice(allocator, str);
}

/// Helper function to write bytes with length prefix
fn writeBytes(buffer: *std.ArrayList(u8), allocator: std.mem.Allocator, bytes: []const u8) !void {
    try writeUvarint(buffer, allocator, @as(u64, bytes.len));
    try buffer.appendSlice(allocator, bytes);
}

/// Helper function to read string with length prefix
fn readString(data: []const u8, offset: *usize, allocator: std.mem.Allocator) ![]const u8 {
    const len64 = try readUvarint(data, offset);
    if (len64 > @as(u64, std.math.maxInt(usize))) return error.InvalidPacket;
    const len: usize = @intCast(len64);

    if (offset.* + len > data.len) return error.InvalidPacket;

    const str = try allocator.alloc(u8, len);
    @memcpy(str, data[offset.* .. offset.* + len]);
    offset.* += len;

    return str;
}

/// Helper function to read bytes with length prefix
fn readBytes(data: []const u8, offset: *usize, allocator: std.mem.Allocator) ![]const u8 {
    const len64 = try readUvarint(data, offset);
    if (len64 > @as(u64, std.math.maxInt(usize))) return error.InvalidPacket;
    const len: usize = @intCast(len64);

    if (offset.* + len > data.len) return error.InvalidPacket;

    const bytes = try allocator.alloc(u8, len);
    @memcpy(bytes, data[offset.* .. offset.* + len]);
    offset.* += len;

    return bytes;
}

fn writeUvarint(buffer: *std.ArrayList(u8), allocator: std.mem.Allocator, value: u64) !void {
    var x = value;
    while (x >= 0x80) {
        const b: u8 = @intCast((x & 0x7f) | 0x80);
        try buffer.append(allocator, b);
        x >>= 7;
    }
    try buffer.append(allocator, @intCast(x));
}

fn readUvarint(data: []const u8, offset: *usize) !u64 {
    var x: u64 = 0;
    var s: u6 = 0;
    var i: usize = 0;
    while (i < 10) : (i += 1) {
        if (offset.* >= data.len) return error.InvalidPacket;
        const b = data[offset.*];
        offset.* += 1;

        if (b < 0x80) {
            if (i == 9 and b > 1) return error.InvalidPacket;
            return x | (@as(u64, b) << s);
        }

        x |= (@as(u64, b & 0x7f) << s);
        s += 7;
    }
    return error.InvalidPacket;
}

test "deserialize data packet (varint lengths)" {
    const allocator = std.testing.allocator;

    // PacketTypeData (0x04), timestamp=1 (big-endian)
    var buf = std.ArrayList(u8).initCapacity(allocator, 64) catch unreachable;
    defer buf.deinit(allocator);

    try buf.append(allocator, 0x04);
    var ts: [8]u8 = undefined;
    std.mem.writeInt(u64, ts[0..8], 1, .big);
    try buf.appendSlice(allocator, &ts);

    // data_type="experiment" (len=10 -> 0x0A)
    try buf.append(allocator, 10);
    try buf.appendSlice(allocator, "experiment");

    // payload="{}" (len=2 -> 0x02)
    try buf.append(allocator, 2);
    try buf.appendSlice(allocator, "{}");

    const packet = try ResponsePacket.deserialize(buf.items, allocator);
    defer packet.deinit(allocator);

    try std.testing.expectEqual(PacketType.data, packet.packet_type);
    try std.testing.expectEqual(@as(u64, 1), packet.timestamp);
    try std.testing.expectEqualStrings("experiment", packet.data_type.?);
    try std.testing.expectEqualStrings("{}", packet.data_payload.?);
}
