package statistic

import (
	"testing"

	"github.com/tidbcloud/observability/o11ymetering/pkg/common/constants"
	"github.com/tidbcloud/observability/o11ymetering/pkg/service/statistic/handler"
)

func Test_pickTaskHandler(t *testing.T) {
	tests := []struct {
		name           string
		statisticClass constants.StatisticClass
		meteringGroup  constants.MeteringGroup
		bizEnv         constants.BizEnv
		expectedType   string // Used to compare the returned handler type
		shouldPanic    bool
	}{
		{
			name:           "MetricsCollection - DedicatedDataplane",
			statisticClass: constants.MetricsCollection,
			meteringGroup:  constants.DedicatedDataplane,
			bizEnv:         constants.Prod,
			expectedType:   "DedicatedDataplaneMetricCollection",
			shouldPanic:    false,
		},
		{
			name:           "MetricsCollection - Serverless",
			statisticClass: constants.MetricsCollection,
			meteringGroup:  constants.Serverless,
			bizEnv:         constants.Dev,
			expectedType:   "ServerlessMetricCollection",
			shouldPanic:    false,
		},
		{
			name:           "MetricsCollection - CloudPlatform",
			statisticClass: constants.MetricsCollection,
			meteringGroup:  constants.CloudPlatform,
			bizEnv:         constants.Staging,
			expectedType:   "CloudPlatformBaseMetricCollection",
			shouldPanic:    false,
		},
		{
			name:           "MetricsCollection - DedicatedRegionalControl",
			statisticClass: constants.MetricsCollection,
			meteringGroup:  constants.DedicatedRegionalControl,
			bizEnv:         constants.Prod,
			expectedType:   "DedicatedRegionalControlMetricCollection",
			shouldPanic:    false,
		},
		{
			name:           "LogsCollectionLoki - DedicatedDataplane",
			statisticClass: constants.LogsCollectionLoki,
			meteringGroup:  constants.DedicatedDataplane,
			bizEnv:         constants.Dev,
			expectedType:   "DedicatedDataplaneLokiCollection",
			shouldPanic:    false,
		},
		{
			name:           "LogsCollectionLoki - Serverless",
			statisticClass: constants.LogsCollectionLoki,
			meteringGroup:  constants.Serverless,
			bizEnv:         constants.Prod,
			expectedType:   "ServerlessLokiCollection",
			shouldPanic:    false,
		},
		{
			name:           "LogsCollectionLoki - CloudPlatform",
			statisticClass: constants.LogsCollectionLoki,
			meteringGroup:  constants.CloudPlatform,
			bizEnv:         constants.Staging,
			expectedType:   "CloudPlatformBaseLokiCollection",
			shouldPanic:    false,
		},
		{
			name:           "LogsCollectionLoki - DedicatedRegionalControl",
			statisticClass: constants.LogsCollectionLoki,
			meteringGroup:  constants.DedicatedRegionalControl,
			bizEnv:         constants.Prod,
			expectedType:   "DedicatedRegionalControlLokiCollection",
			shouldPanic:    false,
		},
		{
			name:           "LogsStorageS3 - DedicatedDataplane",
			statisticClass: constants.LogsStorageS3,
			meteringGroup:  constants.DedicatedDataplane,
			bizEnv:         constants.Dev,
			expectedType:   "DedicatedDataplaneS3Collection",
			shouldPanic:    false,
		},
		{
			name:           "LogsStorageS3 - Serverless",
			statisticClass: constants.LogsStorageS3,
			meteringGroup:  constants.Serverless,
			bizEnv:         constants.Prod,
			expectedType:   "ServerlessS3Collection",
			shouldPanic:    false,
		},
		{
			name:           "LogsStorageS3 - CloudPlatform",
			statisticClass: constants.LogsStorageS3,
			meteringGroup:  constants.CloudPlatform,
			bizEnv:         constants.Staging,
			expectedType:   "CloudPlatformBaseS3Collection",
			shouldPanic:    false,
		},
		{
			name:           "LogsStorageS3 - DedicatedRegionalControl",
			statisticClass: constants.LogsStorageS3,
			meteringGroup:  constants.DedicatedRegionalControl,
			bizEnv:         constants.Prod,
			expectedType:   "DedicatedRegionalControlS3Collection",
			shouldPanic:    false,
		},
		{
			name:           "DryRun",
			statisticClass: constants.DryRunStatisticClass,
			meteringGroup:  constants.DryRunMeteringGroup,
			bizEnv:         constants.Dev,
			expectedType:   "DryRun",
			shouldPanic:    false,
		},
		{
			name:           "Unhandled combination - Panics",
			statisticClass: constants.StatisticClass("UnknownClass"), // Assume a non-existent constant
			meteringGroup:  constants.MeteringGroup("UnknownGroup"),  // Assume a non-existent constant
			bizEnv:         constants.Prod,
			expectedType:   "", // Don't care about return type, because it will panic
			shouldPanic:    true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.shouldPanic {
				defer func() {
					if r := recover(); r == nil {
						t.Errorf("The code did not panic when it was expected to")
					}
				}()
				pickTaskHandler(tt.statisticClass, tt.meteringGroup, tt.bizEnv)
			} else {
				result := pickTaskHandler(tt.statisticClass, tt.meteringGroup, tt.bizEnv)

				// Use type assertion to check the returned handler type
				switch tt.expectedType {
				case "DedicatedDataplaneMetricCollection":
					if _, ok := result.(handler.DedicatedDataplaneMetricCollection); !ok {
						t.Errorf("Expected type handler.DedicatedDataplaneMetricCollection, got %T", result)
					}
				case "ServerlessMetricCollection":
					if _, ok := result.(handler.ServerlessMetricCollection); !ok {
						t.Errorf("Expected type handler.ServerlessMetricCollection, got %T", result)
					}
				case "CloudPlatformBaseMetricCollection":
					if _, ok := result.(handler.CloudPlatformBaseMetricCollection); !ok {
						t.Errorf("Expected type handler.CloudPlatformBaseMetricCollection, got %T", result)
					}
				case "DedicatedRegionalControlMetricCollection":
					if _, ok := result.(handler.DedicatedRegionalControlMetricCollection); !ok {
						t.Errorf("Expected type handler.DedicatedRegionalControlMetricCollection, got %T", result)
					}
				case "DedicatedDataplaneLokiCollection":
					if _, ok := result.(handler.DedicatedDataplaneLokiCollection); !ok {
						t.Errorf("Expected type handler.DedicatedDataplaneLokiCollection, got %T", result)
					}
				case "ServerlessLokiCollection":
					if _, ok := result.(handler.ServerlessLokiCollection); !ok {
						t.Errorf("Expected type handler.ServerlessLokiCollection, got %T", result)
					}
				case "CloudPlatformBaseLokiCollection":
					if _, ok := result.(handler.CloudPlatformBaseLokiCollection); !ok {
						t.Errorf("Expected type handler.CloudPlatformBaseLokiCollection, got %T", result)
					}
				case "DedicatedRegionalControlLokiCollection":
					if _, ok := result.(handler.DedicatedRegionalControlLokiCollection); !ok {
						t.Errorf("Expected type handler.DedicatedRegionalControlLokiCollection, got %T", result)
					}
				case "DedicatedDataplaneS3Collection":
					if _, ok := result.(handler.DedicatedDataplaneS3Collection); !ok {
						t.Errorf("Expected type handler.DedicatedDataplaneS3Collection, got %T", result)
					}
				case "ServerlessS3Collection":
					if _, ok := result.(handler.ServerlessS3Collection); !ok {
						t.Errorf("Expected type handler.ServerlessS3Collection, got %T", result)
					}
				case "CloudPlatformBaseS3Collection":
					if _, ok := result.(handler.CloudPlatformBaseS3Collection); !ok {
						t.Errorf("Expected type handler.CloudPlatformBaseS3Collection, got %T", result)
					}
				case "DedicatedRegionalControlS3Collection":
					if _, ok := result.(handler.DedicatedRegionalControlS3Collection); !ok {
						t.Errorf("Expected type handler.DedicatedRegionalControlS3Collection, got %T", result)
					}
				case "DryRun":
					if _, ok := result.(handler.DryRun); !ok {
						t.Errorf("Expected type handler.DryRun, got %T", result)
					}
				default:
					t.Errorf("Unexpected expectedType: %s", tt.expectedType)
				}

				// Check if the Env field inside the returned handler is correct
				// Note: Here we need to access the Env field based on the specific handler struct
				// Assuming all handlers implement GetEnv() method or Env field is public
				// If Env field is private, you may need to define a GetEnv() method in the handler interface
				// Or check separately after type assertion in each case
				// For convenience, assume Env field can be accessed directly
				switch v := result.(type) {
				case handler.DedicatedDataplaneMetricCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.ServerlessMetricCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.CloudPlatformBaseMetricCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.DedicatedRegionalControlMetricCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.DedicatedDataplaneLokiCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.ServerlessLokiCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.CloudPlatformBaseLokiCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.DedicatedRegionalControlLokiCollection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.DedicatedDataplaneS3Collection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.ServerlessS3Collection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.CloudPlatformBaseS3Collection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.DedicatedRegionalControlS3Collection:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				case handler.DryRun:
					if v.Env != tt.bizEnv {
						t.Errorf("Expected Env %v, got %v", tt.bizEnv, v.Env)
					}
				}
			}
		})
	}
}

// Note: Tests for ReplayStatisticTasks and getFilteredTaskHandlers require
// config initialization via LoadMeteringConfig. These should be tested in
// integration tests with a proper config file.