const asynk = @import("async");
const std = @import("std");
const zml = @import("../../zml.zig");

const assert = std.debug.assert;
const log = std.log.scoped(.zml_io);

pub const GgufErrors = error{
    ValueTypeMismatch,
    InvalidGguf,
    UnsupportedGgufType,
    EndOfMetadata,
    OutOfMemory,
};

// Enums and structures
pub const TensorType = enum(u32) {
    f32 = 0,
    f16 = 1,
    q4_0 = 2,
    q4_1 = 3,
    deprecated_q4_2 = 4,
    deprecated_q4_3 = 5,
    q5_0 = 6,
    q5_1 = 7,
    q8_0 = 8,
    q8_1 = 9,
    // k-quantizations
    q2_k = 10,
    q3_k = 11,
    q4_k = 12,
    q5_k = 13,
    q6_k = 14,
    q8_k = 15,
    i8 = 16,
    i16 = 17,
    i32 = 18,

    const MAX_KNOWN_ENUM = 18;

    pub fn canConvertQuant(self: TensorType) bool {
        return switch (self) {
            .q8_0, .q4_k, .q6_k, .q2_k, .q4_0, .q4_1 => true,
            else => false,
        };
    }

    pub fn toDtype(self: TensorType) ?zml.DataType {
        return switch (self) {
            .f32 => .f32,
            .f16 => .f16,
            .i8 => .i8,
            .i16 => .i16,
            .i32 => .i32,
            else => null,
        };
    }

    pub fn sizeOf(self: TensorType) usize {
        return self.toDtype().?.sizeOf();
    }

    /// Return the tensor type features
    pub fn getFeatures(t: TensorType) TensorTypeFeatures {
        return switch (t) {
            inline else => |val| @field(TENSOR_TYPE_FEATURES, @tagName(val)),
        };
    }
};

/// GGUF tensor type to features lookup table.
pub const TensorTypeFeatures = struct {
    items_per_block: u29,
    bytes_per_block: u29,

    pub fn alignment(features: TensorTypeFeatures) u8 {
        return std.math.log2_int(u29, features.bytes_per_block);
    }
};

pub const TENSOR_TYPE_FEATURES: std.enums.EnumFieldStruct(TensorType, TensorTypeFeatures, null) = .{
    .f32 = .{ .items_per_block = 1, .bytes_per_block = @sizeOf(f32) },
    .f16 = .{ .items_per_block = 1, .bytes_per_block = @sizeOf(f16) },
    .q4_0 = .{ .items_per_block = 32, .bytes_per_block = 18 },
    .q4_1 = .{ .items_per_block = 32, .bytes_per_block = 20 },
    .deprecated_q4_2 = .{ .items_per_block = 0, .bytes_per_block = 0 },
    .deprecated_q4_3 = .{ .items_per_block = 0, .bytes_per_block = 0 },
    .q5_0 = .{ .items_per_block = 32, .bytes_per_block = 22 },
    .q5_1 = .{ .items_per_block = 32, .bytes_per_block = 24 },
    .q8_0 = .{ .items_per_block = 32, .bytes_per_block = 34 },
    .q8_1 = .{ .items_per_block = 32, .bytes_per_block = 40 },
    .q2_k = .{ .items_per_block = 256, .bytes_per_block = 82 },
    .q3_k = .{ .items_per_block = 256, .bytes_per_block = 110 },
    .q4_k = .{ .items_per_block = 256, .bytes_per_block = 144 },
    .q5_k = .{ .items_per_block = 256, .bytes_per_block = 176 },
    .q6_k = .{ .items_per_block = 256, .bytes_per_block = 210 },
    .q8_k = .{ .items_per_block = 256, .bytes_per_block = 292 },
    .i8 = .{ .items_per_block = 1, .bytes_per_block = @sizeOf(i8) },
    .i16 = .{ .items_per_block = 1, .bytes_per_block = @sizeOf(i16) },
    .i32 = .{ .items_per_block = 1, .bytes_per_block = @sizeOf(i32) },
};

