package service

import (
	"fmt"
	"sync"
	"time"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
)

const (
	defaultNodeCapacity  = 1
	defaultGroupMaxQueue = 16
	defaultQueueTimeout  = 30 * time.Second
)

var (
	errQueueFull    = fmt.Errorf("queue is full")
	errQueueTimeout = fmt.Errorf("queue timeout")
)

// errProviderUnavailable is the terminal sentinel delivered to a queued
// waiter when the live resolver returns zero candidates — no connected,
// enabled provider can fulfill the request at all. Capacity or long-context
// saturation does not trigger this; that is a transient block, not a
// terminal absence.
var errProviderUnavailable = fmt.Errorf("provider unavailable")

// resolveOutcome classifies the result of re-resolving a queued candidate
// universe at dispatch time.
type resolveOutcome uint8

const (
	// resolveOk means the resolver returned usable candidates. The pump
	// continues to selection (findAvailableNodeLocked), which may still
	// skip the item if all candidates are at capacity — that is the
	// "temporarily blocked" path, not terminal.
	resolveOk resolveOutcome = iota
	// resolveNoCandidates means the live candidate universe is empty:
	// the resolver found no eligible providers. The item is removed from
	// the queue and the waiter receives errProviderUnavailable.
	resolveNoCandidates
	// resolveResolverError means the live resolver returned a non-empty
	// error — a configuration, programming, or other failure distinct from
	// errProviderUnavailable. The item remains queued for a later live-state
	// re-evaluation and no terminal result is delivered.
	resolveResolverError
	// resolveTerminalError preserves a typed, request-specific admission error
	// and removes the waiter without reserving or dispatching a provider slot.
	resolveTerminalError
)

// candidateNode pairs a registry entry with the per-request capacity derived
// from the node's adapter config (or Runtime.Concurrency as fallback).
// For provider-pool candidates, providerID, adapter, and servedTarget carry
// the globally-unique provider id, the dispatch adapter key, and the concrete
// model name to dispatch. Priority is the provider's dispatch priority (lower
// is preferred) and is used as a tie-break when in_flight counts are equal.
// providerExecutionPath identifies how a provider executes: tunnel/passthrough
// for OpenAI-compatible callers, or normalized for Ollama/CLI/native providers.
type providerExecutionPath string

const (
	// providerExecutionPathTunnel marks OpenAI-compatible (passthrough) execution.
	providerExecutionPathTunnel providerExecutionPath = "provider_tunnel"
	// providerExecutionPathNormalized marks Ollama/CLI/native (normalized) execution.
	providerExecutionPathNormalized providerExecutionPath = "normalized"
)

type candidateNode struct {
	entry        *edgenode.NodeEntry
	capacity     int
	priority     int
	providerID   string // non-empty for provider-pool candidates
	providerType string // non-empty for provider-pool candidates; the raw provider type (e.g. "vllm", "ollama")
	// lifecycleCapabilities are the provider-advertised capabilities used by
	// request-local admission policy. They are copied from the Edge-owned
	// provider catalog for every initial and re-resolved candidate.
	lifecycleCapabilities []string
	adapter               string // non-empty for provider-pool candidates; dispatch adapter key
	servedTarget          string // concrete served model name; used for target rewrite
	// profile is the resolved concrete protocol profile snapshot for this
	// candidate. It is carried through to the tunnel wire request so the Node
	// adapter can resolve operation URLs from the profile's operation paths.
	// Nil for legacy candidates without a profile.
	profile *config.ConcreteProtocolProfile
	// longContextCapacity is the provider's concurrent long-context slot limit.
	// Zero means the provider declares no dedicated long-slot limit and long
	// requests are not gated on it.
	longContextCapacity int
	// executionPath classifies whether the provider executes via tunnel/passthrough
	// (OpenAI-compatible) or normalized (Ollama/CLI/native). Used by downstream
	// dispatch logic to decide execution without re-evaluating the provider type.
	executionPath providerExecutionPath
	leaseID       uint64
	// generation is the registry-assigned connection generation of the node
	// entry this candidate was resolved from. It fences a stale candidate — one
	// resolved before the owning connection disconnected or was superseded by a
	// reconnect — out of reserve and dispatch handoff. Zero means untracked
	// (legacy/direct candidates and hand-built fixtures) and is never fenced.
	generation             uint64
	responseStallTimeoutMS int64
}

