package create

import (
	"encoding/json"
	"testing"
)

func TestGemma4UnifiedImportTransformRegistration(t *testing.T) {
	tests := []struct {
		name       string
		configJSON string
		cfg        sourceModelConfig
		wantErr    bool
		wantLayers int
	}{
		{
			name:       "unified conditional generation architecture",
			configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"text_config":{"num_hidden_layers":48}}`,
			cfg:        sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}},
			wantLayers: 48,
		},
		{
			name:       "unified model type fallback",
			configJSON: `{"model_type":"gemma4_unified","text_config":{"num_hidden_layers":48}}`,
			cfg:        sourceModelConfig{ModelType: "gemma4_unified"},
			wantLayers: 48,
		},
		{
			name:       "malformed config is a hard error",
			configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"num_hidden_layers":"oops"}`,
			cfg:        sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}},
			wantErr:    true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			inv := Inventory{Config: tt.cfg, RawConfig: json.RawMessage(tt.configJSON)}

			transform, err := newTensorImportTransform(inv)
			if tt.wantErr {
				if err == nil {
					t.Fatalf("newTensorImportTransform() error = nil, want error")
				}
				return
			}
			if err != nil {
				t.Fatalf("newTensorImportTransform() error = %v", err)
			}

			gemmaTransform, ok := transform.(gemma4ImportTransform)
			if !ok {
				t.Fatalf("newTensorImportTransform() = %T, want gemma4ImportTransform", transform)
			}
			if gemmaTransform.numLayers != tt.wantLayers {
				t.Fatalf("numLayers = %d, want %d", gemmaTransform.numLayers, tt.wantLayers)
			}
		})
	}
}

