package unifier_test

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/thushan/olla/internal/adapter/registry"
	"github.com/thushan/olla/internal/core/domain"
	"github.com/thushan/olla/internal/logger"
)

func createTestLogger() logger.StyledLogger {
	loggerCfg := &logger.Config{Level: "error", Theme: "default"}
	log, _, _ := logger.New(loggerCfg)
	return logger.NewPlainStyledLogger(log)
}

// Test helper functions
func strPtr(s string) *string {
	return &s
}

func int64Ptr(i int64) *int64 {
	return &i
}

func TestUnifierIntegrationWithRegistry(t *testing.T) {
	log := createTestLogger()
	ctx := context.Background()

	// Create unified registry
	unifiedRegistry := registry.NewUnifiedMemoryModelRegistry(log, nil, nil, nil)

	// Test data - models from different endpoints
	ollamaEndpoint := "http://localhost:11434"
	lmstudioEndpoint := "http://localhost:1234"

	ollamaModels := []*domain.ModelInfo{
		{
			Name: "phi4:latest",
			Details: &domain.ModelDetails{
				Family:            strPtr("phi3"), // Misclassified
				ParameterSize:     strPtr("14.7B"),
				QuantizationLevel: strPtr("Q4_K_M"),
				State:             strPtr("loaded"),
			},
			Size:     8_000_000_000,
			LastSeen: time.Now(),
		},
		{
			Name: "llama3.3:latest",
			Details: &domain.ModelDetails{
				Family:            strPtr("llama"),
				ParameterSize:     strPtr("70.6B"),
				QuantizationLevel: strPtr("Q4_K_M"),
			},
			Size:     40_000_000_000,
			LastSeen: time.Now(),
		},
		{
			Name: "hf.co/unsloth/Qwen3-32B-GGUF:Q4_K_XL",
			Details: &domain.ModelDetails{
				Family:            strPtr("qwen3"),
				ParameterSize:     strPtr("32.8B"),
				QuantizationLevel: strPtr("unknown"),
			},
			Size:     20_000_000_000,
			LastSeen: time.Now(),
		},
	}

	lmstudioModels := []*domain.ModelInfo{
		{
			Name: "microsoft/phi-4-mini-reasoning",
			Details: &domain.ModelDetails{
				Family:            strPtr("phi3"), // Also misclassified
				QuantizationLevel: strPtr("Q4_K_M"),
				MaxContextLength:  int64Ptr(131072),
				Type:              strPtr("llm"),
				State:             strPtr("not-loaded"),
			},
			Size:     8_000_000_000,
			LastSeen: time.Now(),
		},
		{
			Name: "deepseek/deepseek-r1-0528-qwen3-8b",
			Details: &domain.ModelDetails{
				Family:            strPtr("qwen3"),
				QuantizationLevel: strPtr("Q4_K_M"),
				Type:              strPtr("llm"),
			},
			Size:     5_000_000_000,
			LastSeen: time.Now(),
		},
	}

	// Register models from both endpoints
	err := unifiedRegistry.RegisterModels(ctx, ollamaEndpoint, ollamaModels)
	require.NoError(t, err)

	err = unifiedRegistry.RegisterModels(ctx, lmstudioEndpoint, lmstudioModels)
	require.NoError(t, err)

	// Wait for async unification to complete
	time.Sleep(200 * time.Millisecond)

	// Test 1: Check unified models were created
	unifiedModels, err := unifiedRegistry.GetUnifiedModels(ctx)
	require.NoError(t, err)
	assert.NotEmpty(t, unifiedModels)

	// Test 2: Verify phi4 models were unified correctly
	// In the simplified implementation, we use the original model name as ID
	// The two phi4 models have different names, so they won't be unified
	phi4Model, err := unifiedRegistry.GetUnifiedModel(ctx, "phi4:latest")
	require.NoError(t, err)
	require.NotNil(t, phi4Model)

	// Should have endpoint from Ollama only since names differ
	assert.Len(t, phi4Model.SourceEndpoints, 1)

	// Check endpoint details
	for _, ep := range phi4Model.SourceEndpoints {
		if ep.EndpointURL == ollamaEndpoint {
			assert.Equal(t, "phi4:latest", ep.NativeName)
			assert.Equal(t, "loaded", ep.State)
		}
	}

	// Test 3: Verify alias resolution
	// In simplified implementation, we look up by exact name
	testAliases := []string{
		"phi4:latest",
		"microsoft/phi-4-mini-reasoning",
	}

	for _, alias := range testAliases {
		resolved, err := unifiedRegistry.GetUnifiedModel(ctx, alias)
		assert.NoError(t, err, "Failed to resolve alias: %s", alias)
		assert.NotNil(t, resolved)
	}

	// Test 4: Check model availability
	assert.True(t, unifiedRegistry.IsModelAvailable(ctx, "phi4:latest"))
	assert.True(t, unifiedRegistry.IsModelAvailable(ctx, "microsoft/phi-4-mini-reasoning"))
	assert.False(t, unifiedRegistry.IsModelAvailable(ctx, "nonexistent:model"))

	// Test 5: Get endpoints for model
	endpoints, err := unifiedRegistry.GetEndpointsForModel(ctx, "phi4:latest")
	require.NoError(t, err)
	assert.GreaterOrEqual(t, len(endpoints), 1)

	// Test 6: Verify other models
	llamaModel, err := unifiedRegistry.GetUnifiedModel(ctx, "llama3.3:latest")
	require.NoError(t, err)
	assert.Len(t, llamaModel.SourceEndpoints, 1)

	// Test 7: Check statistics
	stats, err := unifiedRegistry.GetUnifiedStats(ctx)
	require.NoError(t, err)
	assert.Positive(t, stats.TotalUnifiedModels)
	assert.Positive(t, stats.UnificationStats.TotalUnified)
	assert.Equal(t, int64(0), stats.UnificationStats.TotalErrors)

	// Test 8: Remove endpoint and verify cleanup
	err = unifiedRegistry.RemoveEndpoint(ctx, lmstudioEndpoint)
	require.NoError(t, err)

	// Phi4 model should still exist (from Ollama) after removing LM Studio
	phi4ModelAfter, err := unifiedRegistry.GetUnifiedModel(ctx, "phi4:latest")
	require.NoError(t, err)
	assert.Len(t, phi4ModelAfter.SourceEndpoints, 1)
	assert.Equal(t, ollamaEndpoint, phi4ModelAfter.SourceEndpoints[0].EndpointURL)
}

