package service

import (
	"context"
	"fmt"
	"sort"
	"time"
)

// normalizeGroupPolicy applies queue defaults while preserving an explicit
// zero queue timeout as "no timeout".
func normalizeGroupPolicy(policy groupPolicy) groupPolicy {
	if policy.maxQueue <= 0 {
		policy.maxQueue = defaultGroupMaxQueue
	}
	if policy.queueTimeout > 0 {
		policy.queueTimeoutSet = true
	}
	if !policy.queueTimeoutSet {
		policy.queueTimeout = defaultQueueTimeout
		policy.queueTimeoutSet = true
	}
	return policy
}

// getOrCreateGroupLocked returns an existing group or creates one with the
// given policy (applying defaults for unset fields). Policy is only applied
// at creation; existing groups retain their current policy until an explicit
// updateGroupPolicyLocked call (used for config-refresh policy propagation).
func (m *modelQueueManager) getOrCreateGroupLocked(key string, policy groupPolicy) *modelQueueGroup {
	g, ok := m.groups[key]
	if !ok {
		policy = normalizeGroupPolicy(policy)
		g = &modelQueueGroup{
			key:          key,
			policy:       policy,
			inflight:     make(map[string]int),
			longInflight: make(map[string]int),
		}
		m.groups[key] = g
	}
	return g
}

// updateGroupPolicyLocked refreshes an existing group's queue policy so that
// provider max_queue / queue_timeout_ms changes applied via config refresh take
// effect on subsequent admissions. Only configured policy fields override the
// current values; in-flight counts and queued items are preserved untouched.
func (m *modelQueueManager) updateGroupPolicyLocked(g *modelQueueGroup, policy groupPolicy) {
	if g == nil {
		return
	}
	if policy.maxQueue > 0 {
		g.policy.maxQueue = policy.maxQueue
	}
	if policy.queueTimeout > 0 {
		g.policy.queueTimeout = policy.queueTimeout
		g.policy.queueTimeoutSet = true
	} else if policy.queueTimeoutSet {
		g.policy.queueTimeout = policy.queueTimeout
		g.policy.queueTimeoutSet = true
	}
}

// findAvailableNodeLocked returns the available candidate with the lowest
// absolute in_flight count. Full candidates are skipped, so a provider that is
// at capacity is deferred behind providers that still have room. When in_flight
// is equal, the candidate with the lower priority value wins. When both
// in_flight and priority are equal, all tied candidates remain in best[] so
// that the rotation logic in the caller picks the next slot after the last
// selected one.
func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode, long bool) *candidateNode {
	var best []candidateNode
	bestInflight := 0
	bestPriority := 0
	for i := range candidates {
		c := candidates[i]
		if c.capacity <= 0 {
			continue
		}
		slot := c.slotKey()
		inflight := group.inflight[slot]
		if inflight >= c.capacity {
			continue
		}
		// Long-context gate: a long request is excluded from a provider whose
		// long slots are full. Providers with long_context_capacity == 0 declare
		// no dedicated long-slot limit and are not gated here. Normal requests
		// (long == false) never consult long slot state.
		if long && c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity {
			continue
		}
		switch {
		case len(best) == 0 || inflight < bestInflight:
			best = []candidateNode{c}
			bestInflight = inflight
			bestPriority = c.priority
		case inflight == bestInflight:
			if c.priority < bestPriority {
				best = []candidateNode{c}
				bestPriority = c.priority
			} else if c.priority == bestPriority {
				best = append(best, c)
			}
		}
	}
	if len(best) == 0 {
		return nil
	}
	sort.Slice(best, func(i, j int) bool {
		return candidateLess(&best[i], &best[j])
	})
	idx := 0
	if group.lastSelectedSlot != "" {
		for i := range best {
			if best[i].slotKey() == group.lastSelectedSlot {
				idx = (i + 1) % len(best)
				break
			}
		}
	}
	return &best[idx]
}

