package autogen

import "sort"

// tensorScan is everything one walk of the GGUF tensor-info section yields.
// The walk is not free (it reads a name + dims + type for every tensor in the
// file), so each new question about the weights is answered by adding a field
// here rather than by a second pass.
type tensorScan struct {
	expertShare float64
	vocabElems  int64
	diffKind    string
	bakedEnc    bool
	// condHidden is the hidden width of the TEXT ENCODER a diffusion transformer
	// expects, read off the tensor that projects caption embeddings into the DiT
	// (ne[0] is that projection input width). It is what lets the encoder pool
	// pick the right --llm/--t5xxl with no per-model table: 4096 = T5-XXL,
	// 3584 = Qwen2.5-VL-7B, 3072 = Ministral-3B, 2560 = Qwen3-4B. 0 for a file
	// with no recognised projection (every non-diffusion model).
	condHidden int64

	// videoKind names the video-diffusion family the tensor layout marks, or ""
	// for anything that is not a video DiT. Structural, not metadata: MiniMax-H3
	// ships a gguf with ZERO KVs (no general.architecture at all), and the one
	// arch string that does say "wan" belongs to ERNIE-Image-Turbo, an IMAGE
	// model. Only the tensor table can tell these apart.
	videoKind string

	// embedWidth is the model's hidden size read off the token-embedding tensor
	// (ne[0] of token_embd.weight, or of HF's model.embed_tokens.weight), and is
	// the fallback for a text encoder whose gguf declares no hyperparameter KVs
	// at all. LTX-2.5's Gemma-4 encoder is converted by ComfyUI tooling that
	// keeps every tensor and writes NO gemma4.embedding_length: without this the
	// file reads as width 0 and the encoder pool discards it as "not a text
	// encoder", silently leaving an LTX model to condition on a stock Gemma.
	embedWidth int64

	// hasAudioOut is true when the video DiT also denoises an audio latent
	// (MiniMax-H3's audio_patch_proj / final_layer.audio_out), i.e. it needs a
	// second --audio-vae and produces a soundtrack.
	hasAudioOut bool

	// typeBytes is on-disk bytes per ggml type id. A type missing from
	// ggmlTypeSize can't be weighed, so it is recorded with 0 bytes: its
	// PRESENCE still has to be visible to quantLabelFrom, which reports an
	// unweighable file as unlabelled rather than as the type it could measure.
	typeBytes map[uint32]int64

	// paramCount is the element count summed over every tensor (Metadata.ParamCount).
	paramCount int64
}

// ggmlTypeName spells a ggml type id the way llama-quantize names it. Ids are
// ggml's own enum (ggml/include/ggml.h), the same numbering ggmlTypeSize keys
// on — keep the two tables in step: a type that can be weighed but not named
// makes an unlabelled file, and one named but not weighed makes a wrong label.
var ggmlTypeName = map[uint32]string{
	0: "F32", 1: "F16", 2: "Q4_0", 3: "Q4_1", 6: "Q5_0", 7: "Q5_1",
	8: "Q8_0", 9: "Q8_1", 10: "Q2_K", 11: "Q3_K", 12: "Q4_K", 13: "Q5_K",
	14: "Q6_K", 15: "Q8_K", 16: "IQ2_XXS", 17: "IQ2_XS", 18: "IQ3_XXS",
	19: "IQ1_S", 20: "IQ4_NL", 21: "IQ3_S", 22: "IQ2_S", 23: "IQ4_XS",
	24: "I8", 25: "I16", 26: "I32", 27: "I64", 28: "F64", 29: "IQ1_M", 30: "BF16", 34: "TQ1_0", 35: "TQ2_0",
	39: "MXFP4", 40: "NVFP4", 41: "Q1_0",
}

// filler types carry no information about how a file was quantized: F32 is the
// norm/bias residue every quant leaves behind (and the type a converter writes
// a 1-D tensor as regardless), and the integer types are indices, not weights.
// Counting them would label a small model by its bias tensors.
var fillerType = map[uint32]bool{0: true, 24: true, 25: true, 26: true, 27: true}

// quantDominance is the share of weight bytes one type must hold to name the
// file by itself. Measured against a real model tree rather than guessed: a
// named recipe leaves its headline type 71-100% of the bytes (Q4_K_M sits at
// 71%, the rest Q6_K; a full SDXL checkpoint at 75%, the rest its F16 text
// encoders), while a file with no single answer spreads much wider - an
// unsloth UD-Q4_K_XL peaks at 45%, a hand-mixed one at 55%. 0.7 is the gap
// between those two populations, not a round number.
const quantDominance = 0.7