func TestUnifierIntegrationEdgeCases(t *testing.T) {
	log := createTestLogger()
	ctx := context.Background()

	// Create unified registry
	unifiedRegistry := registry.NewUnifiedMemoryModelRegistry(log, nil, nil, nil)

	t.Run("empty model list", func(t *testing.T) {
		err := unifiedRegistry.RegisterModels(ctx, "http://localhost:11434", []*domain.ModelInfo{})
		assert.NoError(t, err)

		models, err := unifiedRegistry.GetUnifiedModels(ctx)
		assert.NoError(t, err)
		assert.Empty(t, models)
	})

	t.Run("nil models in list", func(t *testing.T) {
		models := []*domain.ModelInfo{
			{
				Name: "valid:model",
				Details: &domain.ModelDetails{
					Family:        strPtr("test"),
					ParameterSize: strPtr("7B"),
				},
			},
			nil, // Should be handled gracefully
			{
				Name: "another:model",
				Details: &domain.ModelDetails{
					Family:        strPtr("test2"),
					ParameterSize: strPtr("13B"),
				},
			},
		}

		err := unifiedRegistry.RegisterModels(ctx, "http://localhost:11434", models)
		assert.NoError(t, err)

		// Wait for async processing
		time.Sleep(50 * time.Millisecond)

		unifiedModels, err := unifiedRegistry.GetUnifiedModels(ctx)
		assert.NoError(t, err)
		assert.Len(t, unifiedModels, 2) // Only valid models
	})

	t.Run("concurrent endpoint registration", func(t *testing.T) {
		// Register the same model from multiple endpoints concurrently
		model := &domain.ModelInfo{
			Name: "concurrent:test",
			Details: &domain.ModelDetails{
				Family:            strPtr("concurrent"),
				ParameterSize:     strPtr("7B"),
				QuantizationLevel: strPtr("Q4_0"),
			},
			Size: 4_000_000_000,
		}

		endpoints := []string{
			"http://endpoint1:11434",
			"http://endpoint2:11434",
			"http://endpoint3:11434",
		}

		// Register concurrently
		var wg sync.WaitGroup
		for _, endpoint := range endpoints {
			wg.Add(1)
			go func(url string) {
				defer wg.Done()
				err := unifiedRegistry.RegisterModels(ctx, url, []*domain.ModelInfo{model})
				assert.NoError(t, err)
			}(endpoint)
		}
		wg.Wait()

		// Wait for async unification
		time.Sleep(100 * time.Millisecond)

		// Should have one unified model with multiple endpoints
		unified, err := unifiedRegistry.GetUnifiedModel(ctx, "concurrent:test")
		require.NoError(t, err)
		assert.GreaterOrEqual(t, len(unified.SourceEndpoints), 1)
	})
}

