// Package lock implements the two-mode GPU arbitration lock: any number of
// concurrent LLM requests ("readers") or exactly one image job ("writer"),
// with image jobs taking priority over newly arriving LLM requests. An
// external hold (SetExternal) blocks new grants of both kinds while a
// foreign process — e.g. a game — holds the GPU; in-flight work drains.
package lock

import (
	"context"
	"log/slog"
	"sync"
	"time"
)

// State is the current GPU occupancy state.
type State string

const (
	StateIdle     State = "idle"
	StateLLM      State = "llm"
	StateImage    State = "image"
	StateExternal State = "external"
)

type imageWaiter struct{ id uint64 }

// Lock is a writer-preferring two-mode lock. The zero value is not usable;
// construct with New.
type Lock struct {
	mu     sync.Mutex
	change chan struct{} // closed and replaced on every state change

	n           int  // LLM requests in flight
	llmWaiting  int  // LLM requests blocked waiting for the GPU
	imageActive bool // an image job holds the GPU
	imageQ      []imageWaiter
	nextID      uint64
	external    string    // non-empty: a foreign process (e.g. a game) holds the GPU
	detail      string    // what the current holder is doing (best effort)
	since       time.Time // when the current state began

	log *slog.Logger
}

// New returns a ready-to-use Lock. log may be nil; if set, every state
// transition is logged at debug level.
func New(log *slog.Logger) *Lock {
	return &Lock{change: make(chan struct{}), log: log, since: time.Now()}
}

// broadcast wakes all waiters. Call with mu held.
func (l *Lock) broadcast() {
	close(l.change)
	l.change = make(chan struct{})
}

func (l *Lock) logTransition(msg string, args ...any) {
	if l.log != nil {
		l.log.Debug(msg, args...)
	}
}

// SetExternal records that a process outside gpu-turnstile's control (a
// game, another ML job) holds the GPU: new LLM and image grants block until
// ClearExternal. In-flight work is not preempted. holder describes the
// process for logs and busy responses.
func (l *Lock) SetExternal(holder string) {
	l.mu.Lock()
	l.external = holder
	l.since = time.Now()
	l.broadcast()
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateExternal, "holder", holder)
}

// ClearExternal lifts the external hold; waiting LLM and image requests
// proceed.
func (l *Lock) ClearExternal() {
	l.mu.Lock()
	l.external = ""
	l.since = time.Now()
	l.broadcast()
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateIdle)
}

// External returns the current external holder, or "" when none.
func (l *Lock) External() string {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.external
}

// AcquireLLM blocks until no image job is active or pending, then registers
// one in-flight LLM request. Returns ctx.Err() if the context is cancelled
// while waiting; no state is changed in that case.
func (l *Lock) AcquireLLM(ctx context.Context) error {
	l.mu.Lock()
	waiting := false
	for l.imageActive || len(l.imageQ) > 0 || l.external != "" {
		if !waiting {
			l.llmWaiting++
			waiting = true
		}
		ch := l.change
		l.mu.Unlock()
		select {
		case <-ctx.Done():
			l.mu.Lock()
			l.llmWaiting--
			l.mu.Unlock()
			return ctx.Err()
		case <-ch:
		}
		l.mu.Lock()
	}
	if waiting {
		l.llmWaiting--
	}
	l.n++
	if l.n == 1 {
		l.since = time.Now()
	}
	n := l.n
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
	return nil
}

// TryAcquireLLM acquires one in-flight LLM slot without waiting and
// reports whether it succeeded. It fails when an image job is active or
// pending or an external hold is set.
func (l *Lock) TryAcquireLLM() bool {
	l.mu.Lock()
	if l.imageActive || len(l.imageQ) > 0 || l.external != "" {
		l.mu.Unlock()
		return false
	}
	l.n++
	n := l.n
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
	return true
}

// ReleaseLLM marks one LLM request as finished.
func (l *Lock) ReleaseLLM() {
	l.mu.Lock()
	l.n--
	n := l.n
	if l.n == 0 {
		l.since = time.Now()
		l.detail = ""
		l.broadcast()
	}
	l.mu.Unlock()
	if n == 0 {
		l.logTransition("lock transition", "state", StateIdle, "llm_inflight", 0)
	}
}

// AcquireImage queues the caller FIFO behind other image jobs, blocks new
// LLM requests immediately, and waits for in-flight LLM requests to drain
// before granting exclusive GPU access. Returns ctx.Err() if the context is
// cancelled while waiting; the caller is removed from the queue.
func (l *Lock) AcquireImage(ctx context.Context) error {
	l.mu.Lock()
	l.nextID++
	w := imageWaiter{id: l.nextID}
	l.imageQ = append(l.imageQ, w)
	pending := len(l.imageQ)
	l.broadcast() // new LLM requests must block from now on
	l.mu.Unlock()
	l.logTransition("image job pending", "image_queue", pending)

	for {
		l.mu.Lock()
		if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive && l.external == "" {
			l.imageQ = l.imageQ[1:]
			l.imageActive = true
			l.since = time.Now()
			l.mu.Unlock()
			l.logTransition("lock transition", "state", StateImage)
			return nil
		}
		ch := l.change
		l.mu.Unlock()
		select {
		case <-ctx.Done():
			l.mu.Lock()
			for i, q := range l.imageQ {
				if q.id == w.id {
					l.imageQ = append(l.imageQ[:i], l.imageQ[i+1:]...)
					break
				}
			}
			l.broadcast()
			l.mu.Unlock()
			l.logTransition("image job wait cancelled", "image_queue", len(l.imageQ))
			return ctx.Err()
		case <-ch:
		}
	}
}

// ReleaseImage frees the GPU after an image job.
func (l *Lock) ReleaseImage() {
	l.mu.Lock()
	l.imageActive = false
	l.since = time.Now()
	l.detail = ""
	l.broadcast()
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateIdle)
}

// Snapshot reports the current state, the number of in-flight LLM requests,
// and whether an image job is active or waiting.
func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
	l.mu.Lock()
	defer l.mu.Unlock()
	switch {
	case l.imageActive:
		state = StateImage
	case l.n > 0:
		state = StateLLM
	case l.external != "":
		state = StateExternal
	default:
		state = StateIdle
	}
	return state, l.n, l.imageActive || len(l.imageQ) > 0
}

// SetDetail records what the current holder is doing (e.g. the request
// path), for status displays. Best effort: overwritten by each new holder,
// cleared when the GPU goes idle.
func (l *Lock) SetDetail(detail string) {
	l.mu.Lock()
	l.detail = detail
	l.mu.Unlock()
}

// Status is a point-in-time view of the lock for monitoring.
type Status struct {
	State       State
	Detail      string
	LLMInflight int
	LLMWaiting  int
	ImageActive bool
	ImageQueue  int
	External    string
	Since       time.Time
}

// Status reports the full lock state, including waiters and how long the
// current state has held.
func (l *Lock) Status() Status {
	l.mu.Lock()
	defer l.mu.Unlock()
	s := Status{
		Detail:      l.detail,
		LLMInflight: l.n,
		LLMWaiting:  l.llmWaiting,
		ImageActive: l.imageActive,
		ImageQueue:  len(l.imageQ),
		External:    l.external,
		Since:       l.since,
	}
	switch {
	case l.imageActive:
		s.State = StateImage
	case l.n > 0:
		s.State = StateLLM
	case l.external != "":
		s.State = StateExternal
	default:
		s.State = StateIdle
	}
	return s
}