pub const GgufValueType = enum(u32) {
    // The value is a 8-bit unsigned integer.
    uint8 = 0,
    // The value is a 8-bit signed integer.
    int8 = 1,
    // The value is a 16-bit unsigned little-endian integer.
    uint16 = 2,
    // The value is a 16-bit signed little-endian integer.
    int16 = 3,
    // The value is a 32-bit unsigned little-endian integer.
    uint32 = 4,
    // The value is a 32-bit signed little-endian integer.
    int32 = 5,
    // The value is a 32-bit IEEE754 floating point number.
    float32 = 6,
    // The value is a boolean.
    // 1-byte value where 0 is false and 1 is true.
    // Anything else is invalid, and should be treated as either the model
    // being invalid or the reader being buggy.
    bool = 7,
    // The value is a UTF-8 non-null-terminated string, with length prepended.
    string = 8,
    // The value is an array of other values, with the length and type
    // prepended. Arrays can be nested, and the length of the array is the
    // number of elements in the array, not the number of bytes.
    array = 9,
    // The value is a 64-bit unsigned little-endian integer.
    uint64 = 10,
    // The value is a 64-bit signed little-endian integer.
    int64 = 11,
    // The value is a 64-bit IEEE754 floating point number.
    float64 = 12,
    // Special values used by the callbacks of gguf_do_with_value().
    array_start = 100,
    array_end = 101,

    // Allow other values in case GGUF add more types without us noticing
    _,

    pub fn sizeOf(self: GgufValueType) usize {
        return switch (self) {
            .uint8 => @sizeOf(u8),
            .int8 => @sizeOf(i8),
            .uint16 => @sizeOf(u16),
            .int16 => @sizeOf(i16),
            .uint32 => @sizeOf(u32),
            .int32 => @sizeOf(i32),
            .float32 => @sizeOf(f32),
            .bool => @sizeOf(bool),
            .uint64 => @sizeOf(u64),
            .int64 => @sizeOf(i64),
            .float64 => @sizeOf(f64),
            .string => @sizeOf([]u8),
            else => unreachable,
        };
    }

    pub fn arrayTypeCheck(self: GgufValueType, comptime T: type) !void {
        switch (self) {
            .string => if (T != []u8 and T != []const u8) return error.ValueTypeMismatch,
            .uint8 => if (T != u8) return error.ValueTypeMismatch,
            .int8 => if (T != i8) return error.ValueTypeMismatch,
            .uint16 => if (T != u16) return error.ValueTypeMismatch,
            .int16 => if (T != i16) return error.ValueTypeMismatch,
            .uint32 => if (T != u32) return error.ValueTypeMismatch,
            .int32 => if (T != i32) return error.ValueTypeMismatch,
            .float32 => if (T != f32) return error.ValueTypeMismatch,
            .bool => if (T != bool) return error.ValueTypeMismatch,
            .uint64 => if (T != u64) return error.ValueTypeMismatch,
            .int64 => if (T != i64) return error.ValueTypeMismatch,
            .float64 => if (T != f64) return error.ValueTypeMismatch,
            else => {},
        }
    }
};

pub const ValueType = enum {
    uint8,
    int8,
    uint16,
    int16,
    uint32,
    int32,
    float32,
    uint64,
    int64,
    float64,
    boolval,
    string,
    array,
};

// Union of possible values.
pub const GgufValue = union(ValueType) {
    uint8: u8,
    int8: i8,
    uint16: u16,
    int16: i16,
    uint32: u32,
    int32: i32,
    float32: f32,
    uint64: u64,
    int64: i64,
    float64: f64,
    boolval: bool,
    string: []const u8,
    array: Array,

    pub const Array = struct {
        // Any value type is valid, including arrays.
        child: GgufValueType,
        // Number of elements, not bytes
        len: usize,
        data: []u8,
    };

    pub fn asLoaderValue(self: GgufValue) zml.aio.Value {
        return switch (self) {
            .array => |v| .{
                .array = .{
                    .item_type = switch (v.child) {
                        .bool => .boolval,
                        .uint8 => .uint8,
                        .int8 => .int8,
                        .uint16 => .uint16,
                        .int16 => .int16,
                        .uint32 => .uint32,
                        .int32 => .int32,
                        .float32 => .float32,
                        .uint64 => .uint64,
                        .int64 => .int64,
                        .float64 => .float64,
                        .string => .string,
                        // TODO: .array => .array,
                        else => unreachable,
                    },
                    .data = v.data,
                },
            },
            inline else => |v, tag| @unionInit(zml.aio.Value, @tagName(tag), v),
        };
    }
};

