package classification

import (
	"errors"
	"strings"
	"testing"

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

func prefBoolPtr(value bool) *bool {
	return &value
}

func TestPreferenceClassifier_ContrastiveFewShot(t *testing.T) {
	reset := SetEmbeddingFuncForTests(func(text string, modelType string, targetDim int) (*candle_binding.EmbeddingOutput, error) {
		lower := strings.ToLower(text)
		switch {
		case strings.Contains(lower, "bug") || strings.Contains(lower, "fix"):
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.2, 0.9}}, nil
		case strings.Contains(lower, "code"):
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.9, 0.1}}, nil
		default:
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.5, 0.5}}, nil
		}
	})
	defer reset()

	rules := []config.PreferenceRule{
		{Name: "code_generation", Description: "Writes code", Examples: []string{"write python code"}},
		{Name: "bug_fixing", Description: "Finds and fixes bugs", Examples: []string{"fix this bug"}},
	}

	localCfg := &config.PreferenceModelConfig{
		UseContrastive: prefBoolPtr(true),
		EmbeddingModel: "qwen3",
	}

	classifier, err := NewPreferenceClassifier(nil, rules, localCfg)
	if err != nil {
		t.Fatalf("failed to create contrastive preference classifier: %v", err)
	}

	conversation := `[{"role":"user","content":"can you help fix this bug?"}]`
	result, err := classifier.Classify(conversation)
	if err != nil {
		t.Fatalf("contrastive classification failed: %v", err)
	}

	if result.Preference != "bug_fixing" {
		t.Fatalf("expected bug_fixing preference, got %s", result.Preference)
	}

	if result.Confidence <= 0 {
		t.Fatalf("expected positive confidence score, got %f", result.Confidence)
	}
}

func TestContrastivePreferenceClassifier_UsesDescriptionsWhenNoExamples(t *testing.T) {
	reset := SetEmbeddingFuncForTests(func(text string, modelType string, targetDim int) (*candle_binding.EmbeddingOutput, error) {
		switch text {
		case "Writes code":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{1, 0}}, nil
		case "Fixes bugs":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0, 1}}, nil
		case "please write code":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{1, 0}}, nil
		default:
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.1, 0.1}}, nil
		}
	})
	defer reset()

	rules := []config.PreferenceRule{
		{Name: "code_generation", Description: "Writes code"},
		{Name: "bug_fixing", Description: "Fixes bugs"},
	}

	localCfg := &config.PreferenceModelConfig{
		UseContrastive: prefBoolPtr(true),
	}

	classifier, err := NewPreferenceClassifier(nil, rules, localCfg)
	if err != nil {
		t.Fatalf("failed to create contrastive classifier: %v", err)
	}

	conversation := `[{"role":"user","content":"please write code"}]`
	result, err := classifier.Classify(conversation)
	if err != nil {
		t.Fatalf("classification failed: %v", err)
	}

	if result.Preference != "code_generation" {
		t.Fatalf("expected code_generation, got %s", result.Preference)
	}
}

func TestContrastivePreferenceClassifier_NoExamplesError(t *testing.T) {
	rules := []config.PreferenceRule{{Name: "code_generation"}}
	localCfg := &config.PreferenceModelConfig{UseContrastive: prefBoolPtr(true)}

	if _, err := NewPreferenceClassifier(nil, rules, localCfg); err == nil {
		t.Fatalf("expected error when no descriptions or examples are provided")
	}
}

func TestContrastivePreferenceClassifier_EmptyText(t *testing.T) {
	reset := SetEmbeddingFuncForTests(func(text string, modelType string, targetDim int) (*candle_binding.EmbeddingOutput, error) {
		return &candle_binding.EmbeddingOutput{Embedding: []float32{0.1, 0.2}}, nil
	})
	defer reset()

	rules := []config.PreferenceRule{{Name: "code_generation", Description: "Writes code"}}
	localCfg := &config.PreferenceModelConfig{UseContrastive: prefBoolPtr(true)}

	classifier, err := NewPreferenceClassifier(nil, rules, localCfg)
	if err != nil {
		t.Fatalf("failed to create contrastive classifier: %v", err)
	}

	if _, err := classifier.Classify(""); err == nil {
		t.Fatalf("expected error for empty text")
	}
}

func TestPreferenceClassifier_DefaultsToContrastiveWhenConfigOmitted(t *testing.T) {
	reset := SetEmbeddingFuncForTests(func(text string, modelType string, targetDim int) (*candle_binding.EmbeddingOutput, error) {
		switch text {
		case "Writes code":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{1, 0}}, nil
		case "please write code":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{1, 0}}, nil
		default:
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.1, 0.1}}, nil
		}
	})
	defer reset()

	rules := []config.PreferenceRule{{Name: "code_generation", Description: "Writes code"}}

	classifier, err := NewPreferenceClassifier(nil, rules, nil)
	if err != nil {
		t.Fatalf("failed to create default contrastive preference classifier: %v", err)
	}

	result, err := classifier.Classify(`[{"role":"user","content":"please write code"}]`)
	if err != nil {
		t.Fatalf("classification failed: %v", err)
	}

	if result.Preference != "code_generation" {
		t.Fatalf("expected code_generation, got %s", result.Preference)
	}
}

