package convert

import (
	"makarna/pkg/quant"
)

type Spec struct {
	Architecture       string
	TieWordEmbeddings  bool
	BaseQuant          quant.QuantType
	MixMode            bool

	SkipTensor func(name string) bool

	IsQuantizable func(name string, shape []int, baseQuant quant.QuantType) bool

	ResolveQuant func(name string, baseQuant quant.QuantType) quant.QuantType
}

func defaultSpec(architecture string, tieWordEmbeddings bool, baseQuant quant.QuantType, mixMode bool) *Spec {
	s := &Spec{
		Architecture:      architecture,
		TieWordEmbeddings: tieWordEmbeddings,
		BaseQuant:         baseQuant,
		MixMode:           mixMode,
	}

	s.SkipTensor = func(name string) bool {
		return false
	}

	s.IsQuantizable = func(_ string, shape []int, baseQuant quant.QuantType) bool {
		if baseQuant == "" {
			return false
		}
		if len(shape) < 2 {
			return false
		}
		last := shape[len(shape)-1]
		return last%256 == 0
	}

	s.ResolveQuant = func(_ string, baseQuant quant.QuantType) quant.QuantType { return baseQuant }

	return s
}

// NewSpec creates a conversion spec and applies a registered model plugin (if any).
func NewSpec(architecture string, tieWordEmbeddings bool, baseQuant quant.QuantType, mixMode bool) *Spec {
	s := defaultSpec(architecture, tieWordEmbeddings, baseQuant, mixMode)
	if p := pluginForArchitecture(architecture); p != nil {
		p.Apply(s)
	}
	return s
}