// slotKey returns a unique slot key for inflight accounting.
// For provider-pool candidates (providerID non-empty) it uses "nodeID:providerID"
// so that multiple providers on the same node are tracked independently.
// For legacy candidates it falls back to nodeID only.
func (c *candidateNode) slotKey() string {
	if c.providerID != "" {
		return c.entry.NodeID + ":" + c.providerID
	}
	return c.entry.NodeID
}

type groupPolicy struct {
	maxQueue        int
	queueTimeout    time.Duration
	queueTimeoutSet bool
}

// GroupPolicy is the exported view of groupPolicy used by bootstrap/runtime
// to construct a provider-pool policy from config. It is converted into the
// internal groupPolicy at the boundary.
type GroupPolicy = groupPolicy

// NewGroupPolicy constructs a GroupPolicy with the given max_queue and
// queue_timeout. maxQueue <= 0 normalizes to the canonical default.
func NewGroupPolicy(maxQueue int, queueTimeout time.Duration) GroupPolicy {
	if maxQueue <= 0 {
		maxQueue = defaultGroupMaxQueue
	}
	return GroupPolicy{
		maxQueue:        maxQueue,
		queueTimeout:    queueTimeout,
		queueTimeoutSet: true,
	}
}

// GroupPolicySnapshot is an exported snapshot of the root provider-pool
// queue policy. Callers that need to observe max_queue, queue_timeout, and
// the explicit-zero flag without reaching into the service's internal state
// use this type.
type GroupPolicySnapshot struct {
	MaxQueue        int
	QueueTimeout    time.Duration
	QueueTimeoutSet bool
}

// Snapshot returns an exported snapshot of the policy. Safe to call from any
// package for test observability or status reporting.
func (p GroupPolicy) Snapshot() GroupPolicySnapshot {
	return GroupPolicySnapshot{
		MaxQueue:        p.maxQueue,
		QueueTimeout:    p.queueTimeout,
		QueueTimeoutSet: p.queueTimeoutSet,
	}
}

// leaseState tracks where a provider lease is in its lifecycle. A lease moves
// reserved → tracked → released, or straight from reserved to released when a
// dispatch fails before a run id exists. Only the transition into released frees
// provider resources, so every terminal cause converges on exactly one decrement.
type leaseState uint8

const (
	leaseStateReserved leaseState = iota
	leaseStateTracked
	leaseStateReleased
)

// providerLease is the identity of one admitted slot. The lease id is unique for
// the manager's lifetime, so a late release racing a re-admission on the same
// (node, provider) slot can never free the newer request's resources: the stale
// id is simply no longer in the lease map.
type providerLease struct {
	id         uint64
	groupKey   string
	nodeID     string
	providerID string // non-empty for provider-pool dispatches
	generation uint64 // registry connection generation this lease was admitted under
	adapter    string // immutable adapter binding selected for this attempt
	target     string // immutable concrete target binding selected for this attempt
	long       bool   // true when a long-context slot was reserved for this lease
	state      leaseState
	runID      string
}

// providerRuntimeHealthKey scopes runtime health to one provider on one Node
// connection. A reconnect receives a new generation and therefore never
// inherits health evidence observed on the superseded connection.
type providerRuntimeHealthKey struct {
	nodeID     string
	generation uint64
	providerID string
}

// providerRuntimeHealthOverlay is deliberately separate from the config-owned
// provider resource. observationSeq is the high-water mark for every validated
// bound observation, while unavailable changes effective admission/snapshot
// health only. adapter and target retain the exact binding that lowered the
// provider so only the same exact-target status probe may recover it.
type providerRuntimeHealthOverlay struct {
	adapter        string
	target         string
	observationSeq uint64
	unavailable    bool
}