func TestGemma4QuantizationType(t *testing.T) {
	// 26B MoE: 30 layers, 128 experts
	transform26B := gemma4ImportTransform{numLayers: 30, numExperts: 128}
	// 8-expert model (hypothetical)
	transform8E := gemma4ImportTransform{numLayers: 30, numExperts: 8}

	aligned := []int32{2816, 2816} // divisible by 64 (int4/int8 group size) and 16 (nvfp4)

	tests := []struct {
		name      string
		transform gemma4ImportTransform
		tensor    string
		shape     []int32
		quantize  string
		want      string
	}{
		// === embed_tokens: quantize to 8-bit variant (serves as both embed and lm_head) ===
		{"embed_tokens int4", transform26B, "model.embed_tokens.weight", aligned, "int4", "int8"},
		{"embed_tokens nvfp4", transform26B, "model.embed_tokens.weight", aligned, "nvfp4", "mxfp8"},
		{"embed_tokens mxfp4", transform26B, "model.embed_tokens.weight", aligned, "mxfp4", "mxfp8"},
		{"embed_tokens int8", transform26B, "model.embed_tokens.weight", aligned, "int8", "int8"},
		{"embed_tokens mxfp8", transform26B, "model.embed_tokens.weight", aligned, "mxfp8", "mxfp8"},

		// === v_proj: layer-position heuristic for int4/nvfp4 ===
		// Layer 0 is in first 1/8 (30/8=3) → promoted
		{"v_proj int4 promoted layer", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "int4", "int8"},
		// Layer 4 is NOT in useMoreBits → base quant
		{"v_proj int4 non-promoted layer", transform26B, "model.layers.4.self_attn.v_proj.weight", aligned, "int4", "int4"},
		// Layer 29 is in last 1/8 → promoted
		{"v_proj int4 last layer promoted", transform26B, "model.layers.29.self_attn.v_proj.weight", aligned, "int4", "int8"},
		// nvfp4: promote to mxfp8 (cross-family, validated by MLX quantized_matmul)
		{"v_proj nvfp4 promoted layer", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "nvfp4", "mxfp8"},
		{"v_proj nvfp4 non-promoted layer", transform26B, "model.layers.4.self_attn.v_proj.weight", aligned, "nvfp4", "nvfp4"},
		// mxfp4: promoted to mxfp8 at promoted layers (same mxfp family)
		{"v_proj mxfp4 promoted layer", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "mxfp4", "mxfp8"},
		{"v_proj mxfp4 non-promoted layer", transform26B, "model.layers.4.self_attn.v_proj.weight", aligned, "mxfp4", "mxfp4"},
		// int8/mxfp8: no promotion (already 8-bit)
		{"v_proj int8 base", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "int8", "int8"},
		{"v_proj mxfp8 base", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "mxfp8", "mxfp8"},

		// === down_proj (dense MLP): same heuristic as v_proj ===
		{"dense down_proj int4 promoted", transform26B, "model.layers.0.mlp.down_proj.weight", aligned, "int4", "int8"},
		{"dense down_proj int4 non-promoted", transform26B, "model.layers.4.mlp.down_proj.weight", aligned, "int4", "int4"},
		{"dense down_proj nvfp4 promoted", transform26B, "model.layers.0.mlp.down_proj.weight", aligned, "nvfp4", "mxfp8"},
		{"dense down_proj nvfp4 non-promoted", transform26B, "model.layers.4.mlp.down_proj.weight", aligned, "nvfp4", "nvfp4"},
		{"dense down_proj mxfp4 promoted", transform26B, "model.layers.0.mlp.down_proj.weight", aligned, "mxfp4", "mxfp8"},
		{"dense down_proj mxfp4 non-promoted", transform26B, "model.layers.4.mlp.down_proj.weight", aligned, "mxfp4", "mxfp4"},

		// === Expert down_proj: int4→int8, nvfp4→nvfp8 at promoted layers ===
		{"expert down_proj int4 promoted", transform26B, "model.layers.0.moe.experts.42.down_proj.weight", aligned, "int4", "int8"},
		{"expert down_proj int4 non-promoted", transform26B, "model.layers.4.moe.experts.42.down_proj.weight", aligned, "int4", "int4"},
		// nvfp4 experts: promote to mxfp8 (all experts at a layer get same treatment,
		// so GatherQMM sees uniform quant per projection per layer)
		{"expert down_proj nvfp4 promoted layer", transform26B, "model.layers.0.moe.experts.42.down_proj.weight", aligned, "nvfp4", "mxfp8"},
		{"expert down_proj nvfp4 non-promoted layer", transform26B, "model.layers.4.moe.experts.42.down_proj.weight", aligned, "nvfp4", "nvfp4"},
		// mxfp4 experts: promote to mxfp8 (same mxfp family, GatherQMM compatible)
		{"expert down_proj mxfp4 promoted layer", transform26B, "model.layers.0.moe.experts.42.down_proj.weight", aligned, "mxfp4", "mxfp8"},
		{"expert down_proj mxfp4 non-promoted layer", transform26B, "model.layers.4.moe.experts.42.down_proj.weight", aligned, "mxfp4", "mxfp4"},

		// === Expert gate_up_proj: always base quant (not a sensitive tensor) ===
		{"expert gate_up int4", transform26B, "model.layers.0.moe.experts.42.gate_up_proj.weight", aligned, "int4", "int4"},
		{"expert gate_up nvfp4", transform26B, "model.layers.0.moe.experts.42.gate_up_proj.weight", aligned, "nvfp4", "nvfp4"},
		{"expert gate_up mxfp4", transform26B, "model.layers.0.moe.experts.42.gate_up_proj.weight", aligned, "mxfp4", "mxfp4"},

		// === Router projection: expert selection is sensitive; keep source precision ===
		{"router proj int4", transform26B, "model.layers.0.router.proj.weight", aligned, "int4", ""},
		{"router proj nvfp4", transform26B, "model.layers.0.router.proj.weight", aligned, "nvfp4", ""},
		{"router proj mxfp4", transform26B, "model.layers.0.router.proj.weight", aligned, "mxfp4", ""},

		// === k_proj: promoted only for 8-expert models ===
		{"k_proj 128 experts int4", transform26B, "model.layers.0.self_attn.k_proj.weight", aligned, "int4", "int4"},
		{"k_proj 8 experts int4", transform8E, "model.layers.0.self_attn.k_proj.weight", aligned, "int4", "int8"},
		{"k_proj 8 experts nvfp4", transform8E, "model.layers.0.self_attn.k_proj.weight", aligned, "nvfp4", "mxfp8"},
		{"k_proj 8 experts mxfp4", transform8E, "model.layers.0.self_attn.k_proj.weight", aligned, "mxfp4", "mxfp8"},

		// === q_proj, o_proj, gate_proj, up_proj: always base quant ===
		{"q_proj int4", transform26B, "model.layers.0.self_attn.q_proj.weight", aligned, "int4", "int4"},
		{"o_proj int4", transform26B, "model.layers.0.self_attn.o_proj.weight", aligned, "int4", "int4"},
		{"gate_proj int4", transform26B, "model.layers.0.mlp.gate_proj.weight", aligned, "int4", "int4"},
		{"up_proj int4", transform26B, "model.layers.0.mlp.up_proj.weight", aligned, "int4", "int4"},

		// === Non-quantizable tensors: always bf16 ===
		{"embed_tokens per_layer skip", transform26B, "model.embed_tokens_per_layer.weight", aligned, "int4", ""},
		{"norm", transform26B, "model.layers.0.input_layernorm.weight", []int32{2816}, "int4", ""},
		{"router scale", transform26B, "model.layers.0.router.scale", []int32{2816}, "int4", ""},

		// === Audio/vision tower tensors: must pass through unquantized for all quant types ===
		// These contain .v_proj and down_proj but should NOT be intercepted by
		// the sensitive-tensor promotion logic.
		{"audio norm int4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "int4", ""},
		{"audio norm nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "nvfp4", ""},
		{"audio norm int8", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "int8", ""},
		{"audio norm mxfp8", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "mxfp8", ""},
		{"audio conv int4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.conv.weight", []int32{128, 1, 3, 3}, "int4", ""},
		{"audio conv nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.conv.weight", []int32{128, 1, 3, 3}, "nvfp4", ""},
		{"audio linear int4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "int4", ""},
		{"audio linear nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "nvfp4", ""},
		// Audio tower v_proj — must NOT be promoted despite containing .v_proj
		{"audio v_proj int4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", ""},
		{"audio v_proj nvfp4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", ""},
		// Vision tower v_proj — vision tower IS quantized (unlike audio tower),
		// but not intercepted by gemma4's layer-position heuristic.
		// Falls through to GetTensorQuantization which applies uniform promotion.
		{"vision v_proj int4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", "int8"},
		{"vision v_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", "mxfp8"},
		// Audio tower down_proj
		{"audio down_proj int4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "int4", ""},
		{"audio down_proj nvfp4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "nvfp4", ""},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := tt.transform.quantizationType(tt.tensor, tt.shape, tt.quantize)
			if got != tt.want {
				t.Errorf("quantizationType(%q, %v, %q) = %q, want %q",
					tt.tensor, tt.shape, tt.quantize, got, tt.want)
			}
		})
	}
}

func TestUseMoreBits(t *testing.T) {
	// 30 layers: first 1/8 = layers 0-2, last 1/8 = layers 27-29
	// In between: every 3rd from offset (i - n/8) % 3 == 2
	n := 30
	promoted := map[int]bool{}
	for i := range n {
		if useMoreBits(i, n) {
			promoted[i] = true
		}
	}

	// First 1/8 (30/8 = 3): layers 0, 1, 2
	for _, i := range []int{0, 1, 2} {
		if !promoted[i] {
			t.Errorf("layer %d should be promoted (first 1/8)", i)
		}
	}

	// Last 1/8: layers 26, 27, 28, 29 (>= 7*30/8 = 26)
	for _, i := range []int{26, 27, 28, 29} {
		if !promoted[i] {
			t.Errorf("layer %d should be promoted (last 1/8)", i)
		}
	}

	// Some middle layers should NOT be promoted
	for _, i := range []int{3, 4, 6, 7} {
		if promoted[i] {
			t.Errorf("layer %d should NOT be promoted", i)
		}
	}

	// Layer 5 should be promoted: (5 - 3) % 3 == 2
	if !promoted[5] {
		t.Errorf("layer 5 should be promoted (periodic)")
	}
}