const std = @import("std");
const value = @import("value.zig");

/// Parser for GGUF model files (spec versions 2 and 3).
///
/// Layout:
///   magic "GGUF", version u32, tensor_count u64, metadata_kv_count u64,
///   metadata key/value pairs, tensor infos, then the data section starting
///   at `general.alignment` (default 32) alignment.
///
/// Metadata values are walked (not interpreted) except `general.alignment`.
/// Tensor dims arrive in ggml NE order (dim 0 = innermost/contiguous) and are
/// REVERSED into our row-major outermost-first shape convention.
///
/// v1 ggml_type allowlist: F32=0, I32=26, I64=27, F64=28 (verified against
/// llama.cpp ggml.h). Quantized/f16/bf16 tensors are rejected with a clear
/// error rather than misread.
///
/// Hard caps against hostile files: <= 1M tensors, <= 100k metadata pairs,
/// names/strings <= 1 MiB, rank <= 8, all arithmetic overflow-checked.
pub const ParseError = error{
    BadMagic,
    UnsupportedVersion,
    BadHeader,
    Truncated,
    UnsupportedDtype,
    TooLarge,
    OutOfMemory,
};

const max_tensors: u64 = 1_000_000;
const max_kvs: u64 = 100_000;
const max_string_len: u64 = 1024 * 1024;
const max_rank: u32 = 8;
const default_alignment: u32 = 32;

/// One numeric metadata kv captured for higher layers (e.g. LLM configs).
/// Integer-valued GGUF types only (UINT32/INT32/UINT64); `val` is the
/// zero/sign-extended value.
pub const KvNum = struct {
    name: []u8,
    val: u64,
};

/// Full parse output: tensor views plus captured numeric kvs.
pub const ParseResult = struct {
    views: []value.TensorView,
    nums: []KvNum,
};

/// ggml_type codes we can decode (llama.cpp ggml.h).
fn dtypeFromGgml(t: u32) ?value.DType {
    return switch (t) {
        0 => .f32,
        26 => .i32,
        27 => .i64,
        28 => .f64,
        else => null,
    };
}

/// Bounds-checked little-endian reader over a byte slice.
const Cur = struct {
    b: []const u8,
    i: usize = 0,

    fn take(self: *Cur, n: usize) ParseError![]const u8 {
        if (self.i + n > self.b.len or self.i + n < self.i) return error.Truncated;
        const s = self.b[self.i .. self.i + n];
        self.i += n;
        return s;
    }

    fn u32le(self: *Cur) ParseError!u32 {
        const s = try self.take(4);
        return std.mem.readInt(u32, s[0..4], .little);
    }

    fn u64le(self: *Cur) ParseError!u64 {
        const s = try self.take(8);
        return std.mem.readInt(u64, s[0..8], .little);
    }

    /// Reads a u64-length-prefixed byte string (not duped).
    fn rawStr(self: *Cur) ParseError![]const u8 {
        const len = try self.u64le();
        if (len > max_string_len) return error.TooLarge;
        return self.take(@intCast(len));
    }
};

/// Skips one metadata value of type `t`. Variable-length containers are
/// walked so nested sizes stay bounds-checked by the cursor.
fn skipValue(cur: *Cur, t: u32) ParseError!void {
    switch (t) {
        0, 1, 7 => _ = try cur.take(1),
        2, 3 => _ = try cur.take(2),
        4, 5, 6 => _ = try cur.take(4),
        10, 11, 12 => _ = try cur.take(8),
        8 => {
            _ = try cur.rawStr();
        },
        9 => {
            const elem_t = try cur.u32le();
            if (elem_t > 12) return error.BadHeader;
            const count = try cur.u64le();
            var k: u64 = 0;
            while (k < count) : (k += 1) try skipValue(cur, elem_t);
        },
        else => return error.BadHeader,
    }
}

fn alignUp(x: u64, a: u64) ?u64 {
    if (a == 0 or (a & (a - 1)) != 0) return null; // must be a power of two
    const y = std.math.add(u64, x, a - 1) catch return null;
    return y & ~(a - 1);
}

/// Parses `file` (complete file contents) into tensor views. On success the
/// caller owns the returned slice and every `name`/`shape` inside it.
pub fn parse(gpa: std.mem.Allocator, file: []const u8) ParseError![]value.TensorView {
    const r = try parseWithKv(gpa, file);
    return r.views;
}

