package agenttask

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

var (
	ErrAlreadyAdmitted       = errors.New("agenttask work already has an active dispatch ticket")
	ErrProviderLimitMismatch = errors.New("agenttask provider capacity changed within one scheduler revision")
)

type DispatchCandidate struct {
	ProjectID   ProjectID
	WorkspaceID WorkspaceID
	WorkUnitID  WorkUnitID
	AttemptID   AttemptID
	Target      ExecutionTarget
}

func (c DispatchCandidate) key() string {
	return fmt.Sprintf("%s\x00%s\x00%s", c.ProjectID, c.WorkUnitID, c.AttemptID)
}

type schedulerPool struct {
	limit  int
	inUse  int
	notify chan struct{}
}

type Scheduler struct {
	mu      sync.Mutex
	pools   map[string]*schedulerPool
	tickets map[string]struct{}
}

func NewScheduler() *Scheduler {
	return &Scheduler{
		pools:   make(map[string]*schedulerPool),
		tickets: make(map[string]struct{}),
	}
}

func (s *Scheduler) Acquire(
	ctx context.Context,
	candidate DispatchCandidate,
) (*DispatchLease, error) {
	if candidate.Target.Capacity <= 0 {
		return nil, fmt.Errorf("%w: capacity must be positive", ErrProviderLimitMismatch)
	}
	ticketKey := candidate.key()
	poolKey := candidate.Target.PoolKey()
	for {
		s.mu.Lock()
		if _, exists := s.tickets[ticketKey]; exists {
			s.mu.Unlock()
			return nil, ErrAlreadyAdmitted
		}
		pool := s.pools[poolKey]
		if pool == nil {
			pool = &schedulerPool{
				limit:  candidate.Target.Capacity,
				notify: make(chan struct{}),
			}
			s.pools[poolKey] = pool
		} else if pool.limit != candidate.Target.Capacity {
			s.mu.Unlock()
			return nil, ErrProviderLimitMismatch
		}
		if pool.inUse < pool.limit {
			pool.inUse++
			s.tickets[ticketKey] = struct{}{}
			s.mu.Unlock()
			return &DispatchLease{
				scheduler: s,
				poolKey:   poolKey,
				ticketKey: ticketKey,
			}, nil
		}
		notify := pool.notify
		s.mu.Unlock()
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-notify:
		}
	}
}

type DispatchLease struct {
	once      sync.Once
	scheduler *Scheduler
	poolKey   string
	ticketKey string
}

func (l *DispatchLease) Release() {
	if l == nil || l.scheduler == nil {
		return
	}
	l.once.Do(func() {
		s := l.scheduler
		s.mu.Lock()
		defer s.mu.Unlock()
		pool := s.pools[l.poolKey]
		if pool != nil && pool.inUse > 0 {
			pool.inUse--
			close(pool.notify)
			pool.notify = make(chan struct{})
		}
		delete(s.tickets, l.ticketKey)
	})
}

func (s *Scheduler) Active(poolKey string) int {
	s.mu.Lock()
	defer s.mu.Unlock()
	if pool := s.pools[poolKey]; pool != nil {
		return pool.inUse
	}
	return 0
}