package syncz

import (
	"context"
	"fmt"
	"log/slog"

	"google.golang.org/protobuf/proto"
)

type WorkSource[ID comparable, C any] struct {
	ID            ID
	Configuration C
}

type WorkerFactory[ID comparable, C any] interface {
	New(WorkSource[ID, C]) Worker
}

type WorkerManager[ID comparable, C any] struct {
	log           *slog.Logger
	fld           func(ID) slog.Attr
	workerFactory WorkerFactory[ID, C]
	equal         func(c1, c2 C) bool
	workers       map[ID]*workerHolder[ID, C] // source id -> worker holder instance
}

func NewWorkerManager[ID comparable, C any](log *slog.Logger, fld func(ID) slog.Attr, workerFactory WorkerFactory[ID, C], equal func(c1, c2 C) bool) *WorkerManager[ID, C] {
	return &WorkerManager[ID, C]{
		log:           log,
		fld:           fld,
		workerFactory: workerFactory,
		equal:         equal,
		workers:       map[ID]*workerHolder[ID, C]{},
	}
}

func NewProtoWorkerManager[ID comparable, C proto.Message](log *slog.Logger, fld func(ID) slog.Attr, workerFactory WorkerFactory[ID, C]) *WorkerManager[ID, C] {
	return NewWorkerManager(log, fld, workerFactory, func(c1, c2 C) bool {
		return proto.Equal(c1, c2)
	})
}

func (m *WorkerManager[ID, C]) startNewWorker(source WorkSource[ID, C]) {
	id := source.ID
	oldHolder := m.workers[id]
	// Synchronously call from the original goroutine to avoid calling the factory concurrently.
	worker := m.workerFactory.New(source)
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	holder := &workerHolder[ID, C]{
		src:  source,
		done: done,
		stop: cancel,
	}
	go func() {
		defer close(done)
		if oldHolder != nil {
			m.log.Info("Stopping worker", m.fld(id))
			oldHolder.stop()
			m.log.Info("Waiting for worker to stop", m.fld(id))
			<-oldHolder.done
			m.log.Info("Worker stopped", m.fld(id))
			select {
			case <-ctx.Done():
				// Context is done, no need to start the new worker.
				return
			default:
			}
		}
		m.log.Info("Starting worker", m.fld(id))
		worker.Run(ctx)
	}()
	m.workers[id] = holder
}

func (m *WorkerManager[ID, C]) ApplyConfiguration(sources []WorkSource[ID, C]) error {
	newSetOfSources := make(map[ID]struct{}, len(sources))
	var sourcesToStartWorkersFor []WorkSource[ID, C]

	// Collect sources without workers or with updated configuration.
	for _, source := range sources {
		id := source.ID
		if _, ok := newSetOfSources[id]; ok {
			return fmt.Errorf("duplicate source id: %v", id)
		}
		newSetOfSources[id] = struct{}{}
		holder := m.workers[id]
		if holder == nil { // New source added
			sourcesToStartWorkersFor = append(sourcesToStartWorkersFor, source)
		} else { // We have a worker for this source already
			if m.equal(source.Configuration, holder.src.Configuration) {
				// Worker's configuration hasn't changed, nothing to do here
				continue
			}
			m.log.Info("Configuration changed, restarting worker", m.fld(id))
			sourcesToStartWorkersFor = append(sourcesToStartWorkersFor, source)
		}
	}

	// Stop workers for sources which have been removed from the list.
	// Also, remove old workers that have stopped.
	for sourceID, holder := range m.workers {
		select {
		case <-holder.done:
			delete(m.workers, sourceID)
			continue
		default:
		}
		if _, ok := newSetOfSources[sourceID]; ok {
			continue
		}
		m.log.Info("Stopping worker", m.fld(sourceID))
		holder.stop()
		// Don't immediately delete stopped workers so that Wait() can wait for them to stop.
	}

	// Start new workers for new sources or because of updated configuration.
	for _, source := range sourcesToStartWorkersFor {
		m.startNewWorker(source)
	}
	return nil
}

// Stop tells all workers to stop. It does not wait for workers to stop. Use Wait() for that.
func (m *WorkerManager[ID, C]) Stop() {
	for sourceID, holder := range m.workers {
		m.log.Info("Stopping worker", m.fld(sourceID))
		holder.stop()
	}
}

// Wait waits for all workers to stop.
func (m *WorkerManager[ID, C]) Wait() {
	for sourceID, holder := range m.workers {
		m.log.Info("Waiting for worker to stop", m.fld(sourceID))
		<-holder.done
		m.log.Info("Worker stopped", m.fld(sourceID))
	}
	// Clear the workers map to allow for a proper restart of the workers.
	clear(m.workers)
}

func (m *WorkerManager[ID, C]) StopAndWait() {
	m.Stop()
	m.Wait()
}

type workerHolder[ID comparable, C any] struct {
	src  WorkSource[ID, C]
	done <-chan struct{}
	stop context.CancelFunc
}
