package main

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

// Model Pool: decides whether a model may be loaded now, and frees idle models to make room.
// Rules (see docs/DESIGN-TH.md section 3):
//   - a model that is loaded stays loaded for a long time (OLLAMA_KEEP_ALIVE is 24h in the installer): it is unloaded when memory is
//     needed (see below, and freeUp for image/video jobs), when the owner comes back, or when nobody has used it for a day;
//   - a request for another model, when memory is short, evicts idle models (least recently used first);
//   - a model with requests in flight is never evicted; the new request waits, then gets 503;
//   - a freshly loaded model stays at least minResidency so two models cannot ping-pong each other out.

type poolBackend interface {
	loaded() (map[string]int64, error)    // model -> bytes in memory
	installed() (map[string]int64, error) // model -> bytes on disk
	unload(name string) error
	memory() (total, avail uint64)
}

type poolError struct {
	Status int
	Msg    string
}

func (e *poolError) Error() string { return e.Msg }

type reservation struct {
	bytes uint64
	at    time.Time
}

type pool struct {
	b            poolBackend
	now          func() time.Time
	sleep        func(time.Duration)
	reserveBytes uint64        // memory always left for the owner / system
	headroom     uint64        // extra per model for the KV cache and runtime buffers
	minResidency time.Duration // keep a freshly loaded model at least this long
	maxWait      time.Duration // how long a request may wait for busy models to finish

	mu       sync.Mutex
	inflight map[string]int
	lastUse  map[string]time.Time
	seenAt   map[string]time.Time // when we first saw the model loaded
	reserved map[string]reservation
}

func newPool(b poolBackend) *pool {
	return &pool{
		b: b, now: time.Now, sleep: time.Sleep,
		reserveBytes: 8 << 30, headroom: 2 << 30,
		minResidency: 60 * time.Second, maxWait: 90 * time.Second,
		inflight: map[string]int{}, lastUse: map[string]time.Time{},
		seenAt: map[string]time.Time{}, reserved: map[string]reservation{},
	}
}

// resolveModel maps the name a client used to the exact installed name:
// an exact match wins; otherwise "name" may mean "name:latest" (Ollama's convention).
func resolveModel(name string, installed map[string]int64) string {
	name = strings.TrimSpace(name)
	if _, ok := installed[name]; ok {
		return name
	}
	if !strings.Contains(name, ":") {
		if _, ok := installed[name+":latest"]; ok {
			return name + ":latest"
		}
	}
	return name
}

// acquire blocks until the model may be used, then returns a release func.
func (p *pool) acquire(ctx context.Context, model string) (func(), *poolError) {
	name := model
	deadline := p.now().Add(p.maxWait)
	for {
		rel, wait, perr := p.tryAcquire(name)
		if perr != nil {
			return nil, perr
		}
		if rel != nil {
			return rel, nil
		}
		_ = wait
		if p.now().After(deadline) {
			return nil, &poolError{503, "not enough free memory to load " + name + " right now (other models are busy); try again shortly"}
		}
		select {
		case <-ctx.Done():
			return nil, &poolError{499, "request cancelled"}
		default:
		}
		p.sleep(500 * time.Millisecond)
	}
}

func (p *pool) release(name string) func() {
	var once sync.Once
	return func() {
		once.Do(func() {
			p.mu.Lock()
			defer p.mu.Unlock()
			if p.inflight[name] > 0 {
				p.inflight[name]--
			}
			p.lastUse[name] = p.now()
		})
	}
}

