// Package fanout distributes values from a single input channel to multiple output channels.
package fanout

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
)

// FanOutObserver defines callbacks for monitoring fan-out lifecycle events.
type FanOutObserver interface {
	OnDistribute()
	OnOutputClosed()
}

// FanOutStrategy defines how values are distributed across output channels.
type FanOutStrategy int

const (
	// RoundRobin distributes values evenly across output channels in rotation.
	RoundRobin FanOutStrategy = iota
	// Broadcast sends each value to all output channels.
	Broadcast
)

// ErrInvalidStrategy is returned when an unknown FanOutStrategy is provided.
var ErrInvalidStrategy = errors.New("fanout: unknown distribution strategy")

// FanOutConfig holds configuration options for creating a FanOut instance.
type FanOutConfig struct {
	WorkerCount int
	BufferSize  int
	Strategy    FanOutStrategy
	Observer    FanOutObserver
}

// Option is a functional option for configuring a FanOut instance.
type Option func(*FanOutConfig)

// FanOut distributes values from a single input channel to multiple output channels.
type FanOut[T any] struct {
	cfg           *FanOutConfig
	outputs       []chan T
	running       atomic.Bool
	cachedOutputs []<-chan T
	once          sync.Once
}

// NewFanOut creates a new FanOut instance with the provided options.
func NewFanOut[T any](opts ...Option) (*FanOut[T], error) {
	cfg := &FanOutConfig{
		WorkerCount: 1,
		BufferSize:  0,
		Observer:    nil,
		Strategy:    RoundRobin,
	}

	for _, opt := range opts {
		opt(cfg)
	}

	if cfg.WorkerCount < 1 {
		return nil, fmt.Errorf("fanout: worker count must be at least 1, got %d", cfg.WorkerCount)
	}

	if cfg.BufferSize < 0 {
		return nil, fmt.Errorf("fanout: buffer size must not be negative, got %d", cfg.BufferSize)
	}

	if cfg.Strategy != RoundRobin && cfg.Strategy != Broadcast {
		return nil, ErrInvalidStrategy
	}

	outputs := make([]chan T, cfg.WorkerCount)
	for i := range outputs {
		outputs[i] = make(chan T, cfg.BufferSize)
	}

	return &FanOut[T]{
		cfg:     cfg,
		outputs: outputs,
	}, nil
}

// Run starts distributing values from the input channel to all output channels.
func (f *FanOut[T]) Run(ctx context.Context, input <-chan T) {
	if input == nil {
		panic("fanout: input channel must not be nil")
	}

	if !f.running.CompareAndSwap(false, true) {
		panic("fanout: Run() called multiple times")
	}

	go func() {
		idx := 0
		defer f.running.Store(false)
		defer f.closeAllOutput()
		for {
			select {
			case value, ok := <-input:
				if !ok {
					return
				}
				switch f.cfg.Strategy {
				case RoundRobin:
					f.roundRobin(ctx, value, &idx)
				case Broadcast:
					f.broadcast(ctx, value)
				}
			case <-ctx.Done():
				return
			}
		}
	}()
}

// Outputs returns receive-only channel references for the distributed streams.
func (f *FanOut[T]) Outputs() []<-chan T {
	f.once.Do(func() {
		f.cachedOutputs = make([]<-chan T, len(f.outputs))
		for i, ch := range f.outputs {
			f.cachedOutputs[i] = ch
		}
	})

	return f.cachedOutputs
}

func safeObserve(fn func()) {
	defer func() { recover() }()
	fn()
}

func (f *FanOut[T]) roundRobin(ctx context.Context, value T, index *int) {
	target := *index % f.cfg.WorkerCount
	select {
	case f.outputs[target] <- value:
		*index = (*index + 1) % f.cfg.WorkerCount
		if f.cfg.Observer != nil {
			safeObserve(f.cfg.Observer.OnDistribute)
		}
	case <-ctx.Done():
		return
	}
}

func (f *FanOut[T]) broadcast(ctx context.Context, value T) {
	var wg sync.WaitGroup
	for _, ch := range f.outputs {
		wg.Go(func() {
			select {
			case ch <- value:
				if f.cfg.Observer != nil {
					safeObserve(f.cfg.Observer.OnDistribute)
				}
			case <-ctx.Done():
				return
			}
		})
	}
	wg.Wait()
}

func (f *FanOut[T]) closeAllOutput() {
	for _, ch := range f.outputs {
		close(ch)
	}

	if f.cfg.Observer != nil {
		safeObserve(f.cfg.Observer.OnOutputClosed)
	}
}

// WithWorkerCount sets the number of output channels for the fan-out.
func WithWorkerCount(n int) Option {
	return func(cfg *FanOutConfig) {
		cfg.WorkerCount = n
	}
}

// WithBufferSize sets the buffer capacity for each output channel.
func WithBufferSize(size int) Option {
	return func(cfg *FanOutConfig) {
		cfg.BufferSize = size
	}
}

// WithStrategy sets the distribution strategy for the fan-out.
func WithStrategy(strategy FanOutStrategy) Option {
	return func(cfg *FanOutConfig) {
		cfg.Strategy = strategy
	}
}

// WithObserver attaches a FanOutObserver to receive lifecycle event callbacks.
func WithObserver(observer FanOutObserver) Option {
	return func(cfg *FanOutConfig) {
		cfg.Observer = observer
	}
}