type admitResult struct {
	candidate *candidateNode
	err       error
}

// recoveryCandidatePolicy carries the request-local avoided-provider recovery
// hint through immediate admission and every queued re-resolution. It is a pure
// value with no persistence beyond the request/queue item lifetime: the queue
// stores it on the pending item only so the pump reapplies the identical
// preference the caller submitted.
//
// The zero value (empty avoidProviderID, false allowAvoidedProviderFallback)
// disables recovery entirely, so every non-recovery admission path preserves the
// current candidate-selection behavior.
type recoveryCandidatePolicy struct {
	// avoidProviderID, when non-empty, marks the provider the caller wants to
	// avoid. A runtime-eligible alternate is always preferred over it.
	avoidProviderID string
	// allowAvoidedProviderFallback permits re-selecting the avoided provider,
	// but only when no runtime-eligible alternate remains. It is the caller's
	// explicit, probe-backed permission and never derived from overlay state.
	allowAvoidedProviderFallback bool
}

// active reports whether the policy expresses an avoided-provider preference.
func (p recoveryCandidatePolicy) active() bool {
	return p.avoidProviderID != ""
}

// queueItem is one pending admission. candidates carry the request's resource
// identity (node, provider, served target) only: capacity, long-context capacity,
// priority, and the enabled switch are re-read from live state at dispatch time,
// never from the snapshot taken here at enqueue time.
type queueItem struct {
	candidates        []candidateNode
	resolveCandidates func() ([]candidateNode, error)
	waitCh            chan admitResult
	enqueuedAt        time.Time     // zero means fixture-only; used for deadline rebase against current policy timeout
	policyChange      chan struct{} // non-blocking signal from setProviderPoolPolicyLocked to restart the waiter timer
	deadline          time.Time     // current timeout; zero means no queue deadline (waiter waits for ctx cancellation only)
	long              bool          // true when the queued request is a long-context request
	providerPool      bool          // true when this item is enqueued under the provider-pool policy scope
	reason            string
	enqueueSeq        uint64
	// recovery is the request-local avoided-provider preference reapplied on
	// every pump re-resolution. Zero value for non-recovery admissions.
	recovery recoveryCandidatePolicy
}

type modelQueueGroup struct {
	key              string
	policy           groupPolicy
	queue            []*queueItem
	inflight         map[string]int // slotKey → current in-flight count
	longInflight     map[string]int // slotKey → current long-context in-flight count
	lastSelectedSlot string
	adapter          string
	target           string
}

type modelQueueManager struct {
	mu sync.Mutex
	// healthObserver receives immutable post-decision projections only. It is
	// never called while mu is held.
	healthObserver providerHealthObserver
	groups         map[string]*modelQueueGroup
	// leases holds every live lease by id. Admission inserts under the same
	// critical section that reserves the resource, and release deletes under the
	// same critical section that frees it, so the lease map is the single source
	// of truth for "this slot is still held".
	leases map[uint64]*providerLease
	// leaseByRun indexes tracked leases so run terminal events and node
	// disconnects can resolve the lease identity they own.
	leaseByRun map[string]uint64
	leaseSeq   uint64
	// enqueueSeq allocates the global arrival order for pending items. Every
	// queued item across every model group draws from this one counter, so the
	// scheduler can compare waiters that never share a group queue.
	enqueueSeq uint64
	store      *edgenode.NodeStore
	resources  map[providerResourceKey]*providerResourceState
	// runtimeHealth is a generation-scoped overlay. It never mutates NodeStore
	// provider config and is guarded by the same lock as leases/resources so
	// health transitions and admission observe one linearized state.
	runtimeHealth map[providerRuntimeHealthKey]*providerRuntimeHealthOverlay
	// providerPoolPolicy is the canonical root policy shared by every
	// provider-pool admission. It replaces the legacy per-provider first-encounter
	// heuristic and makes max_queue a hard cap across all model groups for the
	// whole pool. Legacy adapter (non-provider-pool) admissions never touch this
	// value so their group-local policy stays intact.
	providerPoolPolicy groupPolicy
}