// candidateLess provides a deterministic ordering for equal-inflight/priority rotation:
// providerID first, then nodeID.
func candidateLess(a, b *candidateNode) bool {
	if b == nil {
		return true
	}
	if a.providerID != b.providerID {
		return a.providerID < b.providerID
	}
	return a.entry.NodeID < b.entry.NodeID
}

// admit selects an available candidate for the given model group, or queues
// the request until a slot opens. Blocks until a candidate is assigned, the
// queue timeout expires, or ctx is cancelled. Returns the selected candidateNode
// so callers can rewrite the dispatch target for provider-pool requests.
// The optional long flag marks the request as long-context so it is gated on
// provider long slots and reserves a long slot on dispatch.
func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, longFlag ...bool) (*candidateNode, error) {
	candidate, _, err := m.admitWithReason(ctx, groupKey, adapter, target, candidates, policy, longFlag...)
	return candidate, err
}

func (m *modelQueueManager) admitWithReason(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, longFlag ...bool) (*candidateNode, string, error) {
	long := len(longFlag) > 0 && longFlag[0]
	m.mu.Lock()

	group := m.getOrCreateGroupLocked(groupKey, policy)
	// Refresh the policy on an already-existing group so capacity/queue changes
	// applied via config refresh (and recomputed per-request from the live store)
	// take effect on this and subsequent admissions.
	m.updateGroupPolicyLocked(group, policy)
	if group.adapter == "" {
		group.adapter = adapter
	}
	if group.target == "" {
		group.target = target
	}

	candidate := m.findAvailableNodeLocked(group, candidates, long)
	if candidate != nil {
		slot := candidate.slotKey()
		group.inflight[slot]++
		if long && candidate.longContextCapacity > 0 {
			group.longInflight[slot]++
		}
		group.lastSelectedSlot = slot
		m.mu.Unlock()
		return candidate, "dispatched", nil
	}

	if len(group.queue) >= group.policy.maxQueue {
		m.mu.Unlock()
		return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueFull)
	}

	// Determine queue reason.
	reason := "capacity_full"
	if long {
		longBlocked := false
		for _, c := range candidates {
			if c.capacity <= 0 {
				continue
			}
			slot := c.slotKey()
			inflight := group.inflight[slot]
			if inflight < c.capacity {
				if c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity {
					longBlocked = true
					break
				}
			}
		}
		if longBlocked {
			reason = "long_context_capacity_full"
		}
	}

	var deadline time.Time
	if group.policy.queueTimeout > 0 {
		deadline = time.Now().Add(group.policy.queueTimeout)
	}
	item := &queueItem{
		candidates: candidates,
		waitCh:     make(chan admitResult, 1),
		deadline:   deadline,
		long:       long,
		reason:     reason,
	}
	group.queue = append(group.queue, item)
	timeout := group.policy.queueTimeout
	m.mu.Unlock()

	if timeout <= 0 {
		select {
		case res := <-item.waitCh:
			return res.candidate, item.reason, res.err
		case <-ctx.Done():
			m.mu.Lock()
			m.removeItemLocked(groupKey, item)
			m.mu.Unlock()
			// Handle race: dispatch may have sent to waitCh just before cancellation.
			select {
			case res := <-item.waitCh:
				if res.candidate != nil {
					slot := res.candidate.slotKey()
					var nodeID, providerID string
					if colonIdx := findLastColon(slot); colonIdx > 0 {
						nodeID = slot[:colonIdx]
						providerID = slot[colonIdx+1:]
					} else {
						nodeID = slot
					}
					longSlot := item.long && res.candidate.longContextCapacity > 0
					m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
				}
			default:
			}
			return nil, "", ctx.Err()
		}
	}

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	select {
	case res := <-item.waitCh:
		return res.candidate, item.reason, res.err
	case <-timer.C:
		m.mu.Lock()
		m.removeItemLocked(groupKey, item)
		m.mu.Unlock()
		// Handle race: dispatch may have sent to waitCh just before timeout fired.
		select {
		case res := <-item.waitCh:
			if res.candidate != nil {
				// Use public releaseSlot (which acquires m.mu) because we are
				// outside the lock here.  This avoids shared-state mutation
				// without mutex that the original releaseSlotByCandidate call
				// would cause.
				slot := res.candidate.slotKey()
				var nodeID, providerID string
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					nodeID = slot[:colonIdx]
					providerID = slot[colonIdx+1:]
				} else {
					nodeID = slot
				}
				longSlot := item.long && res.candidate.longContextCapacity > 0
				m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
			}
		default:
		}
		return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueTimeout)
	case <-ctx.Done():
		m.mu.Lock()
		m.removeItemLocked(groupKey, item)
		m.mu.Unlock()
		// Handle race: dispatch may have sent to waitCh just before cancellation.
		select {
		case res := <-item.waitCh:
			if res.candidate != nil {
				// Same fix as timeout path: use public releaseSlot with
				// explicit mutex to avoid concurrent map/slice mutation.
				slot := res.candidate.slotKey()
				var nodeID, providerID string
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					nodeID = slot[:colonIdx]
					providerID = slot[colonIdx+1:]
				} else {
					nodeID = slot
				}
				longSlot := item.long && res.candidate.longContextCapacity > 0
				m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
			}
		default:
		}
		return nil, "", ctx.Err()
	}
}