/// Parses `file` into tensor views AND captures integer-valued metadata kvs
/// (UINT32/INT32/UINT64). Caller owns everything returned.
pub fn parseWithKv(gpa: std.mem.Allocator, file: []const u8) ParseError!ParseResult {
    var cur = Cur{ .b = file };

    const magic = cur.take(4) catch return error.BadMagic;
    if (!std.mem.eql(u8, magic, "GGUF")) return error.BadMagic;
    const version = try cur.u32le();
    if (version != 2 and version != 3) return error.UnsupportedVersion;

    const tensor_count = try cur.u64le();
    const kv_count = try cur.u64le();
    if (tensor_count > max_tensors or kv_count > max_kvs) return error.TooLarge;

    var alignment: u64 = default_alignment;
    var nums: std.ArrayList(KvNum) = .empty;
    errdefer {
        for (nums.items) |n| gpa.free(n.name);
        nums.deinit(gpa);
    }
    var kv_i: u64 = 0;
    while (kv_i < kv_count) : (kv_i += 1) {
        const key = try cur.rawStr();
        const vt = try cur.u32le();
        if (vt > 12) return error.BadHeader;
        var captured: ?u64 = null;
        switch (vt) {
            4 => captured = try cur.u32le(), // UINT32
            // INT32 sign-extended into i64 storage.
            5 => captured = @bitCast(@as(i64, @as(i32, @bitCast(try cur.u32le())))),
            10 => captured = try cur.u64le(), // UINT64
            else => try skipValue(&cur, vt),
        }
        if (std.mem.eql(u8, key, "general.alignment") and vt == 4) {
            alignment = @intCast(captured.?);
            if (alignment == 0 or (alignment & (alignment - 1)) != 0 or alignment > 1024 * 1024) {
                return error.BadHeader; // spec requires a sane power of two
            }
            continue;
        }
        if (captured) |v| {
            const name = gpa.dupe(u8, key) catch return error.OutOfMemory;
            errdefer gpa.free(name);
            try nums.append(gpa, .{ .name = name, .val = v });
        }
    }

    const Pending = struct {
        name: []u8,
        dtype: value.DType,
        shape: []u64,
        rel_off: u64,
        len: u64,
    };
    var pending: std.ArrayList(Pending) = .empty;
    defer {
        for (pending.items) |p| {
            gpa.free(p.name);
            gpa.free(p.shape);
        }
        pending.deinit(gpa);
    }

    var t_i: u64 = 0;
    while (t_i < tensor_count) : (t_i += 1) {
        const name_raw = try cur.rawStr();

        const n_dims = try cur.u32le();
        if (n_dims == 0 or n_dims > max_rank) return error.BadHeader;

        const shape = gpa.alloc(u64, n_dims) catch return error.OutOfMemory;
        errdefer gpa.free(shape);
        var d: u32 = 0;
        while (d < n_dims) : (d += 1) {
            // ggml NE order: dim 0 is innermost -> reverse into our layout.
            shape[n_dims - 1 - d] = try cur.u64le();
        }

        var numel: u64 = 1;
        for (shape) |dim| {
            numel = std.math.mul(u64, numel, dim) catch return error.TooLarge;
        }

        const gt = try cur.u32le();
        const dtype = dtypeFromGgml(gt) orelse return error.UnsupportedDtype;
        const len = std.math.mul(u64, numel, dtype.size()) catch return error.TooLarge;

        const rel_off = try cur.u64le();

        const name = gpa.dupe(u8, name_raw) catch return error.OutOfMemory;
        errdefer gpa.free(name);

        try pending.append(gpa, .{
            .name = name,
            .dtype = dtype,
            .shape = shape,
            .rel_off = rel_off,
            .len = len,
        });
    }

    // Data section begins at the next alignment boundary after the header.
    const data_start = alignUp(@intCast(cur.i), alignment) orelse return error.TooLarge;

    var list: std.ArrayList(value.TensorView) = .empty;
    errdefer {
        for (list.items) |v| {
            gpa.free(v.name);
            gpa.free(v.shape);
        }
        list.deinit(gpa);
    }
    for (pending.items) |p| {
        const abs = std.math.add(u64, data_start, p.rel_off) catch return error.BadHeader;
        const end = std.math.add(u64, abs, p.len) catch return error.BadHeader;
        if (end > file.len) return error.Truncated;
        try list.append(gpa, .{
            .name = p.name,
            .dtype = p.dtype,
            .shape = p.shape,
            .off = abs,
            .len = p.len,
        });
    }
    // Names/shapes moved into `list`; disarm the pending cleanup.
    pending.clearRetainingCapacity();

    return .{
        .views = try list.toOwnedSlice(gpa),
        .nums = try nums.toOwnedSlice(gpa),
    };
}

// ---------------------------------------------------------------------------
// Tests: fixtures are built in-memory; no filesystem access needed.

fn wU32(list: *std.ArrayList(u8), gpa: std.mem.Allocator, n: u32) !void {
    var b: [4]u8 = undefined;
    std.mem.writeInt(u32, &b, n, .little);
    try list.appendSlice(gpa, &b);
}