func newModelQueueManager(store *edgenode.NodeStore) *modelQueueManager {
	return &modelQueueManager{
		healthObserver: defaultHealthObserver(),
		groups:         make(map[string]*modelQueueGroup),
		leases:         make(map[uint64]*providerLease),
		leaseByRun:     make(map[string]uint64),
		store:          store,
		resources:      make(map[providerResourceKey]*providerResourceState),
		runtimeHealth:  make(map[providerRuntimeHealthKey]*providerRuntimeHealthOverlay),
	}
}

// setStore swaps the node catalog the scheduler resolves live provider state
// from, reconciles the per-provider resource states against it, and re-evaluates
// every pending item. The pump matters as much as the reconcile: a capacity
// increase or a re-enabled provider creates room that no lease release will ever
// announce, so without pumping here the waiters that the refresh just unblocked
// would sleep until their own group happens to see an event, or time out.
func (m *modelQueueManager) setStore(store *edgenode.NodeStore) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.store = store

	if store == nil {
		m.pumpAllLocked()
		return
	}

	for _, res := range m.resources {
		res.orphan = true
	}

	for _, rec := range store.All() {
		for _, prov := range rec.Providers {
			if prov.ID == "" {
				continue
			}
			key := providerResourceKey{nodeID: rec.ID, providerID: prov.ID}
			res, ok := m.resources[key]
			enabled := config.ProviderEnabled(prov)
			if !ok {
				res = &providerResourceState{
					nodeID:       rec.ID,
					providerID:   prov.ID,
					capacity:     prov.Capacity,
					longCapacity: prov.LongContextCapacity,
					enabled:      enabled,
				}
				m.resources[key] = res
			} else {
				res.capacity = prov.Capacity
				res.longCapacity = prov.LongContextCapacity
				res.enabled = enabled
				res.orphan = false
			}
		}
	}

	for key, res := range m.resources {
		if res.orphan && res.inFlight == 0 {
			delete(m.resources, key)
		}
	}

	m.refreshPendingDeadlineLocked()
	m.pumpAllLocked()
}

// setProviderPoolPolicy replaces the manager-global provider-pool policy and the
// store reference in the same critical section. It also triggers a queue pump so
// the new max_queue bound takes effect on pending items immediately. Must be
// called with m.mu held.
func (m *modelQueueManager) setProviderPoolPolicyLocked(store *edgenode.NodeStore, policy groupPolicy) {
	m.store = store
	m.providerPoolPolicy = policy
	if store == nil {
		m.refreshPendingDeadlineLocked()
		m.pumpAllLocked()
		return
	}
	for _, res := range m.resources {
		res.orphan = true
	}
	for _, rec := range store.All() {
		for _, prov := range rec.Providers {
			if prov.ID == "" {
				continue
			}
			key := providerResourceKey{nodeID: rec.ID, providerID: prov.ID}
			res, ok := m.resources[key]
			enabled := config.ProviderEnabled(prov)
			if !ok {
				res = &providerResourceState{
					nodeID:       rec.ID,
					providerID:   prov.ID,
					capacity:     prov.Capacity,
					longCapacity: prov.LongContextCapacity,
					enabled:      enabled,
				}
				m.resources[key] = res
			} else {
				res.capacity = prov.Capacity
				res.longCapacity = prov.LongContextCapacity
				res.enabled = enabled
				res.orphan = false
			}
		}
	}
	for key, res := range m.resources {
		if res.orphan && res.inFlight == 0 {
			delete(m.resources, key)
		}
	}
	m.refreshPendingDeadlineLocked()
	m.pumpAllLocked()
}

// pendingProviderPoolCountLocked returns the total number of pending items that
// were enqueued under a provider-pool admission, summed across all model groups.
// Legacy adapter queue items are excluded so group-local max_queue for those
// paths is unaffected. Must be called with m.mu held.
func (m *modelQueueManager) pendingProviderPoolCountLocked() int {
	count := 0
	for _, g := range m.groups {
		for _, item := range g.queue {
			if item.providerPool {
				count++
			}
		}
	}
	return count
}