func TestPreferenceClassifier_ParsePreferenceOutput(t *testing.T) {
	classifier := &PreferenceClassifier{}

	tests := []struct {
		name       string
		output     string
		want       string
		wantErrSub string
	}{
		{
			name:   "plain JSON",
			output: `{"route":"code_generation"}`,
			want:   "code_generation",
		},
		{
			name:   "prose wrapped JSON",
			output: `The best route is {"route":"bug_fixing"}.`,
			want:   "bug_fixing",
		},
		{
			name:   "single quoted JSON",
			output: `{'route':'other'}`,
			want:   "other",
		},
		{
			name:       "missing JSON",
			output:     "route: code_generation",
			wantErrSub: "no valid JSON",
		},
		{
			name:       "empty route",
			output:     `{"route":""}`,
			wantErrSub: "preference field is empty",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result, err := classifier.parsePreferenceOutput(tt.output)
			if tt.wantErrSub != "" {
				if err == nil || !strings.Contains(err.Error(), tt.wantErrSub) {
					t.Fatalf("expected error containing %q, got result=%+v err=%v", tt.wantErrSub, result, err)
				}
				return
			}

			if err != nil {
				t.Fatalf("unexpected parse error: %v", err)
			}
			if result.Preference != tt.want {
				t.Fatalf("expected preference %q, got %+v", tt.want, result)
			}
		})
	}
}

func TestContrastivePreferenceClassifier_BelowThresholdReturnsNoMatchError(t *testing.T) {
	reset := SetEmbeddingFuncForTests(func(text string, modelType string, targetDim int) (*candle_binding.EmbeddingOutput, error) {
		switch text {
		case "Writes code":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{1, 0}}, nil
		case "please help with this task":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.6, 0.8}}, nil
		default:
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.1, 0.1}}, nil
		}
	})
	defer reset()

	rules := []config.PreferenceRule{{
		Name:        "code_generation",
		Description: "Writes code",
		Threshold:   0.7,
	}}

	localCfg := &config.PreferenceModelConfig{
		UseContrastive: prefBoolPtr(true),
	}

	classifier, err := NewPreferenceClassifier(nil, rules, localCfg)
	if err != nil {
		t.Fatalf("failed to create contrastive classifier: %v", err)
	}

	result, err := classifier.Classify(`[{"role":"user","content":"please help with this task"}]`)
	if err == nil {
		t.Fatal("expected below-threshold preference classification to return an error")
	}
	if result != nil {
		t.Fatalf("expected nil result on below-threshold no-match, got %+v", result)
	}
	if !errors.Is(err, ErrPreferenceBelowThreshold) {
		t.Fatalf("expected ErrPreferenceBelowThreshold, got %v", err)
	}
}

func TestContrastivePreferenceClassifier_MarginThresholdRejectsAmbiguousWinner(t *testing.T) {
	reset := SetEmbeddingFuncForTests(func(text string, modelType string, targetDim int) (*candle_binding.EmbeddingOutput, error) {
		switch text {
		case "Writes code":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{1, 0}}, nil
		case "Fixes bugs":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.98, 0.2}}, nil
		case "please help with implementation":
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.99, 0.05}}, nil
		default:
			return &candle_binding.EmbeddingOutput{Embedding: []float32{0.1, 0.1}}, nil
		}
	})
	defer reset()

	classifier, err := NewContrastivePreferenceClassifierWithConfig(
		[]config.PreferenceRule{
			{Name: "code_generation", Description: "Writes code"},
			{Name: "bug_fixing", Description: "Fixes bugs"},
		},
		"qwen3",
		config.PrototypeScoringConfig{MarginThreshold: 0.01},
	)
	if err != nil {
		t.Fatalf("failed to create contrastive classifier: %v", err)
	}

	details, err := classifier.ClassifyDetailed("please help with implementation")
	if err != nil {
		t.Fatalf("ClassifyDetailed failed: %v", err)
	}
	if details.Margin >= 0.01 {
		t.Fatalf("expected ambiguous preference margin below threshold, got %+v", details)
	}

	result, err := classifier.Classify("please help with implementation")
	if err == nil {
		t.Fatalf("expected margin threshold rejection, got result %+v", result)
	}
	if !errors.Is(err, ErrPreferenceBelowThreshold) {
		t.Fatalf("expected ErrPreferenceBelowThreshold, got %v", err)
	}
}
