import random
import unittest
from unittest.mock import patch
import torch
from lightnest import rv32
from lightnest.learned_patch import MAX_WINDOW
from lightnest.sequential_patch_rl import arithmetic, make_task, passes, rollout, terminal_reward

class FixedPolicy(torch.nn.Module):
    def __init__(self, position):
        super().__init__(); self.position=position; self.calls=0
        self.weight=torch.nn.Parameter(torch.tensor(1.))
    def forward(self,c,b):
        self.calls+=1
        scores=torch.full((len(c),257),-30.)+self.weight*0
        scores[:,self.position]=self.weight*30
        return scores,torch.ones((len(c),256,8))*self.weight

class SequentialTests(unittest.TestCase):
    def test_oracle(self):
        self.assertEqual(arithmetic((4,2,1),'sub'),1)
        self.assertEqual(arithmetic((2147483647,1,1),'add'),-2147483647)
    def test_generator_bounds_and_independent_execution(self):
        for split in ('train','unseen_bias','three_op_combined'):
            for seed in range(40):
                t=make_task(random.Random(seed),split)
                self.assertLessEqual(len(t.binary),MAX_WINDOW)
                # Test-only assembler repair never enters policy or training reward.
                target='sub' if 'to sub' in t.command or 'with sub' in t.command else 'add'
                words=[int.from_bytes(t.binary[i:i+4],'little') for i in range(0,len(t.binary),4)]
                words=[((w & ~(127<<25)) | ((32 if target=='sub' else 0)<<25)) if w&127==51 else w for w in words]
                repaired=b''.join(w.to_bytes(4,'little') for w in words)
                self.assertTrue(passes(repaired,t.cases))
    def test_none_and_budget(self):
        t=make_task(random.Random(3)); p=FixedPolicy(MAX_WINDOW)
        out,edits,h,_=rollout(p,[t.command],[t.binary])
        self.assertEqual(out,[t.binary]); self.assertEqual(edits,[0]); self.assertEqual(p.calls,1)
        p=FixedPolicy(0)
        with patch('lightnest.sequential_patch_rl.rv32.simulate',side_effect=AssertionError('oracle reached policy')):
            _,_,h,lp=rollout(p,[t.command],[t.binary],sample=True)
        self.assertEqual(p.calls,3); self.assertEqual(len(h[0]),3)
        lp.sum().backward(); self.assertIsNotNone(p.weight.grad)
    def test_input_bounds(self):
        p=FixedPolicy(0)
        for n in (0,4):
            with self.assertRaises(ValueError): rollout(p,['x'],[b'1234'],max_proposals=n)
        with self.assertRaises(ValueError): rollout(p,['x'],[b'x'*257])
    def test_terminal_reward(self):
        binary=b''.join(w.to_bytes(4,'little') for w in (rv32.addi('x10','x0',7),rv32.EBREAK))
        cases=[([1,2],7)]
        self.assertEqual(terminal_reward(binary,binary,cases,0),1)
        self.assertEqual(terminal_reward(binary,binary,cases,1),-1)
        self.assertEqual(terminal_reward(binary,binary,[([1,2],8)],0),-1)
        with self.assertRaises(ValueError): terminal_reward(binary,binary,[],0)

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