// tryAcquire makes one admission decision. (nil, true, nil) means "not now, retry".
func (p *pool) tryAcquire(requested string) (func(), bool, *poolError) {
	p.mu.Lock()
	defer p.mu.Unlock()
	now := p.now()

	inst, err := p.b.installed()
	if err != nil {
		return nil, false, &poolError{502, "model runtime not reachable: " + err.Error()}
	}
	name := resolveModel(requested, inst)
	loaded, err := p.b.loaded()
	if err != nil {
		return nil, false, &poolError{502, "model runtime not reachable: " + err.Error()}
	}
	for n := range loaded {
		if _, ok := p.seenAt[n]; !ok {
			p.seenAt[n] = now
		}
	}
	for n := range p.seenAt {
		if _, ok := loaded[n]; !ok {
			delete(p.seenAt, n)
		}
	}
	for n, r := range p.reserved {
		if _, ok := loaded[n]; ok || now.Sub(r.at) > 10*time.Minute {
			delete(p.reserved, n)
		}
	}

	if _, ok := loaded[name]; ok {
		p.inflight[name]++
		p.lastUse[name] = now
		return p.release(name), false, nil
	}
	if _, ok := p.reserved[name]; ok { // another request is already loading it
		p.inflight[name]++
		return p.release(name), false, nil
	}

	size, ok := inst[name]
	if !ok {
		return nil, false, &poolError{404, fmt.Sprintf("model %q is not installed on this machine", name)}
	}
	need := uint64(size) + uint64(size)/8 + p.headroom

	total, avail := p.b.memory()
	if total == 0 { // memory unknown on this platform: do not block
		p.reserved[name] = reservation{need, now}
		p.inflight[name]++
		return p.release(name), false, nil
	}
	if need+p.reserveBytes > total {
		return nil, false, &poolError{503, fmt.Sprintf("model %q needs about %.1f GB, more than this machine can spare", name, float64(need)/(1<<30))}
	}

	var pending uint64
	for _, r := range p.reserved {
		pending += r.bytes
	}
	free := int64(avail) - int64(p.reserveBytes) - int64(pending)

	if free < int64(need) {
		type cand struct {
			name string
			size int64
			last time.Time
		}
		var cs []cand
		for n, sz := range loaded {
			if p.inflight[n] > 0 {
				continue // busy: never cut off a running request
			}
			if now.Sub(p.seenAt[n]) < p.minResidency {
				continue // just loaded, give it a chance to be used
			}
			cs = append(cs, cand{n, sz, p.lastUse[n]})
		}
		sort.Slice(cs, func(i, j int) bool { return cs[i].last.Before(cs[j].last) })
		for _, c := range cs {
			if free >= int64(need) {
				break
			}
			if err := p.b.unload(c.name); err == nil {
				free += c.size
				delete(p.seenAt, c.name)
			}
		}
	}
	if free < int64(need) {
		return nil, true, nil // wait for busy models to finish
	}
	p.reserved[name] = reservation{need, now}
	p.inflight[name]++
	p.lastUse[name] = now
	return p.release(name), false, nil
}

// freeUp makes room for something that is not a chat model (an image or video job) by unloading idle models, the least
// recently used first, and only as many as needed. Models stay in memory as long as there is room for everything: nothing is
// unloaded just because it has been idle for a while. A model with requests in flight is never touched. It reports whether
// `need` bytes (on top of the memory kept for the owner) are free after the unloading, and what it unloaded.
func (p *pool) freeUp(need uint64) (bool, []string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	loaded, err := p.b.loaded()
	if err != nil {
		return false, nil
	}
	_, avail := p.b.memory()
	var pending uint64
	for _, r := range p.reserved {
		pending += r.bytes
	}
	free := int64(avail) - int64(p.reserveBytes) - int64(pending)
	if free >= int64(need) {
		return true, nil
	}
	type cand struct {
		name string
		size int64
		last time.Time
	}
	var cs []cand
	for n, sz := range loaded {
		if p.inflight[n] > 0 {
			continue
		}
		cs = append(cs, cand{n, sz, p.lastUse[n]}) // never used through the pool: the oldest
	}
	sort.Slice(cs, func(i, j int) bool { return cs[i].last.Before(cs[j].last) })
	var gone []string
	for _, c := range cs {
		if free >= int64(need) {
			break
		}
		if err := p.b.unload(c.name); err == nil {
			free += c.size
			delete(p.seenAt, c.name)
			gone = append(gone, c.name)
		}
	}
	return free >= int64(need), gone
}

// totalInflight is the number of requests running on this machine right now.
func (p *pool) totalInflight() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	n := 0
	for _, c := range p.inflight {
		n += c
	}
	return n
}