const std = @import("std");

const zml = @import("zml");
const attention = zml.attention.attention;

const common = @import("common.zig");
const lfm2 = @import("lfm2.zig");
const model = @import("lfm2/model.zig");

pub const std_options: std.Options = .{
    .log_level = .info,
};

const Args = struct {
    model: []const u8,
    activations: []const u8,
    backend: ?attention.Backend = null,

    pub const help =
        \\Use lfm2_tests --model=<path> --activations=<path> [options]
        \\
        \\ Validate the LFM2 implementation against activation fixtures.
        \\
        \\ Options:
        \\   --model=<path>            Path to the model repository
        \\   --activations=<path>      Path to activation safetensors
        \\   --backend=<text>          Attention backend to use
        \\
    ;
};

pub fn main(init: std.process.Init) !void {
    const allocator = init.gpa;
    const io = init.io;
    const args = zml.stdx.flags.parse(init.minimal.args, Args);

    const platform: *zml.Platform = try .auto(allocator, io, .{});
    defer platform.deinit(allocator, io);

    const repo = try zml.safetensors.resolveModelRepo(io, args.model);
    var registry: zml.safetensors.TensorRegistry = try .fromRepo(allocator, io, repo);
    defer registry.deinit();
    var store: zml.io.TensorStore = .fromRegistry(allocator, &registry);
    defer store.deinit();

    var repo_model = try lfm2.LoadedModel.init(allocator, io, repo, store.view(), .{});
    defer repo_model.deinit(allocator);

    var progress = std.Progress.start(io, .{ .root_name = args.model });
    const shardings: common.Shardings = try .init(platform);

    var model_buffers = try repo_model.loadBuffers(allocator, io, platform, &store, &progress, shardings);
    defer repo_model.unloadBuffers(&model_buffers, allocator);

    const backend = args.backend orelse attention.Backend.auto(platform);
    const params = lfm2.CompilationParameters.init(
        repo_model.inner,
        repo_model.parsed_config.value,
        repo_model.parsed_config.value.max_position_embeddings,
        backend,
        shardings,
    );
    progress.end();

    try run(allocator, io, platform, args.activations, repo_model.parsed_config.value, repo_model.inner, &model_buffers, params.attention_metadata, params.attention_parameters);
}

