// Package fanin merges multiple input channels into a single output channel.
package fanin

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

// FanInObserver defines callbacks for monitoring fan-in lifecycle events.
type FanInObserver interface {
	OnInputAdded()
	OnInputClosed()
}

// FanInConfig holds configuration options for creating a FanIn instance.
type FanInConfig struct {
	BufferSize int
	Observer   FanInObserver
}

// Option is a functional option for configuring a FanIn instance.
type Option func(*FanInConfig)

// FanIn merges multiple input channels of type T into a single output channel.
type FanIn[T any] struct {
	cfg     *FanInConfig
	mu      sync.Mutex
	inputs  []<-chan T
	running atomic.Bool
	done    chan struct{}
}

// NewFanIn creates a new FanIn instance with the provided options.
func NewFanIn[T any](opts ...Option) (*FanIn[T], error) {
	cfg := &FanInConfig{}

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

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

	return &FanIn[T]{
		cfg: cfg,
	}, nil
}

// WithBufferSize sets the buffer capacity for the merged output channel.
func WithBufferSize(size int) Option {
	return func(f *FanInConfig) {
		f.BufferSize = size
	}
}

// WithObserver attaches a FanInObserver to receive lifecycle event callbacks.
func WithObserver(observer FanInObserver) Option {
	return func(f *FanInConfig) {
		f.Observer = observer
	}
}

// Add registers an input channel to be merged into the output.
// It must be called before Run, and is safe for concurrent use.
func (f *FanIn[T]) Add(ch <-chan T) {
	if ch == nil {
		panic("fanin: cannot add nil channel")
	}

	if f.running.Load() {
		panic("fanin: Add() called after Run()")
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.inputs = append(f.inputs, ch)

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

// Run starts merging the registered input channels into a single output channel.
func (f *FanIn[T]) Run(ctx context.Context) <-chan T {
	if !f.running.CompareAndSwap(false, true) {
		panic("fanin: Run() called multiple times")
	}

	f.mu.Lock()
	inputs := append([]<-chan T(nil), f.inputs...)
	out := make(chan T, f.cfg.BufferSize)
	f.done = make(chan struct{})
	f.mu.Unlock()

	var wg sync.WaitGroup
	for _, ch := range inputs {
		wg.Go(func() {
			for {
				select {
				case v, ok := <-ch:
					if !ok {
						if f.cfg.Observer != nil {
							safeObserve(f.cfg.Observer.OnInputClosed)
						}
						return
					}
					select {
					case out <- v:
					case <-ctx.Done():
						return
					}
				case <-ctx.Done():
					return
				}
			}
		})
	}

	go func() {
		wg.Wait()
		close(out)
		close(f.done)
	}()

	return out
}

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

// Done returns a channel that is closed when all inputs are processed or the context is cancelled.
func (f *FanIn[T]) Done() <-chan struct{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.done
}