// quantLabelFrom names a file's quantization from what its tensors ACTUALLY
// are, rather than from its filename or its self-declared general.file_type
// (which lies: files shipped as "MXFP4" declare Q8_0, and a hand-mixed file
// declares whatever the last quantize pass set). The heaviest weight type wins;
// when nothing holds quantDominance of the bytes, the label says so ("IQ4_XS
// mix") because there is no single honest answer.
//
// Returns "" when the file cannot be judged - no weighable tensors, or a type
// present that this build can neither name nor size - so callers fall back to
// the name rules instead of showing a confident wrong answer.
func quantLabelFrom(typeBytes map[uint32]int64) string {
	var total int64
	var top uint32
	var topBytes int64
	seen := false
	for typ, b := range typeBytes {
		if _, ok := ggmlTypeName[typ]; !ok {
			return ""
		}
		if fillerType[typ] {
			continue
		}
		seen = true
		total += b
		// Ties break on the lower id (the coarser quant), so a file split
		// exactly between two types labels the same way on every read - map
		// iteration order is random.
		if b > topBytes || (b == topBytes && typ < top) {
			top, topBytes = typ, b
		}
	}
	switch {
	case !seen:
		// An all-F32/F16-of-1-D-tensors file: nothing but filler. Name it by
		// the filler anyway rather than reporting nothing.
		return fillerLabel(typeBytes)
	case total <= 0:
		return ""
	case float64(topBytes)/float64(total) >= quantDominance:
		return ggmlTypeName[top]
	}
	return ggmlTypeName[top] + " mix"
}

// fillerLabel names a file made only of types quantLabelFrom skips - an
// unquantized F32 export, or a tensor section of pure indices.
func fillerLabel(typeBytes map[uint32]int64) string {
	ids := make([]uint32, 0, len(typeBytes))
	for typ := range typeBytes {
		ids = append(ids, typ)
	}
	sort.Slice(ids, func(i, j int) bool {
		if typeBytes[ids[i]] != typeBytes[ids[j]] {
			return typeBytes[ids[i]] > typeBytes[ids[j]]
		}
		return ids[i] < ids[j]
	})
	if len(ids) == 0 {
		return ""
	}
	return ggmlTypeName[ids[0]]
}

// condTensors are the tensors that project text-encoder embeddings into a
// diffusion transformer, ranked most to least specific. Their ne[0] is the
// encoder hidden width the model was trained against, which is the only
// machine-readable statement of "which text encoder do I need". Verified against
// flux (txt_in 4096 = T5-XXL), LongCat and Qwen-Image-Edit (txt_in 3584 =
// Qwen2.5-VL-7B), Z-Image (cap_embedder 2560 = Qwen3-4B), ERNIE (text_proj 3072
// = Ministral-3B) and Krea2 (txtfusion 2560 = Qwen3-VL-4B).
//
// Qwen-Image 2.1 is the first to wrap the projection in a two-layer MLP, so its
// width lives in txt_in.IN_LAYER.weight (4096 = Qwen3-VL-8B) and the bare
// txt_in.weight the 20B line uses is absent. Without this entry the file states
// no width at all: its GGUF carries three KV pairs and no diffusers config blob,
// so captionChannelsFrom has nothing to fall back on and the model is emitted
// with a "needs encoder [llm]" warning that no download can clear.
//
// MOSS-Image projects through text_embedder.proj (4096 = Qwen3.5-9B, a
// recurrent hybrid; see EncoderPool.RecurrentLlm).
var condTensorOrder = []string{
	"txt_in.weight",
	"txt_in.in_layer.weight",
	"cap_embedder.1.weight",
	"text_proj.weight",
	"condition_proj.weight",
	"context_embedder.weight",
	"txtfusion.layerwise_blocks.0.prenorm.scale",
	"text_embedder.proj.weight",
}

var condTensors = func() map[string]struct{} {
	m := make(map[string]struct{}, len(condTensorOrder))
	for _, n := range condTensorOrder {
		m[n] = struct{}{}
	}
	return m
}()

// condHiddenFrom picks the highest-ranked caption projection actually present.
func condHiddenFrom(dims map[string]int64) int64 {
	for _, n := range condTensorOrder {
		if d := dims[n]; d > 0 {
			return d
		}
	}
	return 0
}