package sim

import (
	"container/list"
	"encoding/binary"
	"hash/fnv"
)

// BlockHash identifies one KV cache block, the unit of cache accounting.
type BlockHash uint64

// BlockHashes computes chained block hashes over a token sequence.
//
// Chaining (h_i = hash(h_{i-1}, block_i)) mirrors what vLLM does and is
// required: a block's identity must represent not only its own contents but
// every block before it. Otherwise two requests containing the same sentence
// in different contexts would falsely count as a cache hit.
//
// A trailing partial block is ignored — half blocks are never cached.
func BlockHashes(tokens []Token, blockSize int) []BlockHash {
	if blockSize <= 0 || len(tokens) < blockSize {
		return nil
	}
	n := len(tokens) / blockSize
	out := make([]BlockHash, 0, n)
	var prev BlockHash
	buf := make([]byte, 8)
	for i := range n {
		h := fnv.New64a()
		binary.LittleEndian.PutUint64(buf, uint64(prev))
		_, _ = h.Write(buf)
		for _, t := range tokens[i*blockSize : (i+1)*blockSize] {
			binary.LittleEndian.PutUint32(buf[:4], uint32(t))
			_, _ = h.Write(buf[:4])
		}
		prev = BlockHash(h.Sum64())
		out = append(out, prev)
	}
	return out
}

// PrefixCache models an engine's automatic prefix cache as a block-level LRU.
//
// A real engine uses a radix tree, but what matters for measuring routing
// quality is which blocks are still resident, and LRU models that faithfully.
type PrefixCache struct {
	capacity int // in blocks
	ll       *list.List
	m        map[BlockHash]*list.Element

	queries int64
	hits    int64
}

// NewPrefixCache creates a cache holding at most capacity blocks.
// A capacity of zero or less disables the cache (every lookup misses).
func NewPrefixCache(capacity int) *PrefixCache {
	return &PrefixCache{
		capacity: capacity,
		ll:       list.New(),
		m:        make(map[BlockHash]*list.Element),
	}
}

// Enabled reports whether the cache is active.
func (c *PrefixCache) Enabled() bool { return c != nil && c.capacity > 0 }

// Lookup returns the length, in blocks, of the longest **contiguous prefix**
// of hashes present in the cache, refreshing those blocks in the LRU.
//
// Contiguity matters: if a block in the middle was evicted, everything after
// it must be recomputed too, because attention depends on all preceding blocks.
func (c *PrefixCache) Lookup(hashes []BlockHash) int {
	if !c.Enabled() {
		return 0
	}
	c.queries += int64(len(hashes))
	n := 0
	for _, h := range hashes {
		el, ok := c.m[h]
		if !ok {
			break
		}
		c.ll.MoveToFront(el)
		n++
	}
	c.hits += int64(n)
	return n
}

// Insert stores blocks at the front of the LRU, evicting the oldest ones once
// capacity is exceeded.
func (c *PrefixCache) Insert(hashes []BlockHash) {
	if !c.Enabled() {
		return
	}
	for _, h := range hashes {
		if el, ok := c.m[h]; ok {
			c.ll.MoveToFront(el)
			continue
		}
		c.m[h] = c.ll.PushFront(h)
	}
	for c.ll.Len() > c.capacity {
		el := c.ll.Back()
		if el == nil {
			break
		}
		c.ll.Remove(el)
		delete(c.m, el.Value.(BlockHash))
	}
}

// Stats returns cumulative query and hit counts, in blocks.
func (c *PrefixCache) Stats() (queries, hits int64) {
	if c == nil {
		return 0, 0
	}
	return c.queries, c.hits
}

// Len returns the number of blocks currently cached.
func (c *PrefixCache) Len() int {
	if !c.Enabled() {
		return 0
	}
	return c.ll.Len()
}