pub fn run(
    allocator: std.mem.Allocator,
    io: std.Io,
    platform: *zml.Platform,
    activations_path: []const u8,
    config: lfm2.Config,
    mdl: lfm2.Model,
    model_buffers: *lfm2.Buffers,
    attention_metadata: attention.Metadata,
    attention_parameters: attention.Parameters,
) !void {
    var registry: zml.safetensors.TensorRegistry = try .fromPath(allocator, io, activations_path);
    defer registry.deinit();

    var activation_store: zml.io.TensorStore = .fromRegistry(allocator, &registry);
    defer activation_store.deinit();

    var ctx = TestContext{
        .allocator = allocator,
        .io = io,
        .platform = platform,
        .activations_store = &activation_store,
        .attention_metadata = attention_metadata,
        .attention_parameters = attention_parameters,
        .sharding = platform.replicated_sharding,
    };

    try ctx.testLayer("embed_tokens", .{ .batch, .seq }, mdl.embed_tokens, model_buffers.embed_tokens, .{});

    var num_attn_layers: usize = 0;
    var num_conv_layers: usize = 0;
    for (0..config.num_hidden_layers) |i| {
        const kind = std.meta.stringToEnum(model.OperatorKind, config.layer_types[i]) orelse unreachable;
        const layer = mdl.layers[i];
        const layer_buffers = model_buffers.layers[i];
        switch (kind) {
            .conv => {
                try ctx.testLayerPrint("layers.{d}.conv.in_proj", .{i}, .{ .batch, .seq, .d }, layer.operator.conv.in_proj, layer_buffers.operator.conv.in_proj, .{});
                try ctx.testLayerPrint("layers.{d}.conv.out_proj", .{i}, .{ .batch, .seq, .d }, layer.operator.conv.out_proj, layer_buffers.operator.conv.out_proj, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward.w1", .{i}, .{ .batch, .seq, .d }, layer.feed_forward.w1, layer_buffers.feed_forward.w1, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward.w2", .{i}, .{ .batch, .seq, .d }, layer.feed_forward.w2, layer_buffers.feed_forward.w2, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward.w3", .{i}, .{ .batch, .seq, .d }, layer.feed_forward.w3, layer_buffers.feed_forward.w3, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward", .{i}, .{ .batch, .seq, .d }, layer.feed_forward, layer_buffers.feed_forward, .{});
                try ctx.testConvLayer(i, num_conv_layers, mdl.layers[i].operator.conv, model_buffers.layers[i].operator.conv, .{});
                num_conv_layers += 1;
            },
            .full_attention => {
                try ctx.testLayerPrint("layers.{d}.self_attn.k_proj", .{i}, .{ .batch, .seq, .d }, layer.operator.self_attn.k_proj, layer_buffers.operator.self_attn.k_proj, .{});
                try ctx.testLayerPrint("layers.{d}.self_attn.q_proj", .{i}, .{ .batch, .seq, .d }, layer.operator.self_attn.q_proj, layer_buffers.operator.self_attn.q_proj, .{});
                try ctx.testLayerPrint("layers.{d}.self_attn.out_proj", .{i}, .{ .batch, .seq, .d }, layer.operator.self_attn.out_proj, layer_buffers.operator.self_attn.out_proj, .{});
                try ctx.testLayerPrint("layers.{d}.self_attn.v_proj", .{i}, .{ .batch, .seq, .d }, layer.operator.self_attn.v_proj, layer_buffers.operator.self_attn.v_proj, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward.w1", .{i}, .{ .batch, .seq, .d }, layer.feed_forward.w1, layer_buffers.feed_forward.w1, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward.w2", .{i}, .{ .batch, .seq, .d }, layer.feed_forward.w2, layer_buffers.feed_forward.w2, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward.w3", .{i}, .{ .batch, .seq, .d }, layer.feed_forward.w3, layer_buffers.feed_forward.w3, .{});
                try ctx.testLayerPrint("layers.{d}.feed_forward", .{i}, .{ .batch, .seq, .d }, layer.feed_forward, layer_buffers.feed_forward, .{});
                try ctx.testAttnLayer(i, num_attn_layers, mdl.layers[i].operator.self_attn, model_buffers.layers[i].operator.self_attn, .{});
                num_attn_layers += 1;
            },
        }
    }
}

const TestContext = struct {
    allocator: std.mem.Allocator,
    io: std.Io,
    platform: *zml.Platform,
    activations_store: *zml.io.TensorStore,
    attention_metadata: attention.Metadata,
    attention_parameters: attention.Parameters,
    sharding: zml.Sharding,

    fn testLayerPrint(self: *TestContext, comptime name_fmt: []const u8, name_args: anytype, tagz: anytype, layer: anytype, layer_buffers: anytype, opts: zml.testing.CompareOpts) !void {
        const name = try std.fmt.allocPrint(self.allocator, name_fmt, name_args);
        defer self.allocator.free(name);
        try self.testLayer(name, tagz, layer, layer_buffers, opts);
    }

    fn testLayer(self: *TestContext, name: []const u8, tagz: anytype, layer: anytype, layer_buffers: anytype, opts: zml.testing.CompareOpts) !void {
        std.log.info("Testing layer: {s}", .{name});

        const in_key = try std.fmt.allocPrint(self.allocator, "{s}.in", .{name});
        defer self.allocator.free(in_key);
        var in_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, in_key, self.sharding);
        defer in_buffer.deinit();
        const in_tensor = zml.Tensor.fromShape(in_buffer.shape()).withTags(tagz);

        const out_key = try std.fmt.allocPrint(self.allocator, "{s}.out", .{name});
        defer self.allocator.free(out_key);
        var out_buffer_expected = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, out_key, self.sharding);
        defer out_buffer_expected.deinit();

        const exe = try self.platform.compileFn(self.allocator, self.io, @TypeOf(layer).forward, .{ layer, in_tensor }, .{ .shardings = &.{self.sharding} });
        defer exe.deinit();

        var args = try exe.args(self.allocator);
        defer args.deinit(self.allocator);
        args.set(.{ layer_buffers, in_buffer });

        var res = try exe.results(self.allocator);
        defer res.deinit(self.allocator);

        exe.call(args, &res);

        var out_result = res.get(zml.Buffer);
        defer out_result.deinit();

        try zml.testing.expectClose(self.io, out_result, out_buffer_expected, opts);
        std.log.info("Layer {s} passed!", .{name});
    }

    fn testConvLayer(self: *TestContext, ix: usize, cache_ix: usize, layer: model.ShortConv, layer_buffers: zml.Bufferized(model.ShortConv), opts: zml.testing.CompareOpts) !void {
        const name = try std.fmt.allocPrint(self.allocator, "layers.{d}.conv", .{ix});
        defer self.allocator.free(name);
        std.log.info("Testing layer: {s}", .{name});

        const in_key = try std.fmt.allocPrint(self.allocator, "{s}.in", .{name});
        defer self.allocator.free(in_key);
        var in_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, in_key, self.sharding);
        defer in_buffer.deinit();
        const in_tensor = zml.Tensor.fromShape(in_buffer.shape()).withTags(.{ .batch, .seq, .d });

        const cache_key = try std.fmt.allocPrint(self.allocator, "{s}.cache", .{name});
        defer self.allocator.free(cache_key);
        var cache_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, cache_key, self.sharding);
        defer cache_buffer.deinit();
        const cache_tensor = zml.Tensor.fromShape(cache_buffer.shape()).withTags(.{ .layer, .batch, .seq, .d });

        const cache_pos_key = try std.fmt.allocPrint(self.allocator, "{s}.cache_position", .{name});
        defer self.allocator.free(cache_pos_key);
        var cache_pos_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, cache_pos_key, self.sharding);
        defer cache_pos_buffer.deinit();
        const cache_pos_tensor = zml.Tensor.fromShape(cache_pos_buffer.shape()).withTags(.{.batch});

        const out_key = try std.fmt.allocPrint(self.allocator, "{s}.out", .{name});
        defer self.allocator.free(out_key);
        var out_buffer_expected = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, out_key, self.sharding);
        defer out_buffer_expected.deinit();

        const actual_seq_len_tensor: zml.Tensor = .init(.{}, .u32);
        const actual_seq_len: u32 = @intCast(in_tensor.dim(.seq));
        const cache_index_tensor: zml.Tensor = .init(.{}, .u32);

        const exe = try self.platform.compileFn(self.allocator, self.io, @TypeOf(layer).forward, .{ layer, in_tensor, cache_pos_tensor, actual_seq_len_tensor, model.ConvCache{ .state = cache_tensor }, cache_index_tensor, model.ConvParameters{ .is_prefill = false } }, .{ .shardings = &.{self.sharding} });
        defer exe.deinit();

        var args = try exe.args(self.allocator);
        defer args.deinit(self.allocator);
        const conv_cache: zml.Bufferized(model.ConvCache) = .{ .state = cache_buffer };

        const actual_seq_len_slice: zml.Slice = .init(zml.Shape.init(.{}, .u32), std.mem.sliceAsBytes(&[_]u32{actual_seq_len}));
        var actual_seq_len_buf: zml.Buffer = try .fromSlice(self.io, self.platform, actual_seq_len_slice, self.sharding);
        defer actual_seq_len_buf.deinit();

        const cache_index_slice: zml.Slice = .init(zml.Shape.init(.{}, .u32), std.mem.sliceAsBytes(&[_]u32{@intCast(cache_ix)}));
        var cache_index_buf: zml.Buffer = try .fromSlice(self.io, self.platform, cache_index_slice, self.sharding);
        defer cache_index_buf.deinit();

        args.set(.{ layer_buffers, in_buffer, cache_pos_buffer, actual_seq_len_buf, conv_cache, cache_index_buf });

        var res = try exe.results(self.allocator);
        defer res.deinit(self.allocator);
        exe.call(args, &res);

        var out_result = res.get(zml.Buffer);
        defer out_result.deinit();

        try zml.testing.expectClose(self.io, out_result, out_buffer_expected, opts);
        std.log.info("Layer {s} passed!", .{name});
    }

    fn testAttnLayer(self: *TestContext, ix: usize, cache_ix: usize, layer: model.Attention, layer_buffers: zml.Bufferized(model.Attention), opts: zml.testing.CompareOpts) !void {
        const name = try std.fmt.allocPrint(self.allocator, "layers.{d}.self_attn", .{ix});
        defer self.allocator.free(name);
        std.log.info("Testing layer: {s}", .{name});

        const in_key = try std.fmt.allocPrint(self.allocator, "{s}.in", .{name});
        defer self.allocator.free(in_key);
        var in_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, in_key, self.sharding);
        defer in_buffer.deinit();
        const in_tensor = zml.Tensor.fromShape(in_buffer.shape()).withTags(.{ .batch, .seq, .d });

        const key_cache_key = try std.fmt.allocPrint(self.allocator, "{s}.cache.key", .{name});
        defer self.allocator.free(key_cache_key);
        var key_cache_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, key_cache_key, self.sharding);
        defer key_cache_buffer.deinit();
        const key_cache_tensor = zml.Tensor.fromShape(key_cache_buffer.shape()).withTags(.{ .layer, .batch, .h, .k, .hd });

        const value_cache_key = try std.fmt.allocPrint(self.allocator, "{s}.cache.value", .{name});
        defer self.allocator.free(value_cache_key);
        var value_cache_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, value_cache_key, self.sharding);
        defer value_cache_buffer.deinit();
        const value_cache_tensor = zml.Tensor.fromShape(value_cache_buffer.shape()).withTags(.{ .layer, .batch, .h, .k, .hd });

        const cache_pos_key = try std.fmt.allocPrint(self.allocator, "{s}.cache_position", .{name});
        defer self.allocator.free(cache_pos_key);
        var cache_pos_buffer = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, cache_pos_key, self.sharding);
        defer cache_pos_buffer.deinit();
        const cache_pos_tensor = zml.Tensor.fromShape(cache_pos_buffer.shape()).withTags(.{.batch});

        const out_key = try std.fmt.allocPrint(self.allocator, "{s}.out", .{name});
        defer self.allocator.free(out_key);
        var out_buffer_expected = try loadBufferFromStore(self.allocator, self.io, self.platform, self.activations_store, out_key, self.sharding);
        defer out_buffer_expected.deinit();

        const cache_index_tensor: zml.Tensor = .init(.{}, .u32);

        const exe = try self.platform.compileFn(self.allocator, self.io, @TypeOf(layer).forward, .{ layer, in_tensor, cache_pos_tensor, model.KvCache{ .k = key_cache_tensor, .v = value_cache_tensor }, cache_index_tensor, self.attention_metadata, self.attention_parameters }, .{ .shardings = &.{self.sharding} });
        defer exe.deinit();

        var args = try exe.args(self.allocator);
        defer args.deinit(self.allocator);
        const kv_cache: zml.Bufferized(model.KvCache) = .{ .k = key_cache_buffer, .v = value_cache_buffer };

        var attention_metadata_buffers = try self.attention_metadata.initBuffer(self.io, self.platform, self.sharding);
        defer attention.Metadata.deinitBuffer(&attention_metadata_buffers);

        var cache_index_buf: zml.Buffer = try .scalar(self.io, self.platform, cache_ix, .u32);
        defer cache_index_buf.deinit();

        args.set(.{ layer_buffers, in_buffer, cache_pos_buffer, kv_cache, cache_index_buf, attention_metadata_buffers });

        var res = try exe.results(self.allocator);
        defer res.deinit(self.allocator);
        exe.call(args, &res);

        var out_result, var updated_kv = res.get(struct { zml.Buffer, zml.Bufferized(model.KvCache) });
        defer out_result.deinit();
        defer model.KvCache.unloadBuffers(&updated_kv);

        try zml.testing.expectClose(self.io, out_result, out_buffer_expected, opts);
        std.log.info("Layer {s} passed!", .{name});
    }
};

fn loadBufferFromStore(allocator: std.mem.Allocator, io: anytype, platform: *zml.Platform, store: *zml.io.TensorStore, key: []const u8, sharding: zml.Sharding) !zml.Buffer {
    const shape = store.view().getShape(key) orelse return error.NotFound;

    const host_bytes = try allocator.alloc(u8, shape.byteSize());
    defer allocator.free(host_bytes);

    var io_buffer: [8 * 1024]u8 = undefined;
    var reader = try store.view().getReader(key, io, &io_buffer);
    defer reader.deinit();

    _ = try reader.interface.readSliceAll(host_bytes);

    return zml.Buffer.fromBytes(io, platform, shape, sharding, host_bytes);
}