// trackInflight records a dispatched run so release events can find it.
// providerID is non-empty for provider-pool dispatches and is used to
// attribute in-flight counts to the correct provider snapshot. The optional
// longSlot flag records that a long-context slot was reserved so the terminal
// release also frees the long slot.
func (m *modelQueueManager) trackInflight(groupKey, runID, nodeID, providerID string, longSlot ...bool) {
	long := len(longSlot) > 0 && longSlot[0]
	m.mu.Lock()
	m.inflightByRun[runID] = inflightRec{groupKey: groupKey, nodeID: nodeID, providerID: providerID, long: long}
	m.mu.Unlock()
}

// tryDispatchLocked scans the queue from the front and dispatches every item
// that currently has an available candidate. Unlike a head-only pump, a blocked
// head does not stop dispatchable items behind it: when the head is a long
// request whose provider long slots are full, later normal requests that fit
// ordinary capacity are still dispatched, avoiding head-of-line blocking (SDD
// D05/D06). FIFO order is preserved among items with the same dispatchability
// because the scan always picks the earliest dispatchable item first. Expired
// items are timed out in place regardless of position. Must be called with
// m.mu held.
func (m *modelQueueManager) tryDispatchLocked(group *modelQueueGroup) {
	for {
		progressed := false
		for i := 0; i < len(group.queue); i++ {
			item := group.queue[i]

			if !item.deadline.IsZero() && time.Now().After(item.deadline) {
				group.queue = append(group.queue[:i], group.queue[i+1:]...)
				select {
				case item.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", group.key, errQueueTimeout)}:
				default:
				}
				progressed = true
				break
			}

			candidate := m.findAvailableNodeLocked(group, item.candidates, item.long)
			if candidate == nil {
				// No slot for this item under current provider state; skip it so a
				// later dispatchable item is not blocked behind it.
				continue
			}

			group.queue = append(group.queue[:i], group.queue[i+1:]...)
			slot := candidate.slotKey()
			reserveLong := item.long && candidate.longContextCapacity > 0
			group.inflight[slot]++
			if reserveLong {
				group.longInflight[slot]++
			}

			select {
			case item.waitCh <- admitResult{candidate: candidate}:
				group.lastSelectedSlot = slot
			default:
				// Caller already timed out or cancelled; release the reserved slot.
				group.inflight[slot]--
				if reserveLong {
					group.longInflight[slot]--
				}
			}
			progressed = true
			break
		}
		if !progressed {
			break
		}
	}
}

func (m *modelQueueManager) removeItemLocked(groupKey string, item *queueItem) {
	group, ok := m.groups[groupKey]
	if !ok {
		return
	}
	for i, it := range group.queue {
		if it == item {
			group.queue = append(group.queue[:i], group.queue[i+1:]...)
			return
		}
	}
}