import importlib.util
import json
import os
import unittest

from beamboot.chip8 import Chip8, Chip8Error, FONT_START, PROGRAM_START, Quirks
from beamboot.chip8_asm import (SUPPORTED, AsmError, assemble, check_subset, longest_path,
                                 op_class)

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
WIN_INSTRUCTION_BUDGET = 60000   # measured: 56100 for the scripted run
WIN_FRAME_BUDGET = 200           # measured: 187


def machine(src, **kw):
    rom, syms = assemble(src)
    m = Chip8(**kw)
    m.load(rom)
    m.syms = syms
    return m


def run(src, n, **kw):
    m = machine(src, **kw)
    m.run(n)
    return m


class OpcodeTests(unittest.TestCase):
    def test_flow_and_skips(self):
        m = run("JP 0x208\nLD V0,1\nLD V0,2\nLD V0,3\nLD V1,7", 1)
        self.assertEqual(m.pc, 0x208)
        m = run("LD V0,5\nSE V0,5\nLD V1,1\nLD V2,2", 3)
        self.assertEqual((m.v[1], m.v[2]), (0, 2))
        m = run("LD V0,5\nSE V0,6\nLD V1,1\nLD V2,2", 4)
        self.assertEqual((m.v[1], m.v[2]), (1, 2))
        m = run("LD V0,5\nSNE V0,6\nLD V1,1\nLD V2,2", 3)
        self.assertEqual((m.v[1], m.v[2]), (0, 2))
        m = run("LD V0,5\nLD V1,5\nSE V0,V1\nLD V2,1\nLD V3,2", 4)
        self.assertEqual((m.v[2], m.v[3]), (0, 2))
        m = run("LD V0,5\nLD V1,6\nSNE V0,V1\nLD V2,1\nLD V3,2", 4)
        self.assertEqual((m.v[2], m.v[3]), (0, 2))

    def test_loads_and_adds(self):
        m = run("LD V3,0xAB\nLD V4,V3\nADD V3,0x60\nADD V5,1", 4)
        self.assertEqual((m.v[3], m.v[4], m.v[5], m.v[0xF]), (0x0B, 0xAB, 1, 0))

    def test_alu(self):
        m = run("LD V0,0xF0\nLD V1,0x3C\nLD V2,V0\nOR V2,V1\nLD V3,V0\nAND V3,V1\n"
                "LD V4,V0\nXOR V4,V1", 8)
        self.assertEqual((m.v[2], m.v[3], m.v[4]), (0xFC, 0x30, 0xCC))
        m = run("LD V0,200\nLD V1,100\nADD V0,V1", 3)
        self.assertEqual((m.v[0], m.v[0xF]), (44, 1))
        m = run("LD V0,10\nLD V1,3\nADD V0,V1", 3)
        self.assertEqual((m.v[0], m.v[0xF]), (13, 0))
        m = run("LD V0,5\nLD V1,7\nSUB V0,V1", 3)
        self.assertEqual((m.v[0], m.v[0xF]), (254, 0))
        m = run("LD V0,7\nLD V1,5\nSUB V0,V1", 3)
        self.assertEqual((m.v[0], m.v[0xF]), (2, 1))
        m = run("LD V0,5\nLD V1,7\nSUBN V0,V1", 3)
        self.assertEqual((m.v[0], m.v[0xF]), (2, 1))
        m = run("LD V0,7\nLD V1,5\nSUBN V0,V1", 3)
        self.assertEqual((m.v[0], m.v[0xF]), (254, 0))

    def test_shift_quirk_vip_and_modern(self):
        src = "LD V1,0x81\nLD V0,0x00\nSHR V0,V1"
        m = run(src, 3)
        self.assertEqual((m.v[0], m.v[0xF]), (0x40, 1))          # VIP: Vx = Vy >> 1
        m = run(src, 3, quirks=Quirks(shift_uses_vy=False))
        self.assertEqual((m.v[0], m.v[0xF]), (0x00, 0))          # in-place on Vx
        m = run("LD V1,0x81\nLD V0,0\nSHL V0,V1", 3)
        self.assertEqual((m.v[0], m.v[0xF]), (0x02, 1))
        m = run("LD V0,0x81\nSHL V0\nLD V1,3\nSHR V1", 4)       # one-operand form
        self.assertEqual((m.v[0], m.v[1], m.v[0xF]), (0x02, 1, 1))

    def test_vf_reset_quirk(self):
        m = run("LD VF,1\nLD V0,1\nLD V1,2\nOR V0,V1", 4)
        self.assertEqual(m.v[0xF], 0)
        m = run("LD VF,1\nLD V0,1\nLD V1,2\nOR V0,V1", 4, quirks=Quirks(vf_reset=False))
        self.assertEqual(m.v[0xF], 1)

    def test_index_register_and_jump_offset(self):
        m = run("LD I,0x345\nLD V0,5\nADD I,V0", 3)
        self.assertEqual(m.i, 0x34A)
        m = run("LD V0,4\nJP V0,0x300", 2)
        self.assertEqual(m.pc, 0x304)

    def test_random_is_deterministic(self):
        a = run("RND V0,0xFF\nRND V1,0x0F", 2, seed=7)
        b = run("RND V0,0xFF\nRND V1,0x0F", 2, seed=7)
        self.assertEqual((a.v[0], a.v[1]), (b.v[0], b.v[1]))
        self.assertLess(a.v[1], 16)

    def test_bcd_font_and_load_store(self):
        m = run("LD I,0x400\nLD V0,239\nLD B,V0", 3)
        self.assertEqual(list(m.mem[0x400:0x403]), [2, 3, 9])
        m = run("LD V0,0xA\nLD F,V0", 2)
        self.assertEqual(m.i, FONT_START + 50)
        self.assertEqual(m.mem[m.i], 0xF0)
        m = run("LD I,0x400\nLD V0,1\nLD V1,2\nLD V2,3\nLD [I],V2\nLD I,0x400\nLD V2,[I]", 7)
        self.assertEqual(list(m.mem[0x400:0x403]), [1, 2, 3])
        self.assertEqual(m.i, 0x403)                             # VIP: I += x+1
        m = run("LD I,0x400\nLD V2,[I]", 2, quirks=Quirks(increment_i=False))
        self.assertEqual(m.i, 0x400)

    def test_timers_only_tick_explicitly(self):
        m = run("LD V0,3\nLD DT,V0\nLD ST,V0\nLD V1,DT", 4)
        self.assertEqual((m.dt, m.st, m.v[1]), (3, 3, 3))
        m.tick_timers()
        self.assertEqual((m.dt, m.st), (2, 2))
        for _ in range(5):
            m.tick_timers()
        self.assertEqual((m.dt, m.st), (0, 0))

    def test_keys(self):
        src = "LD V0,5\nSKP V0\nLD V1,1\nLD V2,2\nSKNP V0\nLD V3,3\nLD V4,4"
        m = machine(src)
        m.set_key(5, True)
        m.run(6)
        self.assertEqual((m.v[1], m.v[2], m.v[3], m.v[4]), (0, 2, 3, 4))
        m = run(src, 6)
        self.assertEqual((m.v[1], m.v[2], m.v[3], m.v[4]), (1, 2, 0, 4))

    def test_wait_for_key(self):
        m = machine("LD V0,K\nLD V1,9")
        for _ in range(3):
            rec = m.step()
            self.assertTrue(rec.waiting)
            self.assertEqual(m.pc, PROGRAM_START)
        m.set_key(0xB, True)
        m.step()
        m.step()
        self.assertEqual((m.v[0], m.v[1]), (0xB, 9))

    def test_cls(self):
        m = machine("LD I,0x050\nDRW V0,V0,5\nCLS")
        m.run(2)
        self.assertTrue(any(any(r) for r in m.display))
        rec = m.step()
        self.assertTrue(rec.display_changed)
        self.assertFalse(any(any(r) for r in m.display))

    def test_bad_opcode_raises(self):
        m = Chip8()
        m.load(bytes([0x5A, 0xB1]))
        with self.assertRaises(Chip8Error):
            m.step()