func TestUnifierWithRealWorldScenarios(t *testing.T) {
	log := createTestLogger()
	ctx := context.Background()
	unifiedRegistry := registry.NewUnifiedMemoryModelRegistry(log, nil, nil, nil)

	// Scenario 1: Model migration between endpoints
	t.Run("model migration", func(t *testing.T) {
		oldEndpoint := "http://old-server:11434"
		newEndpoint := "http://new-server:11434"

		model := &domain.ModelInfo{
			Name: "migrating:model",
			Details: &domain.ModelDetails{
				Family:            strPtr("test"),
				ParameterSize:     strPtr("7B"),
				QuantizationLevel: strPtr("Q4_K_M"),
			},
		}

		// Initially on old endpoint
		err := unifiedRegistry.RegisterModels(ctx, oldEndpoint, []*domain.ModelInfo{model})
		require.NoError(t, err)

		time.Sleep(50 * time.Millisecond)

		// Add to new endpoint
		err = unifiedRegistry.RegisterModels(ctx, newEndpoint, []*domain.ModelInfo{model})
		require.NoError(t, err)

		time.Sleep(50 * time.Millisecond)

		// Should be available on both
		endpoints, err := unifiedRegistry.GetEndpointsForModel(ctx, "migrating:model")
		require.NoError(t, err)
		assert.Len(t, endpoints, 2)

		// Remove from old endpoint
		err = unifiedRegistry.RemoveEndpoint(ctx, oldEndpoint)
		require.NoError(t, err)

		// Should only be on new endpoint
		endpoints, err = unifiedRegistry.GetEndpointsForModel(ctx, "migrating:model")
		require.NoError(t, err)
		assert.Len(t, endpoints, 1)
		assert.Equal(t, newEndpoint, endpoints[0])
	})

	// Scenario 2: Model version updates
	t.Run("model version updates", func(t *testing.T) {
		endpoint := "http://localhost:11434"

		// Register old version
		oldModel := &domain.ModelInfo{
			Name: "llama2:13b",
			Details: &domain.ModelDetails{
				Family:            strPtr("llama"),
				ParameterSize:     strPtr("13B"),
				QuantizationLevel: strPtr("Q4_0"),
				Digest:            strPtr("old-digest"),
			},
			LastSeen: time.Now().Add(-24 * time.Hour),
		}

		err := unifiedRegistry.RegisterModels(ctx, endpoint, []*domain.ModelInfo{oldModel})
		require.NoError(t, err)

		time.Sleep(50 * time.Millisecond)

		// Update with new version
		newModel := &domain.ModelInfo{
			Name: "llama2:13b",
			Details: &domain.ModelDetails{
				Family:            strPtr("llama"),
				ParameterSize:     strPtr("13B"),
				QuantizationLevel: strPtr("Q4_K_M"), // Better quantization
				Digest:            strPtr("new-digest"),
			},
			LastSeen: time.Now(),
		}

		err = unifiedRegistry.RegisterModels(ctx, endpoint, []*domain.ModelInfo{newModel})
		require.NoError(t, err)

		time.Sleep(50 * time.Millisecond)

		// Should have updated to the latest version
		unified, err := unifiedRegistry.GetUnifiedModel(ctx, "llama2:13b")
		require.NoError(t, err)
		// In simplified implementation, we don't parse quantization
		assert.NotNil(t, unified)
	})
}