// Header
const GgufHeader = extern struct {
    // Magic number to announce that this is a GGUF file. Must be `GUFF`.
    magic: [4]u8,
    // The version of the format implemented.
    // Must be `3` for version described in this spec.
    version: u32,
    // The number of tensors in the file.
    // This is explicit, instead of being included in the metadata, to ensure
    // it is always present for loading the tensors.
    tensor_count: usize,
    // The number of metadata key-value pairs.
    metadata_kv_count: usize,

    pub fn validate(self: GgufHeader) !void {
        if (!std.mem.eql(u8, &self.magic, "GGUF")) {
            log.err("Invalid GGUF file: wrong header {s}", .{self.magic});
            return error.InvalidGguf;
        }
    }
};

// Key representation in this library API.
pub const GgufMetadataKv = struct {
    name: []const u8,
    type_: GgufValueType,
    val: GgufValue,
};

// Tensor representation in this library API.
const GGUF_TENSOR_MAX_DIM: usize = 8; // Future-proof: actual limit is 4.
pub const GgufTensorInfo = struct {
    name: []const u8,
    t: TensorType, // Tensor type (enum TensorType).
    rank: usize, // Number of dimensions of the tensor.
    dims: [GGUF_TENSOR_MAX_DIM]i64, // Dimensions (Eg. [512, 1024, 1, 1]).
    start: usize, // Offset from start of data section.
    byte_len: usize, // Total size in bytes.
    num_weights: usize, // Total number of parameters.

    pub inline fn shape(info: GgufTensorInfo) []const i64 {
        return info.dims[0..info.rank];
    }
};

// Return the value type name given the type ID.
fn getValueTypeName(t: u32) []const u8 {
    if (@as(usize, @intCast(t)) >= GGUF_VALUE_NAME.len) return "unknown";
    return GGUF_VALUE_NAME[@intCast(t)];
}

const GGUF_VALUE_NAME = [_][]const u8{
    "uint8",   "int8", "uint16", "int16", "uint32", "int32",
    "float32", "bool", "string", "array", "uint64", "int64",
    "float64",
};

