package memory

import (
	"fmt"
	"strings"

	candle_binding "github.com/vllm-project/semantic-router/candle-binding"
)

// EmbeddingModelType represents supported embedding model types
type EmbeddingModelType string

const (
	EmbeddingModelBERT   EmbeddingModelType = "bert"
	EmbeddingModelMMBERT EmbeddingModelType = "mmbert"
	EmbeddingModelMulti  EmbeddingModelType = "multimodal"
	EmbeddingModelQwen3  EmbeddingModelType = "qwen3"
	EmbeddingModelGemma  EmbeddingModelType = "gemma"
)

// EmbeddingConfig holds the embedding model configuration
type EmbeddingConfig struct {
	Model     EmbeddingModelType
	Dimension int // Target dimension for Matryoshka models (default: 256 for mmbert)
	Layer     int // Target layer for 2D Matryoshka early exit (0 = full model, recommended for search/RAG)
}

// GenerateEmbedding generates an embedding using the configured model
func GenerateEmbedding(text string, cfg EmbeddingConfig) ([]float32, error) {
	if deterministicEmbeddingsEnabled() {
		return generateDeterministicEmbedding(text, cfg), nil
	}

	modelName := strings.ToLower(strings.TrimSpace(string(cfg.Model)))
	switch modelName {
	case "qwen3":
		return generateQwen3Embedding(text, modelName)
	case "gemma":
		return generateGemmaEmbedding(text, modelName)
	case "mmbert":
		return generateMmBertEmbedding(text, modelName, cfg)
	case "multimodal":
		return generateMultiModalEmbedding(text, modelName, cfg)
	case "bert", "":
		return generateBertEmbedding(text)
	default:
		return nil, fmt.Errorf("unsupported embedding model: %s (must be 'bert', 'qwen3', 'gemma', 'mmbert', or 'multimodal')", modelName)
	}
}

func generateQwen3Embedding(text, modelName string) ([]float32, error) {
	output, err := candle_binding.GetEmbeddingBatched(text, modelName, 0)
	if err != nil {
		return nil, fmt.Errorf("qwen3 embedding failed: %w", err)
	}
	return output.Embedding, nil
}

func generateGemmaEmbedding(text, modelName string) ([]float32, error) {
	output, err := candle_binding.GetEmbeddingWithModelType(text, modelName, 0)
	if err != nil {
		return nil, fmt.Errorf("gemma embedding failed: %w", err)
	}
	return output.Embedding, nil
}

func generateMmBertEmbedding(text, modelName string, cfg EmbeddingConfig) ([]float32, error) {
	targetDim := cfg.Dimension
	if targetDim <= 0 {
		targetDim = 256
	}
	output, err := candle_binding.GetEmbedding2DMatryoshka(text, modelName, cfg.Layer, targetDim)
	if err != nil {
		return nil, fmt.Errorf("mmbert embedding failed: %w", err)
	}
	return output.Embedding, nil
}

func generateMultiModalEmbedding(text, modelName string, cfg EmbeddingConfig) ([]float32, error) {
	targetDim := cfg.Dimension
	if targetDim <= 0 {
		targetDim = 384
	}
	output, err := candle_binding.GetEmbeddingWithModelType(text, modelName, targetDim)
	if err != nil {
		return nil, fmt.Errorf("multimodal embedding failed: %w", err)
	}
	return output.Embedding, nil
}

func generateBertEmbedding(text string) ([]float32, error) {
	embedding, err := candle_binding.GetEmbedding(text, 0)
	if err != nil {
		return nil, fmt.Errorf("bert embedding failed: %w", err)
	}
	return embedding, nil
}