// refreshPendingDeadlineLocked recomputes every pending provider-pool item's
// deadline from its enqueue time and the current root provider-pool queue
// timeout. Items with providerPool=false are left untouched so legacy queue
// items keep their original deadline. Items whose enqueuedAt is zero
// (fixture-only items built directly without going through admit) skip rebase.
// A zero timeout disables the queue deadline and turns the waiter into a pure
// ctx-cancellation wait. A shorter timeout shortens the deadline and the
// notifier is drained then signalled so the waiter's select loop can pick up
// the fresh value. Must be called with m.mu held.
func (m *modelQueueManager) refreshPendingDeadlineLocked() {
	for _, g := range m.groups {
		for _, item := range g.queue {
			if !item.providerPool {
				continue
			}
			if item.enqueuedAt.IsZero() {
				continue
			}
			timeout := g.policy.queueTimeout
			if m.providerPoolPolicy.queueTimeoutSet {
				timeout = m.providerPoolPolicy.queueTimeout
			}
			newDeadline := deadlineFrom(item.enqueuedAt, timeout)
			if newDeadline.Equal(item.deadline) {
				continue
			}
			item.deadline = newDeadline
			if item.policyChange == nil {
				continue
			}
			// Drain before sending so the waiter loop always sees exactly one
			// signal after a policy change, regardless of how many refreshes
			// happened while the loop was processing a previous signal.
			select {
			case <-item.policyChange:
			default:
			}
			select {
			case item.policyChange <- struct{}{}:
			default:
			}
		}
	}
}

// activateNodeGenerationLocked restores every provider resource owned by nodeID
// to an available state for the accepted connection generation: it clears the
// orphan marker a disconnect set and advances the resource generation to the
// reconnect's (never lowering it), so a candidate resolved at the new generation
// is eligible to reserve and any leftover older-generation candidate stays
// fenced. In-flight counters are left as they are — the disconnect already
// settled the old generation's leases exactly once and this reconnect holds none
// yet — and a provider resource that no longer exists is intentionally not
// resurrected: the live resolver recreates it from the store on first admission.
// Must be called with m.mu held.
func (m *modelQueueManager) activateNodeGenerationLocked(nodeID string, generation uint64) {
	if nodeID == "" {
		return
	}
	// Preserve a same-generation duplicate activation, but discard every older
	// generation's runtime evidence. Configuration remains untouched.
	for key := range m.runtimeHealth {
		if key.nodeID == nodeID && key.generation != generation {
			delete(m.runtimeHealth, key)
		}
	}
	for _, res := range m.resources {
		if res.nodeID != nodeID {
			continue
		}
		res.orphan = false
		if generation > res.generation {
			res.generation = generation
		}
	}
}

func (m *modelQueueManager) providerRuntimeUnavailableLocked(nodeID string, generation uint64, providerID string) bool {
	if providerID == "" {
		return false
	}
	overlay := m.runtimeHealth[providerRuntimeHealthKey{
		nodeID: nodeID, generation: generation, providerID: providerID,
	}]
	return overlay != nil && overlay.unavailable
}

func (m *modelQueueManager) candidateRuntimeHealthyLocked(candidate *candidateNode) bool {
	if candidate == nil || candidate.entry == nil || candidate.providerID == "" {
		return true
	}
	return !m.providerRuntimeUnavailableLocked(candidate.entry.NodeID, candidate.generation, candidate.providerID)
}

// filterRuntimeHealthyCandidatesLocked removes runtime-unavailable provider
// candidates while preserving candidate order and legacy candidates. It is
// used by both immediate admission and queued re-resolution so effective
// provider eligibility has one source of truth.
func (m *modelQueueManager) filterRuntimeHealthyCandidatesLocked(candidates []candidateNode) []candidateNode {
	filtered := make([]candidateNode, 0, len(candidates))
	for i := range candidates {
		if m.candidateRuntimeHealthyLocked(&candidates[i]) {
			filtered = append(filtered, candidates[i])
		}
	}
	return filtered
}

