package workerload

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"
)

type Phase string

const (
	PhasePreparing       Phase = "preparing"
	PhaseWaitingUpstream Phase = "waiting_upstream"
	PhaseFinalizing      Phase = "finalizing"
)

type PressureState string

const (
	PressureNormal   PressureState = "normal"
	PressureBusy     PressureState = "busy"
	PressureCritical PressureState = "critical"
)

const (
	ModeAdaptive = "adaptive"
	ModeLegacy   = "legacy"
)

var ErrReleased = errors.New("worker load lease already released")

type Config struct {
	Mode           string
	HardLimit      int
	InitialActive  int
	InitialHeavy   int
	HealthySamples int
}

type ResourceSample struct {
	MemoryCurrentBytes int64
	MemoryLimitBytes   int64
	CPUUtilization     float64
	CPUThrottled       bool
	DBConnections      int32
	DBMaxConnections   int32
	SampledAt          time.Time
}

type Snapshot struct {
	Mode                 string
	ActiveLimit          int
	HeavyLimit           int
	ClaimLimit           int
	SafeCapacity         int
	ActiveTasks          int
	PreparingTasks       int
	WaitingUpstreamTasks int
	FinalizingTasks      int
	PressureState        PressureState
	PressureReason       string
	MemoryUtilization    float64
	CPUUtilization       float64
	DBUtilization        float64
	SampledAt            time.Time
}

type Controller struct {
	mu sync.Mutex

	mode           string
	hardLimit      int
	activeLimit    int
	heavyLimit     int
	claimLimit     int
	healthySamples int
	healthyCount   int

	preparing  int
	waiting    int
	finalizing int

	last Snapshot
	wake chan struct{}
}

type Lease struct {
	controller *Controller
	phase      Phase
	released   bool
}

func New(config Config) *Controller {
	mode := strings.ToLower(strings.TrimSpace(config.Mode))
	if mode != ModeLegacy {
		mode = ModeAdaptive
	}
	if config.HardLimit < 1 {
		config.HardLimit = 1
	}
	if config.InitialActive < 1 {
		config.InitialActive = 4
	}
	if config.InitialHeavy < 1 {
		config.InitialHeavy = 1
	}
	if config.HealthySamples < 1 {
		config.HealthySamples = 3
	}
	active := min(config.InitialActive, config.HardLimit)
	heavy := min(config.InitialHeavy, active)
	if mode == ModeLegacy {
		active = config.HardLimit
		heavy = config.HardLimit
	}
	controller := &Controller{
		mode: mode, hardLimit: config.HardLimit,
		activeLimit: active, heavyLimit: heavy, claimLimit: active,
		healthySamples: config.HealthySamples,
		wake:           make(chan struct{}, 1),
	}
	controller.last = controller.snapshotLocked(ResourceSample{SampledAt: time.Now()})
	return controller
}

func (c *Controller) Observe(sample ResourceSample) Snapshot {
	c.mu.Lock()
	defer c.mu.Unlock()
	if sample.SampledAt.IsZero() {
		sample.SampledAt = time.Now()
	}
	memory := utilization(sample.MemoryCurrentBytes, sample.MemoryLimitBytes)
	database := utilization(int64(sample.DBConnections), int64(sample.DBMaxConnections))
	cpu := bounded(sample.CPUUtilization)
	state, reason := pressure(memory, cpu, database, sample.CPUThrottled)
	if c.mode == ModeLegacy {
		c.activeLimit = c.hardLimit
		c.heavyLimit = c.hardLimit
		state = PressureNormal
		reason = "legacy"
	} else {
		c.adjustLocked(state, memory, cpu, database)
	}
	c.last = c.snapshotLocked(sample)
	c.last.PressureState = state
	c.last.PressureReason = reason
	c.last.MemoryUtilization = memory
	c.last.CPUUtilization = cpu
	c.last.DBUtilization = database
	c.last.SafeCapacity = c.activeLimit
	if state == PressureCritical {
		c.last.SafeCapacity = 0
	}
	c.signalLocked()
	return c.last
}

func (c *Controller) SetClaimLimit(limit int) Snapshot {
	c.mu.Lock()
	defer c.mu.Unlock()
	if limit < 0 {
		limit = 0
	}
	c.claimLimit = min(limit, c.hardLimit)
	c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt})
	c.signalLocked()
	return c.last
}

func (c *Controller) Snapshot() Snapshot {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt})
}

func (c *Controller) TryStart() (*Lease, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	activeLimit := min(c.activeLimit, c.claimLimit)
	if activeLimit <= 0 || c.activeLocked() >= activeLimit || c.preparing+c.finalizing >= c.heavyLimit {
		return nil, false
	}
	c.preparing++
	c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt})
	return &Lease{controller: c, phase: PhasePreparing}, true
}

func (l *Lease) EnterWaiting() error {
	if l == nil || l.controller == nil {
		return nil
	}
	c := l.controller
	c.mu.Lock()
	defer c.mu.Unlock()
	if l.released {
		return ErrReleased
	}
	if l.phase == PhaseWaitingUpstream {
		return nil
	}
	c.decrementPhaseLocked(l.phase)
	c.waiting++
	l.phase = PhaseWaitingUpstream
	c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt})
	c.signalLocked()
	return nil
}

