// Package loader handles MAK file format reading and writing
package loader

// MAK Format Constants
const (
	Magic     = "MAKA"
	Version   = 1
	Alignment = 32 // Default data alignment
)

// Header is the fixed-size header at the start of MAK files (32 bytes)
type Header struct {
	Magic        [4]byte  // "MAK2"
	Version      uint16   // Format version (2)
	Flags        uint16   // Bit flags
	SectionCount uint32   // Number of sections
	Alignment    uint32   // Data alignment
	Reserved     [16]byte // Future use
}

// Header flags
const (
	FlagLittleEndian uint16 = 0x0001
	FlagBigEndian    uint16 = 0x0002
)

// SectionType identifies different sections in the file
type SectionType uint32

const (
	SectionMetadata    SectionType = 1  // Model config JSON
	SectionTensorIndex SectionType = 2  // Tensor registry
	SectionTensorData  SectionType = 3  // Raw tensor bytes
	SectionVisionEnc   SectionType = 10 // Vision encoder weights
	SectionAudioEnc    SectionType = 11 // Audio encoder (future)
	SectionTokenizer   SectionType = 20 // Embedded tokenizer
	SectionLoRA        SectionType = 30 // LoRA adapter weights
)

// SectionEntry describes a section in the section table
type SectionEntry struct {
	Type   SectionType // Section type
	Offset uint64      // Offset from file start
	Size   uint64      // Section size in bytes
	Flags  uint32      // Section-specific flags
}

// DType represents tensor data types including quantization
type DType uint16

const (
	// Float types
	F32  DType = 0
	F16  DType = 1
	BF16 DType = 2
	F64  DType = 3

	// Integer types
	I8  DType = 5
	I16 DType = 6
	I32 DType = 7
	I64 DType = 8

	// K-quants (llama.cpp style)
	Q4_K DType = 22
	Q3_K DType = 23
	Q5_K DType = 24
	Q6_K DType = 26
	Q8_K DType = 27
	Q2_K DType = 28
)

// DTypeInfo provides metadata about each data type
type DTypeInfo struct {
	Name       string
	BytesPerEl float32 // May be fractional for sub-byte quants
	BlockSize  int     // Elements per quant block (0 = no blocking)
}

var dtypeInfos = map[DType]DTypeInfo{
	F32:  {"F32", 4, 0},
	F16:  {"F16", 2, 0},
	BF16: {"BF16", 2, 0},
	F64:  {"F64", 8, 0},
	I8:   {"I8", 1, 0},
	I16:  {"I16", 2, 0},
	I32:  {"I32", 4, 0},
	I64:  {"I64", 8, 0},

	Q4_K: {"Q4_K", 0.5625, 256},     // 144 bytes per 256 elements
	Q3_K: {"Q3_K", 0.4296875, 256},  // 110 bytes per 256 elements
	Q5_K: {"Q5_K", 0.6875, 256},     // 176 bytes per 256 elements
	Q6_K: {"Q6_K", 0.8203125, 256},  // 210 bytes per 256 elements
	Q8_K: {"Q8_K", 1.140625, 256},   // 292 bytes per 256 elements
	Q2_K: {"Q2_K", 0.328125, 256},   // 84 bytes per 256 elements
}

// Info returns metadata about a DType
func (d DType) Info() DTypeInfo {
	if info, ok := dtypeInfos[d]; ok {
		return info
	}
	return DTypeInfo{Name: "unknown", BytesPerEl: 0}
}

// String returns the name of the DType
func (d DType) String() string {
	return d.Info().Name
}

// TensorEntry describes a tensor in the tensor registry
type TensorEntry struct {
	Name      string
	DType     DType
	Shape     []uint64
	Offset    uint64 // Offset within tensor data section
	Size      uint64 // Size in bytes
	QuantMeta []byte // Optional quantization metadata
}

// Metadata contains model configuration as JSON
type Metadata struct {
	ModelConfig ModelConfig            `json:"model_config"`
	Tensors     map[string]TensorEntry `json:"tensors,omitempty"` // Legacy compatibility
}

// ModelConfig holds model architecture and parameters
type ModelConfig struct {
	Architecture string                 `json:"architecture"`
	Params       map[string]interface{} `json:"params"`
}