package case_driver

import (
	"sync"

	"github.com/tidbcloud/dataflow-service-e2e/cases"
	"github.com/tidbcloud/dataflow-service-e2e/internal/log"
	"github.com/tidbcloud/dataflow-service-e2e/internal/misc"
	. "github.com/tidbcloud/dataflow-service-e2e/internal/model"
)

type Dispatcher interface {
	RegisterCase(caseName string, theCase cases.E2ECase) <-chan error
	Dispatch(entities ...DependencyEntity)
	SetDependencyError(depType string, err error)
}

type depDispatcher struct {
	lgr log.SimpleLogger

	mu sync.Mutex

	partitions     map[string]*depPartition
	depSubscribers map[string]*depSubscriber

	dependencyErrors map[string]error
	dependencyDone   map[string]bool

	onAssigned func(caseName, reqName string, entity DependencyEntity)
}

type depPartition struct {
	depType              string
	dimensions           map[string]string
	hasSharedRequirement bool

	sharedEntity      DependencyEntity
	dedicatedEntities []DependencyEntity
	nextDedicatedIdx  int
}

type depSubscriber struct {
	caseName string
	theCase  cases.E2ECase

	pending  map[string]DependencyRequirement
	resultCh chan error
	closed   bool
}

func newDepDispatcher(
	lgr log.SimpleLogger,
	onAssigned func(caseName, reqName string, entity DependencyEntity),
) *depDispatcher {
	a := &depDispatcher{
		lgr:              lgr,
		partitions:       map[string]*depPartition{},
		depSubscribers:   map[string]*depSubscriber{},
		dependencyErrors: map[string]error{},
		dependencyDone:   map[string]bool{},
		onAssigned:       onAssigned,
	}

	return a
}

func (a *depDispatcher) RegisterCase(caseName string, theCase cases.E2ECase) <-chan error {
	a.mu.Lock()
	defer a.mu.Unlock()

	subscriber := a.newDepSubscriberLocked(caseName, theCase)
	for reqName, req := range theCase.DeclareDependencyRequirements() {
		subscriber.pending[reqName] = req
		partition := a.getOrCreatePartitionLocked(req.Type(), req.Dimensions())
		if !req.IsDedicated() {
			partition.hasSharedRequirement = true
		}
	}

	if a.tryAssignLocked(subscriber) {
		a.markCaseReadyLocked(subscriber, nil)
	}
	return subscriber.resultCh
}

func (a *depDispatcher) Dispatch(entities ...DependencyEntity) {
	a.mu.Lock()
	defer a.mu.Unlock()

	for _, entity := range entities {
		partition := a.getOrCreatePartitionLocked(entity.Type(), entity.Dimensions())
		if partition.hasSharedRequirement && partition.sharedEntity == nil {
			partition.sharedEntity = entity
			continue
		}
		partition.dedicatedEntities = append(partition.dedicatedEntities, entity)
	}

	var ready []*depSubscriber
	for _, subscriber := range a.depSubscribers {
		if a.tryAssignLocked(subscriber) {
			ready = append(ready, subscriber)
		}
	}
	for _, subscriber := range ready {
		a.markCaseReadyLocked(subscriber, nil)
	}
}

func (a *depDispatcher) SetDependencyError(depType string, err error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.dependencyErrors[depType] = err
	for _, subscriber := range a.depSubscribers {
		if subscriber.closed {
			continue
		}
		for _, req := range subscriber.pending {
			if req.Type() == depType {
				a.markCaseReadyLocked(subscriber, err)
				break
			}
		}
	}
}

func (a *depDispatcher) newDepSubscriberLocked(caseName string, theCase cases.E2ECase) *depSubscriber {
	if existed, ok := a.depSubscribers[caseName]; ok {
		a.lgr.Info("case %s already registered, returning existing result channel", caseName)
		return existed
	}
	subscriber := &depSubscriber{
		caseName: caseName,
		theCase:  theCase,
		pending:  map[string]DependencyRequirement{},
		resultCh: make(chan error, 1),
	}
	a.depSubscribers[caseName] = subscriber
	return subscriber
}

func (a *depDispatcher) tryAssignLocked(subscriber *depSubscriber) bool {
	if subscriber.closed {
		return false
	}

	for reqName, req := range subscriber.pending {
		p := a.partitions[partitionKey(req.Type(), req.Dimensions())]
		entity, ok := p.pick(req)
		if !ok {
			continue
		}

		subscriber.theCase.FulfillDependencyRequirement(reqName, entity)
		if a.onAssigned != nil {
			a.onAssigned(subscriber.caseName, reqName, entity)
		}
		delete(subscriber.pending, reqName)
	}

	return len(subscriber.pending) == 0
}

func (a *depDispatcher) markCaseReadyLocked(subscriber *depSubscriber, err error) {
	if subscriber.closed {
		return
	}
	a.lgr.Info("case %s is ready to run, closing result channel, err=%s", subscriber.caseName, err)
	subscriber.closed = true
	subscriber.resultCh <- err
	close(subscriber.resultCh)
}

func (a *depDispatcher) getOrCreatePartitionLocked(depType string, dimensions map[string]string) *depPartition {
	key := partitionKey(depType, dimensions)
	if partition, ok := a.partitions[key]; ok {
		return partition
	}
	partition := &depPartition{
		depType:    depType,
		dimensions: cloneDimensions(dimensions),
	}
	a.partitions[key] = partition
	return partition
}

func (partition *depPartition) pick(req DependencyRequirement) (DependencyEntity, bool) {
	if partition == nil {
		return nil, false
	}

	if req.IsDedicated() {
		if partition.nextDedicatedIdx >= len(partition.dedicatedEntities) {
			return nil, false
		}
		entity := partition.dedicatedEntities[partition.nextDedicatedIdx]
		partition.nextDedicatedIdx++
		return entity, true
	}

	return partition.sharedEntity, partition.sharedEntity != nil
}

func partitionKey(depType string, dimensions map[string]string) string {
	return depType + ":" + dimensionsKey(dimensions)
}

func dimensionsKey(dimensions map[string]string) string {
	return misc.PrettyJSON(dimensions)
}

func cloneDimensions(dimensions map[string]string) map[string]string {
	copied := make(map[string]string, len(dimensions))
	for key, value := range dimensions {
		copied[key] = value
	}
	return copied
}