class DrawTests(unittest.TestCase):
    def test_xor_and_collision(self):
        m = machine("LD V0,2\nLD V1,1\nLD I,0x300\nDRW V0,V1,2\nDRW V0,V1,2")
        m.mem[0x300:0x302] = bytes([0xC0, 0x81])
        m.run(4)
        self.assertEqual(m.v[0xF], 0)
        self.assertEqual(m.display[1][2:4], [1, 1])
        self.assertEqual(m.display[2][2], 1)
        self.assertEqual(m.display[2][3], 0)
        self.assertEqual(m.display[2][9], 1)
        rec = m.step()                                    # draw again -> erases, VF=1
        self.assertEqual(m.v[0xF], 1)
        self.assertFalse(any(any(r) for r in m.display))
        self.assertEqual(rec.vf, 1)
        self.assertTrue(rec.display_changed)

    def test_partial_overlap_sets_vf(self):
        m = machine("LD I,0x300\nLD V0,0\nLD V1,0\nDRW V0,V1,1\nLD V0,4\nDRW V0,V1,1")
        m.mem[0x300] = 0xFF
        m.run(6)
        self.assertEqual(m.v[0xF], 1)
        self.assertEqual(m.display[0], [1] * 4 + [0] * 4 + [1] * 4 + [0] * 52)

    def test_clipping_and_wrapped_start(self):
        m = machine("LD I,0x300\nLD V0,60\nLD V1,30\nDRW V0,V1,4")
        m.mem[0x300:0x304] = bytes([0xFF] * 4)
        m.run(4)
        self.assertEqual(sum(map(sum, m.display)), 4 * 2)   # 4 px wide x 2 rows survive
        m = machine("LD I,0x300\nLD V0,0x44\nLD V1,0x22\nDRW V0,V1,1")  # 68%64=4, 34%32=2
        m.mem[0x300] = 0x80
        m.run(4)
        self.assertEqual(m.display[2][4], 1)

    def test_draw_wrap_quirk_off(self):
        m = machine("LD I,0x300\nLD V0,62\nLD V1,0\nDRW V0,V1,1",
                    quirks=Quirks(clip_sprites=False))
        m.mem[0x300] = 0xF0
        m.run(4)
        self.assertEqual(m.display[0][62:], [1, 1])
        self.assertEqual(m.display[0][:2], [1, 1])

    def test_display_rows_hex(self):
        m = machine("LD I,0x300\nLD V0,0\nDRW V0,V0,1")
        m.mem[0x300] = 0x81
        m.run(3)
        rows = m.display_rows()
        self.assertEqual(len(rows), 32)
        self.assertEqual(rows[0], "8100000000000000")