// findLastColon returns the index of the last ':' in s, or -1 if not found.
func findLastColon(s string) int {
	for i := len(s) - 1; i >= 0; i-- {
		if s[i] == ':' {
			return i
		}
	}
	return -1
}

type providerResourceKey struct {
	nodeID     string
	providerID string
}

type providerResourceState struct {
	nodeID       string
	providerID   string
	capacity     int
	longCapacity int
	inFlight     int
	longInFlight int
	enabled      bool
	orphan       bool
	// generation is the highest connection generation that has reserved this
	// resource. It never decreases, so a disconnect fence for an older generation
	// cannot zero a slot a newer reconnect now owns, and a stale candidate cannot
	// reserve against a resource a newer generation already claimed.
	generation uint64
}

func (r *providerResourceState) canReserve(long bool) bool {
	if !r.enabled {
		return false
	}
	if r.inFlight >= r.capacity {
		return false
	}
	if long && r.longCapacity > 0 && r.longInFlight >= r.longCapacity {
		return false
	}
	return true
}

func (r *providerResourceState) reserve(long bool) {
	r.inFlight++
	if long && r.longCapacity > 0 {
		r.longInFlight++
	}
}

func (r *providerResourceState) release(long bool) {
	r.inFlight--
	if r.inFlight < 0 {
		r.inFlight = 0
	}
	if long {
		r.longInFlight--
		if r.longInFlight < 0 {
			r.longInFlight = 0
		}
	}
}

func (r *providerResourceState) snapshotCounts() (int, int) {
	return r.inFlight, r.longInFlight
}

// resourceForCandidateLocked returns the provider resource state for the given
// candidate and whether it is eligible for admission. Eligibility requires the
// candidate's (nodeID, providerID) to exist in the current NodeStore catalog
// and the resource not to be orphaned. When store is nil (legacy unit fixtures
// that rely on candidate-based resource creation), the previous behaviour of
// creating a fresh enabled resource from the candidate is preserved so existing
// tests continue to pass.
func (m *modelQueueManager) resourceForCandidateLocked(c *candidateNode) (*providerResourceState, bool) {
	key := providerResourceKey{nodeID: c.entry.NodeID, providerID: c.providerID}
	r, ok := m.resources[key]
	if ok {
		return r, true
	}

	if m.store == nil {
		initInflight := 0
		initLongInflight := 0
		slotKey := c.slotKey()
		for _, g := range m.groups {
			if val, ok := g.inflight[slotKey]; ok && val > initInflight {
				initInflight = val
			}
			if val, ok := g.longInflight[slotKey]; ok && val > initLongInflight {
				initLongInflight = val
			}
		}

		r = &providerResourceState{
			nodeID:       c.entry.NodeID,
			providerID:   c.providerID,
			capacity:     c.capacity,
			longCapacity: c.longContextCapacity,
			inFlight:     initInflight,
			longInFlight: initLongInflight,
			enabled:      true,
		}
		m.resources[key] = r
		return r, true
	}

	_, exists := m.store.FindByID(c.entry.NodeID)
	if !exists {
		return nil, false
	}

	rec, exists := m.store.FindByID(c.entry.NodeID)
	if !exists || rec == nil {
		return nil, false
	}
	hasProvider := false
	for _, prov := range rec.Providers {
		if prov.ID == c.providerID {
			hasProvider = true
			break
		}
	}
	if !hasProvider {
		return nil, false
	}

	initInflight := 0
	initLongInflight := 0
	slotKey := c.slotKey()
	for _, g := range m.groups {
		if val, ok := g.inflight[slotKey]; ok && val > initInflight {
			initInflight = val
		}
		if val, ok := g.longInflight[slotKey]; ok && val > initLongInflight {
			initLongInflight = val
		}
	}

	r = &providerResourceState{
		nodeID:       c.entry.NodeID,
		providerID:   c.providerID,
		capacity:     c.capacity,
		longCapacity: c.longContextCapacity,
		inFlight:     initInflight,
		longInFlight: initLongInflight,
		enabled:      true,
	}
	m.resources[key] = r
	return r, true
}