fn wU64(list: *std.ArrayList(u8), gpa: std.mem.Allocator, n: u64) !void {
    var b: [8]u8 = undefined;
    std.mem.writeInt(u64, &b, n, .little);
    try list.appendSlice(gpa, &b);
}

fn wStr(list: *std.ArrayList(u8), gpa: std.mem.Allocator, s: []const u8) !void {
    try wU64(list, gpa, s.len);
    try list.appendSlice(gpa, s);
}

test "parse minimal gguf v3 fixture" {
    const gpa = std.testing.allocator;

    var f: std.ArrayList(u8) = .empty;
    defer f.deinit(gpa);

    try f.appendSlice(gpa, "GGUF");
    try wU32(&f, gpa, 3);
    try wU64(&f, gpa, 1); // tensor_count
    try wU64(&f, gpa, 1); // kv_count

    // kv: general.alignment (u32) = 32
    try wStr(&f, gpa, "general.alignment");
    try wU32(&f, gpa, 4);
    try wU32(&f, gpa, 32);

    // tensor info: "x", 2 dims in NE order [3, 2] -> shape [2, 3], F32, off 0
    try wStr(&f, gpa, "x");
    try wU32(&f, gpa, 2);
    try wU64(&f, gpa, 3);
    try wU64(&f, gpa, 2);
    try wU32(&f, gpa, 0); // GGML_TYPE_F32
    try wU64(&f, gpa, 0); // relative offset

    // Pad to the 32-byte alignment, then 6 f32 elements.
    const pad: usize = (32 - f.items.len % 32) % 32;
    try f.appendNTimes(gpa, 0, pad);
    const data_start: u64 = @intCast(f.items.len);
    const vals = [_]f32{ 1, 2, 3, 4, 5, 6 };
    for (vals) |v| try wU32(&f, gpa, @bitCast(v));

    const views = try parse(gpa, f.items);
    defer {
        for (views) |v| {
            gpa.free(v.name);
            gpa.free(v.shape);
        }
        gpa.free(views);
    }

    try std.testing.expectEqual(@as(usize, 1), views.len);
    try std.testing.expectEqualStrings("x", views[0].name);
    try std.testing.expectEqual(value.DType.f32, views[0].dtype);
    try std.testing.expectEqual(@as(u64, 2), views[0].shape[0]); // reversed!
    try std.testing.expectEqual(@as(u64, 3), views[0].shape[1]);
    try std.testing.expectEqual(data_start, views[0].off);
    try std.testing.expectEqual(@as(u64, 24), views[0].len);

    const raw = std.mem.readInt(u32, f.items[@intCast(views[0].off)..][0..4], .little);
    try std.testing.expectEqual(@as(f32, 1.0), @as(f32, @bitCast(raw)));
}

test "gguf v2 accepted" {
    const gpa = std.testing.allocator;
    var f: std.ArrayList(u8) = .empty;
    defer f.deinit(gpa);
    try f.appendSlice(gpa, "GGUF");
    try wU32(&f, gpa, 2);
    try wU64(&f, gpa, 0);
    try wU64(&f, gpa, 0);
    const views = try parse(gpa, f.items);
    defer gpa.free(views);
    try std.testing.expectEqual(@as(usize, 0), views.len);
}

test "gguf rejects bad magic" {
    const gpa = std.testing.allocator;
    try std.testing.expectError(error.BadMagic, parse(gpa, "NOTG"));
}

test "gguf rejects unsupported version" {
    const gpa = std.testing.allocator;
    var f: std.ArrayList(u8) = .empty;
    defer f.deinit(gpa);
    try f.appendSlice(gpa, "GGUF");
    try wU32(&f, gpa, 4);
    try wU64(&f, gpa, 0);
    try wU64(&f, gpa, 0);
    try std.testing.expectError(error.UnsupportedVersion, parse(gpa, f.items));
}

test "gguf rejects f16 tensors" {
    const gpa = std.testing.allocator;
    var f: std.ArrayList(u8) = .empty;
    defer f.deinit(gpa);
    try f.appendSlice(gpa, "GGUF");
    try wU32(&f, gpa, 3);
    try wU64(&f, gpa, 1);
    try wU64(&f, gpa, 0);
    try wStr(&f, gpa, "h");
    try wU32(&f, gpa, 1); // 1 dim
    try wU64(&f, gpa, 4);
    try wU32(&f, gpa, 1); // GGML_TYPE_F16 -> unsupported in v1
    try wU64(&f, gpa, 0);
    try std.testing.expectError(error.UnsupportedDtype, parse(gpa, f.items));
}