/// GGUF file API
/// A memory-mapped view of a .gguf file.
/// Format used by GGML models: https://github.com/ggerganov/ggml/
pub const GgufFile = struct {
    header: GgufHeader, // GUFF file header info.
    size: usize, // Total file size.
    file: zml.aio.MemoryMappedFile,
    left_kv: usize, // Number of key-value pairs yet to read.
    left_tensors: usize, // Number of tensors yet to read.
    off: usize, // Offset of the next item to parse.
    alignment: usize = 32, // File data alignment. Default: 32 bytes.

    /// Open and memmap the given file.
    pub fn open(path: []const u8) !GgufFile {
        const file = try asynk.File.open(path, .{});
        const header = try file.reader().readStruct(GgufHeader);
        try header.validate();
        return .{
            .header = header,
            .size = (try file.stat()).size,
            .file = try zml.aio.MemoryMappedFile.init(file),
            .off = @sizeOf(GgufHeader),
            .left_kv = header.metadata_kv_count,
            .left_tensors = header.tensor_count,
        };
    }

    ///  Unmap the file memory and close the file handle.
    pub fn close(self: *GgufFile) void {
        self.file.deinit();
    }

    /// Set the context to read the first key-value entry in the GGUF
    /// file and then all the rest. Is used when creating a new context
    /// and also when you want to restart scanning the key-value
    /// items in the file.
    fn rewind(ctx: *GgufFile) void {
        ctx.off = @sizeOf(GgufHeader);
        ctx.left_kv = ctx.header.metadata_kv_count;
        ctx.left_tensors = ctx.header.tensor_count;
    }

    pub fn seek(self: *GgufFile, pos: usize) void {
        assert(pos < self.size);
        self.off = pos;
    }

    fn readInt(self: *GgufFile, comptime T: type) !T {
        if (self.off + @sizeOf(T) >= self.size) return error.InvalidGguf;
        const res = self.file.file.reader().readInt(T, .little);
        self.off += @sizeOf(T);
        return res;
    }

    fn readTensorType(self: *GgufFile) !TensorType {
        const raw = try self.readInt(u32);
        if (raw > TensorType.MAX_KNOWN_ENUM) {
            log.err("Unsupported GGUF tensor type: {d}", .{raw});
            return error.UnsupportedGgufType;
        }
        return @enumFromInt(raw);
    }

    fn readValueType(self: *GgufFile) !GgufValueType {
        const raw = try self.readInt(u32);
        const t: GgufValueType = @enumFromInt(raw);
        switch (t) {
            .uint8, .int8, .uint16, .int16, .uint32, .int32, .float32, .bool, .string, .array, .uint64, .int64, .float64, .array_start, .array_end => {},
            else => {
                log.err("Unsupported GGUF value type: {s}", .{@tagName(t)});
                return error.UnsupportedGgufType;
            },
        }
        return t;
    }

    pub fn readAlloc(self: *GgufFile, allocator: std.mem.Allocator, len: usize) ![]u8 {
        const data = try allocator.alloc(u8, len);
        const read = try self.file.file.reader().readAll(data);
        if (read != data.len) return error.InvalidGguf;
        self.off += len;
        return data;
    }

    pub fn skipBytes(self: *GgufFile, len: usize) !void {
        try self.file.file.seekBy(@intCast(len));
        self.off += len;
    }

    /// Read the len then the actual bytes.
    pub fn readString(self: *GgufFile, allocator: std.mem.Allocator) ![]u8 {
        const len: usize = try self.readInt(u64);
        return self.readAlloc(allocator, len);
    }

    pub fn skipString(self: *GgufFile) !void {
        const len: usize = try self.readInt(u64);
        return self.skipBytes(len);
    }

    fn readArrayHeader(self: *GgufFile, allocator: std.mem.Allocator) !GgufValue.Array {
        const child = try self.readValueType();
        const len: usize = try self.readInt(u64);
        const data = switch (child) {
            // Since strings have variable lenghts, we need to read them one by one
            .string => str: {
                var data = try allocator.alloc([]u8, len);
                for (0..len) |i| data[i] = try self.readString(allocator);
                break :str std.mem.sliceAsBytes(data);
            },
            else => try self.readAlloc(allocator, len * child.sizeOf()),
        };
        return .{
            .child = child,
            .len = len,
            .data = data,
        };
    }

    fn readTypedValue(self: *GgufFile, allocator: std.mem.Allocator, t: GgufValueType) !GgufValue {
        return switch (t) {
            .uint8 => .{ .uint8 = try self.readInt(u8) },
            .int8 => .{ .int8 = try self.readInt(i8) },
            .uint16 => .{ .uint16 = try self.readInt(u16) },
            .int16 => .{ .int16 = try self.readInt(i16) },
            .uint32 => .{ .uint32 = try self.readInt(u32) },
            .int32 => .{ .int32 = try self.readInt(i32) },
            .float32 => .{ .float32 = @bitCast(try self.readInt(u32)) },
            .bool => .{ .boolval = try self.readInt(u8) != 0 },
            .string => .{ .string = try self.readString(allocator) },
            .array => .{ .array = try self.readArrayHeader(allocator) },
            .uint64 => .{ .uint64 = try self.readInt(u64) },
            .int64 => .{ .int64 = try self.readInt(i64) },
            .float64 => .{ .float64 = @bitCast(try self.readInt(u64)) },
            else => error.UnsupportedGgufType,
        };
    }

    /// Parses the next metadata entry.
    /// Returns error.EndOfMetadata if there are no longer metadata to process in this GGUF file.
    pub fn readMetadata(self: *GgufFile, allocator: std.mem.Allocator) !GgufMetadataKv {
        if (self.left_kv == 0) return error.EndOfMetadata;
        self.left_kv -= 1;
        const name = try self.readString(allocator);
        const type_ = try self.readValueType();
        const val: GgufValue = try self.readTypedValue(allocator, type_);
        return .{ .name = name, .type_ = type_, .val = val };
    }

    // Set the data section offset. This function must be called exactly when
    // all the key-values are consumed, in the context of the first call of
    // ctx.getTensor(): this way we will be able to return tensor offsets
    // as absolute positions and pointers to the mmapped file.
    fn setDataOffset(self: *GgufFile) !void {
        const base_off = self.off;

        assert(self.left_kv == 0 and self.left_tensors == self.header.tensor_count);

        for (0..self.left_tensors) |_| try self.skipTensor();
        const padding: usize = getAlignmentPadding(self.alignment, self.off);
        self.file.data_offset = self.off + padding;

        try self.file.file.seekTo(base_off);
        self.off = base_off;
    }

    pub fn skipTensor(self: *GgufFile) !void {
        try self.skipString(); // Skip name
        const num_dim: u32 = try self.readInt(u32);
        // dimensions, type, and offset.
        try self.skipBytes(8 * num_dim + 4 + 8);
    }

    /// Parses the next tensor entry.
    /// Returns error.EndOfMetadata if there are no longer tensor metadata to process in this GGUF file.
    pub fn readTensorInfo(self: *GgufFile, allocator: std.mem.Allocator) !GgufTensorInfo {
        if (self.left_tensors == 0 or self.left_kv != 0) {
            return error.EndOfMetadata;
        }

        // We want to return tensor data with offsets relative to the start
        // of the file, so that the user of the API is able to access tensors
        // as it iterates over them. To do so, we need to perform a full
        // scan if this is the first tensor info we are reading.
        // TODO: explicitly set the data offset in
        if (self.file.data_offset == 0) try self.setDataOffset();
        self.left_tensors -= 1;
        const name = try self.readString(allocator);
        const num_dim = try self.readInt(u32);
        assert(@as(usize, @intCast(num_dim)) <= GGUF_TENSOR_MAX_DIM);
        // Read the dimentions; unused dimensions are left `undefined`.
        // Note: we reverse the order of the dimensions to match zml convention.
        var dims: [GGUF_TENSOR_MAX_DIM]i64 = undefined;
        var num_weights: usize = 1;
        for (0..num_dim) |j| {
            const d = try self.readInt(u64);
            dims[num_dim - 1 - j] = @intCast(d);
            num_weights *= d;
        }
        const t: TensorType = try self.readTensorType();
        const start = try self.readInt(u64);
        // To accurately calculate the bytes used by this tensor on the GGUF
        // file, we need to take into account that quantization methods store
        // tensors as block of N weights. So first of all we need to understand
        // the number of padding weights (since the last block may have just
        // fewer weights stored inside, but still requires to be stored to its full
        // length). Then we can do the math to see how many blocks we need, and
        // multiply by the block size to obtain the final total size.
        const tf = t.getFeatures();
        const byte_len: usize = (std.math.divCeil(usize, num_weights, tf.items_per_block) catch unreachable) * tf.bytes_per_block;
        return .{
            .name = name,
            .t = t,
            .rank = num_dim,
            .dims = dims,
            .start = start,
            .byte_len = byte_len,
            .num_weights = num_weights,
        };
    }
};

/// Given an offset or a length, returns the padding needed to align it to alignment.
fn getAlignmentPadding(alignment: usize, offset: usize) usize {
    return @rem((alignment - @rem(offset, alignment)), alignment);
}