class StackTests(unittest.TestCase):
    def test_call_ret(self):
        m = machine("CALL sub\nLD V1,2\nJP end\nsub: LD V0,1\nRET\nend: JP end")
        m.run(6)
        self.assertEqual((m.v[0], m.v[1], m.sp), (1, 2, 0))

    def test_nested_16_levels_and_overflow(self):
        src = "start: CALL start"
        m = machine(src)
        m.run(16)
        self.assertEqual(m.sp, 16)
        with self.assertRaises(Chip8Error):
            m.step()
        with self.assertRaises(Chip8Error):
            machine("RET").step()

    def test_return_addresses_are_lifo(self):
        m = machine("CALL a\nHALT: JP HALT\na: CALL b\nRET\nb: RET")
        m.run(2)
        self.assertEqual(m.stack[:2], [0x202, 0x206])
        m.run(3)
        self.assertEqual((m.sp, m.pc), (0, 0x202))


class StepRecordTests(unittest.TestCase):
    def test_record_contents(self):
        m = machine("LD V3,9\nLD I,0x400\nLD [I],V3\nLD V3,[I]")
        r = m.step()
        self.assertEqual((r.pc, r.opcode, r.next_pc, r.regs), (0x200, 0x6309, 0x202, {"V3": 9}))
        r = m.step()
        self.assertEqual(r.regs, {"I": 0x400})
        r = m.step()
        self.assertEqual(r.mem, {0x403: 9})          # only bytes that changed are listed
        self.assertEqual(r.regs["I"], 0x404)
        self.assertEqual(m.instructions, 3)

    def test_replay_determinism(self):
        rom, _ = assemble("LD V0,1\nloop: ADD V0,3\nRND V1,0xFF\nJP loop")
        a, b = Chip8(seed=3), Chip8(seed=3)
        a.load(rom)
        b.load(rom)
        for _ in range(200):
            ra, rb = a.step(), b.step()
            self.assertEqual(ra, rb)


