package dispatch

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

type Outcome uint8

// Outcomes make queue admission success and overload visible to callers.
const (
	Accepted Outcome = iota
	Full
	Stopping
)

type Job func(context.Context) error

type Stats struct {
	Accepted  uint64
	Rejected  uint64
	Completed uint64
	Failed    uint64
}

// Pool owns a bounded queue, worker lifecycle, and admission counters.
type Pool struct {
	mu        sync.Mutex
	accepting bool
	jobs      chan Job
	ctx       context.Context
	cancel    context.CancelFunc
	workers   sync.WaitGroup
	done      chan struct{}

	accepted  atomic.Uint64
	rejected  atomic.Uint64
	completed atomic.Uint64
	failed    atomic.Uint64
}

// New starts a fixed worker set and ties job cancellation to parent.
func New(parent context.Context, workers, capacity int) *Pool {
	if workers < 1 {
		workers = 1
	}
	if capacity < 0 {
		capacity = 0
	}
	ctx, cancel := context.WithCancel(parent)
	p := &Pool{accepting: true, jobs: make(chan Job, capacity), ctx: ctx, cancel: cancel, done: make(chan struct{})}
	p.workers.Add(workers)
	for i := 0; i < workers; i++ {
		go p.worker()
	}
	go func() { p.workers.Wait(); close(p.done) }()
	return p
}

// Submit never waits for queue capacity. The caller gets an explicit Full result.
// The caller context can cancel admission before it enters the queue.
func (p *Pool) Submit(ctx context.Context, job Job) (Outcome, error) {
	if job == nil {
		return Stopping, errors.New("nil job")
	}
	if err := ctx.Err(); err != nil {
		return Stopping, err
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if !p.accepting {
		return Stopping, nil
	}

	select {

	case p.jobs <- job:
		p.accepted.Add(1)
		return Accepted, nil
	default:
		p.rejected.Add(1)
		return Full, nil
	}
}

func (p *Pool) StopAdmission() {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.accepting {
		p.accepting = false
		close(p.jobs)
	}
}

// Shutdown stops new submissions and drains queued work until ctx expires.
// Jobs must observe their context so deadline cancellation can stop slow work.
func (p *Pool) Shutdown(ctx context.Context) error {
	p.StopAdmission()
	select {
	case <-p.done:
		p.cancel()
		return nil
	case <-ctx.Done():
		p.cancel()
		return ctx.Err()
	}
}

func (p *Pool) Stats() Stats {
	return Stats{Accepted: p.accepted.Load(), Rejected: p.rejected.Load(), Completed: p.completed.Load(), Failed: p.failed.Load()}
}

// Workers execute admitted jobs and count their final result.
func (p *Pool) worker() {
	defer p.workers.Done()
	for {
		select {
		case <-p.ctx.Done():
			return
		case job, ok := <-p.jobs:
			if !ok {
				return
			}
			if err := job(p.ctx); err != nil {
				p.failed.Add(1)
				continue
			}
			p.completed.Add(1)
		}
	}
}