package nextgenmetering

import (
	"github.com/tidbcloud/metering/pkg/common/models"
	"go.uber.org/zap"
)

type SKUHandlerRegistry struct {
	handlers map[string]map[string]CategorySKUHandler
}

func NewSKUHandlerRegistry(gen *SKUGenerator, log *zap.Logger) *SKUHandlerRegistry {
	registry := &SKUHandlerRegistry{
		handlers: make(map[string]map[string]CategorySKUHandler),
	}

	registry.registerPremiumHandlers(gen)
	registry.registerReservedHandlers(gen)
	registry.registerEssentialHandlers(gen, log)
	registry.registerBYOCHandlers(gen)

	return registry
}

func (r *SKUHandlerRegistry) registerPremiumHandlers(gen *SKUGenerator) {
	handlers := map[string]CategorySKUHandler{
		models.NextGenCategoryResourceManager: NewPremiumResourceManagerHandler(gen),
		models.NextGenCategoryProxy:           NewPremiumProxyHandler(gen),
		models.NextGenCategoryDXF:             NewPremiumDXFHandler(gen),
		models.NextGenCategoryStorageSize:     NewPremiumStorageSizeHandler(gen),
		models.NextGenCategoryControlPlane:    NewPremiumControlPlaneHandler(gen),
		models.NextGenCategoryDFSStats:        NewPremiumDFSStatsHandler(gen),
		models.NextGenCategoryWorkerNode:      NewPremiumWorkerNodeHandler(gen),
		models.NextGenCategoryManualBackup:    NewPremiumManualBackupHandler(gen),
		models.NextGenCategoryFlexBackup:      NewPremiumFlexBackupHandler(gen),
	}
	r.handlers[models.PlanPremium] = handlers
}

func (r *SKUHandlerRegistry) registerReservedHandlers(gen *SKUGenerator) {
	handlers := map[string]CategorySKUHandler{
		models.NextGenCategoryResourceManager: NewPremiumResourceManagerHandler(gen),
		models.NextGenCategoryProxy:           NewPremiumProxyHandler(gen),
		models.NextGenCategoryDXF:             NewPremiumDXFHandler(gen),
		models.NextGenCategoryStorageSize:     NewPremiumStorageSizeHandler(gen),
		models.NextGenCategoryControlPlane:    NewPremiumControlPlaneHandler(gen),
		models.NextGenCategoryDFSStats:        NewPremiumDFSStatsHandler(gen),
		models.NextGenCategoryManualBackup:    NewPremiumManualBackupHandler(gen),
		models.NextGenCategoryFlexBackup:      NewPremiumFlexBackupHandler(gen),
	}
	r.handlers[models.PlanReserved] = handlers
}

func (r *SKUHandlerRegistry) registerEssentialHandlers(gen *SKUGenerator, log *zap.Logger) {
	handlers := map[string]CategorySKUHandler{
		models.NextGenCategoryResourceManager: NewEssentialResourceManagerHandler(gen, log),
		models.NextGenCategoryProxy:           NewEssentialProxyHandler(gen, log),
		models.NextGenCategoryStorageSize:     NewEssentialStorageSizeHandler(gen, log),
		models.NextGenCategoryDFSStats:        NewEssentialDFSStatsHandler(gen, log),
		models.NextGenCategoryFlexBackup:      NewEssentialFlexBackupHandler(gen, log),
		models.NextGenCategoryManualBackup:    NewEssentialManualBackupHandler(gen, log),
	}
	r.handlers[models.PlanEssential] = handlers
}

func (r *SKUHandlerRegistry) registerBYOCHandlers(gen *SKUGenerator) {
	handlers := map[string]CategorySKUHandler{
		models.NextGenCategoryWorkerNode: NewBYOCWorkerNodeHandler(gen),
	}
	r.handlers[models.PlanBYOC] = handlers
}

func (r *SKUHandlerRegistry) GetHandler(plan, category string) (CategorySKUHandler, bool) {
	planHandlers, ok := r.handlers[plan]
	if !ok {
		return nil, false
	}
	handler, ok := planHandlers[category]
	return handler, ok
}

func (r *SKUHandlerRegistry) ShouldProcessCategory(plan, category string) bool {
	_, ok := r.GetHandler(plan, category)
	return ok
}