class AssemblerTests(unittest.TestCase):
    def test_labels_data_equ_and_expressions(self):
        rom, syms = assemble("N EQU 4\nstart: LD V0,N\nJP start+2\nDW 0xBEEF\nDB 1,2,N+1\nend:")
        self.assertEqual(rom, bytes([0x60, 0x04, 0x12, 0x02, 0xBE, 0xEF, 1, 2, 5]))
        self.assertEqual((syms["start"], syms["end"], syms["N"]), (0x200, 0x209, 4))

    def test_every_mnemonic_encodes(self):
        src = """CLS
RET
SYS 0x123
JP 0x234
CALL 0x345
SE V1,0x22
SNE V1,0x22
SE V1,V2
SNE V1,V2
LD V1,0x22
ADD V1,0x22
LD V1,V2
OR V1,V2
AND V1,V2
XOR V1,V2
ADD V1,V2
SUB V1,V2
SHR V1,V2
SUBN V1,V2
SHL V1,V2
LD I,0x456
JP V0,0x567
RND V1,0x22
DRW V1,V2,5
SKP V1
SKNP V1
LD V1,DT
LD V1,K
LD DT,V1
LD ST,V1
ADD I,V1
LD F,V1
LD B,V1
LD [I],V1
LD V1,[I]"""
        rom, _ = assemble(src)
        ops = [rom[i] << 8 | rom[i + 1] for i in range(0, len(rom), 2)]
        self.assertEqual(ops, [
            0x00E0, 0x00EE, 0x0123, 0x1234, 0x2345, 0x3122, 0x4122, 0x5120, 0x9120,
            0x6122, 0x7122, 0x8120, 0x8121, 0x8122, 0x8123, 0x8124, 0x8125, 0x8126,
            0x8127, 0x812E, 0xA456, 0xB567, 0xC122, 0xD125, 0xE19E, 0xE1A1, 0xF107,
            0xF10A, 0xF115, 0xF118, 0xF11E, 0xF129, 0xF133, 0xF155, 0xF165])

    def test_errors(self):
        from beamboot.chip8_asm import AsmError
        for bad in ("JP nowhere", "LD V0,300", "FOO V1", "a: CLS\na: CLS"):
            with self.assertRaises(AsmError):
                assemble(bad)


SELF_CHECK = """
; self-written conformance ROM: writes 1 to RAM[0xEFE] only if every check passes
        LD V0, 0xFF
        ADD V0, V0          ; 0xFE with carry
        SE VF, 1
        JP fail
        SE V0, 0xFE
        JP fail
        LD V1, 3
        LD V2, 5
        SUB V1, V2          ; borrow: 254, VF 0
        SE VF, 0
        JP fail
        SE V1, 254
        JP fail
        LD V3, 123
        LD I, 0xE00
        LD B, V3
        LD V2, [I]          ; V0..V2 = 1,2,3 (I += 3)
        SE V0, 1
        JP fail
        SE V1, 2
        JP fail
        SE V2, 3
        JP fail
        LD V4, 6
        LD V5, 7
        LD I, 0xE10
        LD [I], V5          ; store V0..V5
        LD I, 0xE10
        LD V5, [I]
        SE V4, 6
        JP fail
        SE V5, 7
        JP fail
        LD V0, 1
        LD I, 0xEFE
        LD [I], V0
done:   JP done
fail:   JP fail
"""


class SelfWrittenConformanceTest(unittest.TestCase):
    def test_self_check_rom_sets_flag(self):
        m = machine(SELF_CHECK)
        m.run(200)
        self.assertEqual(m.mem[0xEFE], 1)
        self.assertEqual(m.pc, m.syms["done"])


def _load_play():
    spec = importlib.util.spec_from_file_location(
        "play_chip8", os.path.join(ROOT, "scripts", "play_chip8.py"))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


class SuperBeamTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.play = _load_play()
        cls.rom, cls.syms = cls.play.build()

    def test_committed_rom_matches_source_and_is_small(self):
        with open(os.path.join(ROOT, "roms", "superbeam.ch8"), "rb") as f:
            self.assertEqual(f.read(), self.rom)
        self.assertLess(len(self.rom), 1536)

    def test_scripted_inputs_reach_goal_within_budget(self):
        with open(os.path.join(ROOT, "roms", "superbeam_inputs.json")) as f:
            script = json.load(f)
        res = self.play.run_game(self.rom, self.syms, script, max_frames=WIN_FRAME_BUDGET)
        m = res["machine"]
        self.assertTrue(res["won"])
        self.assertEqual(m.mem[self.syms["WIN"]], 1)
        self.assertEqual(res["deaths"], 0)
        self.assertLessEqual(res["instructions"], WIN_INSTRUCTION_BUDGET)
        self.assertLess(max(res["work"][:-1]), self.play.IPF)   # game logic fits a frame
        # replay is bit-exact deterministic
        res2 = self.play.run_game(self.rom, self.syms, script, capture=True)
        res3 = self.play.run_game(self.rom, self.syms, script, capture=True)
        self.assertEqual(res2["frames"], res3["frames"])

    def test_without_input_player_does_not_win(self):
        res = self.play.run_game(self.rom, self.syms, [], max_frames=120)
        self.assertFalse(res["won"])

    def _frames(self, rom, per_frame):
        """Run per-frame key sets; yield the machine after each frame (checks SP == 0)."""
        m = Chip8()
        m.load(rom)
        for keys in per_frame:
            for k in (4, 5, 6):
                m.set_key(k, k in keys)
            m.run(self.play.IPF)
            m.tick_timers()
            self.assertEqual(m.sp, 0)            # every frame boundary is outside any subroutine
            yield m

    def _level(self, cols, walkers=()):
        from beamboot import superbeam_levels as L
        lv = L._flat(cols)
        return L.build_rom(lv, walkers)

    def test_holding_right_is_stopped_by_the_first_pipe(self):
        res = self.play.run_game(self.rom, self.syms, [[0, 6, "down"]], max_frames=120)
        m = res["machine"]
        self.assertEqual(res["deaths"], 0)
        self.assertFalse(res["won"])
        self.assertEqual(m.v[0xB], 18)                       # still on the ground ...
        self.assertLess(m.v[0xD] * 8 + m.v[0xA] + 6, 9 * 8)  # ... left of the pipe at column 9

    def test_walking_into_a_pit_kills(self):
        rom = self._level({12: 0, 13: 0})
        deaths = [m.mem[0xEFD] for m in self._frames(rom, [{6}] * 80)]
        self.assertGreaterEqual(max(deaths), 1)

    def test_pipe_is_solid_and_can_be_stood_on(self):
        from beamboot import superbeam_levels as L
        rom = self._level({12: L.PIPE})
        ms = list(self._frames(rom, [{6}] * 40))
        self.assertEqual(ms[-1].v[0xB], 18)
        self.assertLessEqual(ms[-1].v[0xD] * 8 + ms[-1].v[0xA] + 6, 12 * 8 - 1)   # blocked by its side
        # jump while pushing right: lands on top (y = 10), then walks off the far side
        ys = [m.v[0xB] for m in self._frames(rom, [{6}] * 40 + [{5, 6}] + [{6}] * 40)]
        self.assertIn(10, ys[41:])
        self.assertEqual(ys[-1], 18)

    def test_tall_pipe_needs_a_step(self):
        from beamboot import superbeam_levels as L
        self.assertIsNone(L.solve(self._level({20: L.TALL}), max_nodes=600))
        self.assertIsNotNone(L.solve(self._level({20: L.PIPE, 21: L.TALL})))

    def test_coins_count_and_disappear(self):
        from beamboot import superbeam_levels as L
        rom = self._level({3: L.GROUND | L.COIN_LOW, 4: L.GROUND | L.COIN_LOW,
                           6: L.GROUND | L.COIN_MID})
        ms = list(self._frames(rom, [{6}] * 30))
        m = ms[-1]
        self.assertEqual(m.mem[0xEFC], 2)                     # low coins: walk into them
        self.assertEqual(L.level_from_ram(m.mem)[3:7], [1, 1, 1, L.GROUND | L.COIN_MID])
        self.assertEqual(m.display[30][56:64], [1, 0, 1, 0, 0, 0, 0, 0])   # coin bar: 2 dots
        # the mid coin needs a jump
        plan = [{6}] * 16 + [{5, 6}] + [{6}] * 20             # jump at x = 40: apex over column 6
        self.assertEqual(list(self._frames(rom, plan))[-1].mem[0xEFC], 3)

    def test_walkers_turn_at_pipes_and_pits(self):
        from beamboot import superbeam_levels as L
        rom = self._level({5: L.PIPE, 9: L.PIT}, walkers=[(1, 7, 0)])
        et = self.syms["ETAB"]
        xs = [m.mem[et + 1] * 8 + m.mem[et + 2] for m in self._frames(rom, [set()] * 80)]
        self.assertEqual(min(xs), 6 * 8)                      # turned at the pipe (column 5)
        self.assertEqual(max(xs), 8 * 8)                      # turned at the pit (column 9)

    def test_rom_uses_only_modelled_opcodes(self):
        code = check_subset(self.rom, (0x200, self.syms["halt"]))
        self.assertGreater(len(code), 300)
        self.assertTrue({op_class(o) for o in code.values()} <= set(SUPPORTED))
        self.assertFalse(any(o >> 12 == 2 for o in code.values()))   # no CALL: SP is always 0
        with self.assertRaises(AsmError):                          # the check does fire
            check_subset(assemble("RND V0, 1\nJP 0x200")[0])

    def test_static_worst_case_frame_fits_ipf(self):
        # longest static path from the DT-wait exit back to the wait loop (death/reset and win
        # paths included); 4 = worst spin-loop latency after the timer tick
        bound = longest_path(self.rom, self.syms["wait"] + 6, {self.syms["wait"]})
        self.assertLessEqual(bound + 4, self.play.IPF)
        self.assertLessEqual(longest_path(self.rom, 0x200, {self.syms["wait"]}) + 4, self.play.IPF)


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