"""Run with python -m unittest -v. Tests exercise deliberate failure controls."""
import unittest
import numpy as np
from kv_lab import (TinyDecoder, Config, PrefixCache, PagePool, softmax,
                    online_attention, pack_int4, unpack_int4, max_error)

class KVLabTests(unittest.TestCase):
    def setUp(self):
        self.model = TinyDecoder()
        self.tokens = (list(range(31)) * 2)[:53]

    def test_incremental_matches_full(self):
        full, _ = self.model.forward(self.tokens)
        _, pre = self.model.forward(self.tokens[:32])
        tail, _ = self.model.forward(self.tokens[32:], pre)
        self.assertLess(max_error(full[32:], tail), 1e-10)

    def test_prefix_hit_processes_only_suffix(self):
        cache = PrefixCache(self.model)
        cache.request(self.tokens)
        b = self.tokens[:42] + [39, 40, 38]
        before = self.model.processed_positions
        logits, _, usage = cache.request(b)
        self.assertEqual(usage['hit_tokens'], 32)
        self.assertEqual(self.model.processed_positions-before, 13)
        ref,_ = self.model.forward(b)
        self.assertLess(max_error(logits,ref[-1]),1e-10)

    def test_full_block_hit_leaves_one_block(self):
        cache = PrefixCache(self.model)
        cache.request(self.tokens[:48])
        _,_,u = cache.request(self.tokens[:48])
        self.assertEqual(u['hit_tokens'], 32)
        self.assertEqual(u['computed_tokens'], 16)

    def test_early_change_breaks_hash_chain(self):
        cache = PrefixCache(self.model)
        cache.request(self.tokens)
        changed = self.tokens.copy();changed[0]=40
        self.assertEqual(cache.lookup(changed)[0], 0)

    def test_identity_includes_weights_and_namespace(self):
        c1=PrefixCache(self.model,namespace='alpha')
        c2=PrefixCache(TinyDecoder(Config(seed=8)),namespace='alpha')
        c3=PrefixCache(self.model,namespace='beta')
        self.assertNotEqual(c1.base,c2.base)
        self.assertNotEqual(c1.base,c3.base)

    def test_past_not_mutated_and_no_private_alias(self):
        _, past = self.model.forward(self.tokens[:32])
        saved=past[0][0].copy()
        self.model.forward(self.tokens[32:],past)
        np.testing.assert_array_equal(saved,past[0][0])
        cache=PrefixCache(self.model);cache.request(self.tokens)
        _,hit=cache.lookup(self.tokens)
        hit[0][0][:]=0
        _,hit2=cache.lookup(self.tokens)
        self.assertGreater(np.max(np.abs(hit2[0][0])),0)

    def test_cache_shape_counts_kv_heads_not_q_heads(self):
        _,state=self.model.forward(self.tokens)
        c=self.model.cfg
        size=sum(k.nbytes+v.nbytes for k,v in state)
        self.assertEqual(size,2*c.layers*len(self.tokens)*c.kv_heads*c.head_dim*8)

    def test_online_global_softmax(self):
        rng=np.random.default_rng(2);q=rng.normal(size=8)
        k=rng.normal(size=(29,8));v=rng.normal(size=(29,5))
        ref=softmax(k@q/np.sqrt(8))@v
        for block in [1,3,7,29,50]:
            self.assertLess(max_error(ref,online_attention(q,k,v,block)),1e-12)

    def test_packed_int4_bytes_and_bound(self):
        x=np.linspace(-3,3,64,dtype=np.float32).reshape(4,16)
        packed,scales,shape=pack_int4(x)
        y=unpack_int4(packed,scales,shape)
        self.assertEqual(packed.nbytes,32)
        self.assertEqual(scales.nbytes,16)
        self.assertLessEqual(max_error(x,y),float(scales.max())/2+1e-6)
        z=np.zeros(17,dtype=np.float32)
        np.testing.assert_array_equal(unpack_int4(*pack_int4(z)),z)

    def test_reference_count_protects_second_reader(self):
        pool=PagePool(1);p=pool.allocate(np.array([1.0,2.0]))
        pool.retain(p);pool.release(p)
        with self.assertRaises(MemoryError):pool.allocate(np.array([9.0]))
        np.testing.assert_array_equal(pool.pages[p],[1.0,2.0])
        pool.release(p)
        self.assertEqual(pool.allocate(np.array([9.0])),p)

    def test_copy_on_write_separates_branches(self):
        pool=PagePool(2);p=pool.allocate(np.array([1.0,2.0,0.0]))
        pool.retain(p);branch=pool.private_for_write(p)
        pool.pages[branch][2]=9.0
        self.assertEqual(pool.pages[p][2],0.0)
        self.assertEqual(pool.refs[p],1);self.assertEqual(pool.refs[branch],1)

    def test_double_release_rejected(self):
        pool=PagePool();p=pool.allocate(np.zeros(3));pool.release(p)
        with self.assertRaises(RuntimeError):pool.release(p)

    def test_wrong_splice_is_detectable(self):
        doc=[8,9,10,11,12]
        _,old=self.model.forward([1,2]+doc)
        _,newlead=self.model.forward([3,4])
        wrong=[(np.concatenate([newlead[i][0],k[:,2:]],axis=1),
                np.concatenate([newlead[i][1],v[:,2:]],axis=1))
               for i,(k,v) in enumerate(old)]
        bad,_=self.model.forward([13],wrong)
        good,_=self.model.forward([3,4]+doc+[13])
        self.assertGreater(max_error(bad[-1],good[-1]),1e-6)

    def test_eviction_recovers_correctly(self):
        cache=PrefixCache(self.model,capacity_blocks=2)
        cache.request(self.tokens)
        self.assertLessEqual(len(cache.blocks),2)
        out,_,_=cache.request(self.tokens)
        full,_=self.model.forward(self.tokens)
        self.assertLess(max_error(out,full[-1]),1e-10)

    def test_invalid_inputs_rejected(self):
        for t in [[],[-1],[41],[1.2]]:
            with self.assertRaises(ValueError):self.model.forward(t)

if __name__=='__main__':unittest.main()