func (l *Lease) EnterFinalizing(ctx context.Context) error {
	if l == nil || l.controller == nil {
		return nil
	}
	c := l.controller
	for {
		c.mu.Lock()
		if l.released {
			c.mu.Unlock()
			return ErrReleased
		}
		if l.phase == PhaseFinalizing {
			c.mu.Unlock()
			return nil
		}
		// Even under critical pressure, let one submitted task at a time finish
		// and release its provider lease instead of deadlocking the drain path.
		heavyLimit := max(c.heavyLimit, 1)
		if c.preparing+c.finalizing < heavyLimit {
			c.decrementPhaseLocked(l.phase)
			c.finalizing++
			l.phase = PhaseFinalizing
			c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt})
			c.mu.Unlock()
			return nil
		}
		wake := c.wake
		c.mu.Unlock()
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-wake:
		}
	}
}

func (l *Lease) Release() {
	if l == nil || l.controller == nil {
		return
	}
	c := l.controller
	c.mu.Lock()
	defer c.mu.Unlock()
	if l.released {
		return
	}
	c.decrementPhaseLocked(l.phase)
	l.released = true
	c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt})
	c.signalLocked()
}

func (l *Lease) Phase() Phase {
	if l == nil || l.controller == nil {
		return ""
	}
	l.controller.mu.Lock()
	defer l.controller.mu.Unlock()
	return l.phase
}

func (c *Controller) adjustLocked(state PressureState, memory, cpu, database float64) {
	switch state {
	case PressureCritical:
		c.healthyCount = 0
		c.activeLimit = max(1, c.activeLimit/2)
		c.heavyLimit = 1
	case PressureBusy:
		c.healthyCount = 0
		step := max(1, c.activeLimit/4)
		c.activeLimit = max(1, c.activeLimit-step)
		c.heavyLimit = min(c.heavyLimit, max(1, (c.activeLimit+3)/4))
	default:
		if memory > .60 || cpu > .65 || database > .65 || c.activeLocked() < min(c.activeLimit, c.claimLimit) {
			c.healthyCount = 0
			return
		}
		c.healthyCount++
		if c.healthyCount < c.healthySamples {
			return
		}
		c.healthyCount = 0
		step := max(1, c.activeLimit/4)
		c.activeLimit = min(c.hardLimit, c.activeLimit+step)
		c.heavyLimit = min(c.activeLimit, max(1, (c.activeLimit+3)/4))
	}
}

func (c *Controller) snapshotLocked(sample ResourceSample) Snapshot {
	sampledAt := sample.SampledAt
	if sampledAt.IsZero() {
		sampledAt = c.last.SampledAt
	}
	safeCapacity := c.activeLimit
	if c.last.PressureState == PressureCritical {
		safeCapacity = 0
	}
	return Snapshot{
		Mode:                 c.mode,
		ActiveLimit:          c.activeLimit,
		HeavyLimit:           c.heavyLimit,
		ClaimLimit:           c.claimLimit,
		SafeCapacity:         safeCapacity,
		ActiveTasks:          c.activeLocked(),
		PreparingTasks:       c.preparing,
		WaitingUpstreamTasks: c.waiting,
		FinalizingTasks:      c.finalizing,
		PressureState:        c.last.PressureState,
		PressureReason:       c.last.PressureReason,
		MemoryUtilization:    c.last.MemoryUtilization,
		CPUUtilization:       c.last.CPUUtilization,
		DBUtilization:        c.last.DBUtilization,
		SampledAt:            sampledAt,
	}
}

func (c *Controller) activeLocked() int { return c.preparing + c.waiting + c.finalizing }

func (c *Controller) decrementPhaseLocked(phase Phase) {
	switch phase {
	case PhasePreparing:
		c.preparing = max(0, c.preparing-1)
	case PhaseWaitingUpstream:
		c.waiting = max(0, c.waiting-1)
	case PhaseFinalizing:
		c.finalizing = max(0, c.finalizing-1)
	}
}

func (c *Controller) signalLocked() {
	select {
	case c.wake <- struct{}{}:
	default:
	}
}

func pressure(memory, cpu, database float64, throttled bool) (PressureState, string) {
	switch {
	case memory >= .90:
		return PressureCritical, "memory"
	case database >= .90:
		return PressureCritical, "database"
	case cpu >= .95 && throttled:
		return PressureCritical, "cpu_throttled"
	case memory >= .75:
		return PressureBusy, "memory"
	case database >= .80:
		return PressureBusy, "database"
	case cpu >= .80 || throttled:
		return PressureBusy, "cpu"
	default:
		return PressureNormal, ""
	}
}

func utilization(current, limit int64) float64 {
	if current <= 0 || limit <= 0 {
		return 0
	}
	return bounded(float64(current) / float64(limit))
}

func bounded(value float64) float64 {
	if value < 0 {
		return 0
	}
	if value > 1 {
		return 1
	}
	return value
}