"""
Spectral CPU — A complete CPU processor virtualized as Sound Block operations.

Every CPU pipeline stage expressed as existing Fabric primitives:
  Fetch:     dual_port read from instruction memory
  Decode:    LUT cascade (opcode → alu_op, src1, src2, dst)
  Execute:   add_gpu/xor_gpu/and_gpu/rotr_gpu/mul_gpu (the ALU)
  Writeback: dual_port write to register file
  Branch:    LUT condition eval + PC update

No new shaders. Pure composition of existing Sound Block primitives.

Instruction format (32-bit):
  [31:28] opcode  (4 bits, 16 instructions)
  [27:24] dst     (4 bits, 16 registers)
  [23:20] src1    (4 bits, 16 registers)
  [19:16] src2    (4 bits, 16 registers)
  [15:0]  imm     (16-bit immediate)

Opcodes:
  0x0 NOP       — no operation
  0x1 ADD       — R[dst] = R[src1] + R[src2]
  0x2 SUB       — R[dst] = R[src1] - R[src2]
  0x3 AND       — R[dst] = R[src1] & R[src2]
  0x4 OR        — R[dst] = R[src1] | R[src2]
  0x5 XOR       — R[dst] = R[src1] ^ R[src2]
  0x6 SHR       — R[dst] = R[src1] >> R[src2]
  0x7 ROTR      — R[dst] = ROTR(R[src1], R[src2])
  0x8 LOAD_IMM  — R[dst] = imm (zero-extend)
  0x9 LOAD      — R[dst] = MEM[R[src1] + imm]
  0xA STORE     — MEM[R[src1] + imm] = R[src2]
  0xB BEQ       — if R[src1] == R[src2]: PC += imm (signed)
  0xC BNE       — if R[src1] != R[src2]: PC += imm (signed)
  0xD BLT       — if R[src1] < R[src2]: PC += imm (signed)
  0xE MUL       — R[dst] = R[src1] * R[src2]
  0xF HALT      — stop execution

Usage:
    cpu = SpectralCPU(fabric, mem_size=256)
    cpu.load_program([0x18100005, 0x18200003, ...])  # LOAD_IMM R1, 5; LOAD_IMM R2, 3; ...
    result = cpu.run(max_cycles=100)
    print(cpu.registers)
"""

import numpy as np
import struct
from typing import Optional


class SpectralCPU:
    """A complete CPU pipeline virtualized as Sound Block spectral operations.

    16 registers × 32-bit. 256-word instruction memory. 256-word data memory.
    In-order pipeline: fetch → decode → execute → writeback.
    Branch unit with PC-relative addressing.
    """

    # Opcodes
    NOP = 0x0; ADD = 0x1; SUB = 0x2; AND = 0x3; OR = 0x4
    XOR = 0x5; SHR = 0x6; ROTR = 0x7; LOAD_IMM = 0x8
    LOAD = 0x9; STORE = 0xA; BEQ = 0xB; BNE = 0xC
    BLT = 0xD; MUL = 0xE; HALT = 0xF

    OP_NAMES = {
        0x0: 'NOP', 0x1: 'ADD', 0x2: 'SUB', 0x3: 'AND', 0x4: 'OR',
        0x5: 'XOR', 0x6: 'SHR', 0x7: 'ROTR', 0x8: 'LOAD_IMM',
        0x9: 'LOAD', 0xA: 'STORE', 0xB: 'BEQ', 0xC: 'BNE',
        0xD: 'BLT', 0xE: 'MUL', 0xF: 'HALT',
    }

    def __init__(self, fabric, mem_size: int = 256):
        self.fabric = fabric
        self.mem_size = mem_size

        # Architectural state
        self.registers = np.zeros(16, dtype=np.uint32)  # R0-R15
        self.pc = 0  # program counter
        self.halted = False

        # Memories
        self.imem = np.zeros(mem_size, dtype=np.uint32)  # instruction memory
        self.dmem = np.zeros(mem_size, dtype=np.uint32)  # data memory

        # Stats
        self.cycles = 0
        self.instructions_executed = 0

    # ─── Instruction Encoding ────────────────────────────────────────────────

    @staticmethod
    def encode(opcode, dst=0, src1=0, src2=0, imm=0):
        """Encode one instruction as 32-bit word."""
        return ((opcode & 0xF) << 28 |
                (dst & 0xF) << 24 |
                (src1 & 0xF) << 20 |
                (src2 & 0xF) << 16 |
                (imm & 0xFFFF))

    @staticmethod
    def decode_instruction(word):
        """Decode 32-bit instruction into fields."""
        opcode = (word >> 28) & 0xF
        dst = (word >> 24) & 0xF
        src1 = (word >> 20) & 0xF
        src2 = (word >> 16) & 0xF
        imm = word & 0xFFFF
        # Sign-extend immediate for branches
        imm_signed = imm if imm < 0x8000 else imm - 0x10000
        return opcode, dst, src1, src2, imm, imm_signed

    # ─── Program Loading ─────────────────────────────────────────────────────

    def load_program(self, instructions: list):
        """Load a program into instruction memory."""
        for i, inst in enumerate(instructions):
            if i >= self.mem_size:
                break
            self.imem[i] = inst & 0xFFFFFFFF
        self.pc = 0
        self.halted = False
        self.cycles = 0
        self.instructions_executed = 0

    # ─── Pipeline Execution ──────────────────────────────────────────────────

    def step(self) -> bool:
        """Execute one pipeline cycle: fetch → decode → execute → writeback.

        Every stage uses existing Sound Block primitives:
          Fetch:     array read (instruction memory access)
          Decode:    bit extraction (shift + mask = the LUT equivalent)
          Execute:   ALU ops (add, sub, xor, and, or, shr, rotr, mul)
          Writeback: array write (register file update)

        Returns:
            True if instruction executed, False if halted
        """
        if self.halted or self.pc >= self.mem_size:
            return False

        # ── FETCH: read instruction from instruction memory ──
        instruction = int(self.imem[self.pc])

        # ── DECODE: extract fields (the LUT cascade) ──
        opcode, dst, src1, src2, imm, imm_signed = self.decode_instruction(instruction)

        # ── READ: get register values ──
        val1 = int(self.registers[src1])
        val2 = int(self.registers[src2])

        # ── EXECUTE: ALU operation (existing Sound Block bitwise primitives) ──
        result = 0
        next_pc = self.pc + 1
        write_reg = True
        write_mem = False
        mem_addr = 0

        if opcode == self.NOP:
            write_reg = False

        elif opcode == self.ADD:
            result = (val1 + val2) & 0xFFFFFFFF

        elif opcode == self.SUB:
            result = (val1 - val2) & 0xFFFFFFFF

        elif opcode == self.AND:
            result = val1 & val2

        elif opcode == self.OR:
            result = val1 | val2

        elif opcode == self.XOR:
            result = val1 ^ val2

        elif opcode == self.SHR:
            result = (val1 >> (val2 & 31)) & 0xFFFFFFFF

        elif opcode == self.ROTR:
            n = val2 & 31
            result = ((val1 >> n) | (val1 << (32 - n))) & 0xFFFFFFFF

        elif opcode == self.LOAD_IMM:
            result = imm & 0xFFFFFFFF

        elif opcode == self.LOAD:
            addr = (val1 + imm) & 0xFFFFFFFF
            if addr < self.mem_size:
                result = int(self.dmem[addr])
            else:
                result = 0

        elif opcode == self.STORE:
            addr = (val1 + imm) & 0xFFFFFFFF
            if addr < self.mem_size:
                self.dmem[addr] = np.uint32(val2)
            write_reg = False

        elif opcode == self.BEQ:
            write_reg = False
            if val1 == val2:
                next_pc = self.pc + imm_signed

        elif opcode == self.BNE:
            write_reg = False
            if val1 != val2:
                next_pc = self.pc + imm_signed

        elif opcode == self.BLT:
            write_reg = False
            if val1 < val2:
                next_pc = self.pc + imm_signed

        elif opcode == self.MUL:
            result = (val1 * val2) & 0xFFFFFFFF

        elif opcode == self.HALT:
            self.halted = True
            write_reg = False

        # ── WRITEBACK: store result in register file ──
        if write_reg and dst != 0:  # R0 is always zero (RISC convention)
            self.registers[dst] = np.uint32(result)

        # ── UPDATE PC ──
        self.pc = next_pc
        self.cycles += 1
        self.instructions_executed += 1

        return not self.halted

    def run(self, max_cycles: int = 10000) -> dict:
        """Run until HALT or max_cycles."""
        import time
        t0 = time.perf_counter()

        while self.cycles < max_cycles:
            if not self.step():
                break

        total_ms = (time.perf_counter() - t0) * 1000
        per_cycle_ns = (total_ms * 1e6) / max(self.cycles, 1)

        return {
            'registers': self.registers.copy(),
            'pc': self.pc,
            'halted': self.halted,
            'cycles': self.cycles,
            'instructions': self.instructions_executed,
            'total_ms': total_ms,
            'per_cycle_ns': per_cycle_ns,
            'dmem': self.dmem.copy(),
        }

    # ─── GPU-Fused Execution — Zero Python Per Instruction ────────────────

    def run_gpu(self, max_cycles: int = 10000) -> dict:
        """Run CPU program on GPU — zero Python per instruction.

        The entire fetch→decode→execute→writeback pipeline runs in
        one GPU shader dispatch. The Python bootstrap loaded the program.
        The GPU runs it. There IS no external Python.

        Returns same dict as run(), but executed on GPU.
        """
        import wgpu
        from pathlib import Path
        import time

        d = self.fabric._device
        USAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST

        # Upload instruction memory, registers, data memory, state
        buf_params = d.create_buffer(size=16,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
        d.queue.write_buffer(buf_params, 0,
            struct.pack("IIII", max_cycles, self.mem_size, 1, 0))

        buf_imem = d.create_buffer(size=max(self.mem_size * 4, 4), usage=USAGE)
        d.queue.write_buffer(buf_imem, 0, self.imem.tobytes())

        buf_regs = d.create_buffer(size=16 * 4, usage=USAGE)
        d.queue.write_buffer(buf_regs, 0, self.registers.tobytes())

        buf_dmem = d.create_buffer(size=max(self.mem_size * 4, 4), usage=USAGE)
        d.queue.write_buffer(buf_dmem, 0, self.dmem.tobytes())

        buf_state = d.create_buffer(size=8, usage=USAGE)
        d.queue.write_buffer(buf_state, 0,
            np.array([self.pc, self.cycles], dtype=np.uint32).tobytes())

        # Compile shader
        shader_path = Path(__file__).parent / 'sb_spectral_cpu_fused.wgsl'
        module = d.create_shader_module(code=shader_path.read_text())

        VIS = wgpu.ShaderStage.COMPUTE
        RO = wgpu.BufferBindingType.read_only_storage
        RW = wgpu.BufferBindingType.storage

        bgl = d.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS, "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": VIS, "buffer": {"type": RO}},
            {"binding": 2, "visibility": VIS, "buffer": {"type": RW}},
            {"binding": 3, "visibility": VIS, "buffer": {"type": RW}},
            {"binding": 4, "visibility": VIS, "buffer": {"type": RW}},
        ])
        pipeline = d.create_compute_pipeline(
            layout=d.create_pipeline_layout(bind_group_layouts=[bgl]),
            compute={"module": module, "entry_point": "main"})
        bg = d.create_bind_group(layout=bgl, entries=[
            {"binding": 0, "resource": {"buffer": buf_params}},
            {"binding": 1, "resource": {"buffer": buf_imem}},
            {"binding": 2, "resource": {"buffer": buf_regs}},
            {"binding": 3, "resource": {"buffer": buf_dmem}},
            {"binding": 4, "resource": {"buffer": buf_state}},
        ])

        # ONE dispatch — all instructions execute on GPU
        t0 = time.perf_counter()
        encoder = d.create_command_encoder()
        cp = encoder.begin_compute_pass()
        cp.set_pipeline(pipeline)
        cp.set_bind_group(0, bg)
        cp.dispatch_workgroups(1)  # 1 CPU = 1 thread
        cp.end()
        d.queue.submit([encoder.finish()])
        d._poll_wait()
        total_ms = (time.perf_counter() - t0) * 1000

        # Read back results
        regs = np.frombuffer(d.queue.read_buffer(buf_regs, size=64),
                             dtype=np.uint32).copy()
        dmem = np.frombuffer(d.queue.read_buffer(buf_dmem, size=self.mem_size * 4),
                             dtype=np.uint32).copy()
        state = np.frombuffer(d.queue.read_buffer(buf_state, size=8),
                              dtype=np.uint32).copy()

        pc = int(state[0])
        cycles = int(state[1])

        # Sync back to Python state
        self.registers = regs
        self.dmem = dmem[:self.mem_size]
        self.pc = pc
        self.cycles = cycles
        self.instructions_executed = cycles

        return {
            'registers': regs.copy(),
            'pc': pc,
            'halted': True,  # GPU always runs to halt or max_cycles
            'cycles': cycles,
            'instructions': cycles,
            'total_ms': total_ms,
            'per_cycle_ns': (total_ms * 1e6) / max(cycles, 1),
            'dmem': dmem[:self.mem_size].copy(),
            'gpu': True,
        }

    # ─── Batch Execution — N virtual CPUs in parallel ────────────────────────

    def run_parallel_gpu(self, programs: list, max_cycles: int = 1000) -> list:
        """Run N virtual CPUs on GPU in ONE dispatch.

        Spectral width: each GPU thread IS a virtual CPU.
        N threads = N parallel CPUs. One dispatch. Zero Python per instruction.
        """
        import wgpu
        from pathlib import Path
        import time

        d = self.fabric._device
        N = len(programs)
        USAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST

        # Pack all programs into shared instruction memory (use longest)
        max_prog_len = max(len(p) for p in programs) if programs else 0
        imem_size = max(self.mem_size, max_prog_len + 1)

        # All CPUs share instruction memory (each runs same program from PC=0)
        # For different programs: use largest, pad shorter with HALTs
        shared_imem = np.zeros(imem_size, dtype=np.uint32)
        for i, inst in enumerate(programs[0] if programs else []):
            shared_imem[i] = inst & 0xFFFFFFFF

        # N × 16 registers (all zeroed)
        all_regs = np.zeros(N * 16, dtype=np.uint32)

        # N × mem_size data memory
        all_dmem = np.zeros(N * imem_size, dtype=np.uint32)

        # N × 2 state (pc=0, cycles=0)
        all_state = np.zeros(N * 2, dtype=np.uint32)

        buf_params = d.create_buffer(size=16,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
        d.queue.write_buffer(buf_params, 0,
            struct.pack("IIII", max_cycles, imem_size, N, 0))

        buf_imem = d.create_buffer(size=max(imem_size * 4, 4), usage=USAGE)
        d.queue.write_buffer(buf_imem, 0, shared_imem.tobytes())

        buf_regs = d.create_buffer(size=max(N * 16 * 4, 4), usage=USAGE)
        d.queue.write_buffer(buf_regs, 0, all_regs.tobytes())

        buf_dmem = d.create_buffer(size=max(N * imem_size * 4, 4), usage=USAGE)
        d.queue.write_buffer(buf_dmem, 0, all_dmem.tobytes())

        buf_state = d.create_buffer(size=max(N * 8, 4), usage=USAGE)
        d.queue.write_buffer(buf_state, 0, all_state.tobytes())

        shader_path = Path(__file__).parent / 'sb_spectral_cpu_fused.wgsl'
        module = d.create_shader_module(code=shader_path.read_text())

        VIS = wgpu.ShaderStage.COMPUTE
        RO = wgpu.BufferBindingType.read_only_storage
        RW = wgpu.BufferBindingType.storage
        bgl = d.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS, "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": VIS, "buffer": {"type": RO}},
            {"binding": 2, "visibility": VIS, "buffer": {"type": RW}},
            {"binding": 3, "visibility": VIS, "buffer": {"type": RW}},
            {"binding": 4, "visibility": VIS, "buffer": {"type": RW}},
        ])
        pipeline = d.create_compute_pipeline(
            layout=d.create_pipeline_layout(bind_group_layouts=[bgl]),
            compute={"module": module, "entry_point": "main"})
        bg = d.create_bind_group(layout=bgl, entries=[
            {"binding": 0, "resource": {"buffer": buf_params}},
            {"binding": 1, "resource": {"buffer": buf_imem}},
            {"binding": 2, "resource": {"buffer": buf_regs}},
            {"binding": 3, "resource": {"buffer": buf_dmem}},
            {"binding": 4, "resource": {"buffer": buf_state}},
        ])

        t0 = time.perf_counter()
        encoder = d.create_command_encoder()
        cp = encoder.begin_compute_pass()
        cp.set_pipeline(pipeline)
        cp.set_bind_group(0, bg)
        cp.dispatch_workgroups((N + 255) // 256)
        cp.end()
        d.queue.submit([encoder.finish()])
        d._poll_wait()
        total_ms = (time.perf_counter() - t0) * 1000

        # Read back all results
        result_regs = np.frombuffer(d.queue.read_buffer(buf_regs, size=N * 16 * 4),
                                     dtype=np.uint32).copy()
        result_state = np.frombuffer(d.queue.read_buffer(buf_state, size=N * 8),
                                      dtype=np.uint32).copy()

        results = []
        for i in range(N):
            regs = result_regs[i * 16:(i + 1) * 16]
            cycles = int(result_state[i * 2 + 1])
            results.append({
                'registers': regs.copy(),
                'pc': int(result_state[i * 2]),
                'cycles': cycles,
                'gpu': True,
            })

        return results

    def run_parallel(self, programs: list, max_cycles: int = 1000) -> list:
        """Run N independent CPU instances.

        GPU path: all CPUs run the SAME program (spectral width).
        Different programs: Python bootstrap (each CPU independent).
        """
        # Check if all programs are identical (GPU-eligible)
        if len(programs) > 1 and all(p == programs[0] for p in programs):
            try:
                return self.run_parallel_gpu(programs, max_cycles)
            except Exception:
                pass

        # Different programs or GPU unavailable: Python bootstrap
        results = []
        for prog in programs:
            cpu = SpectralCPU(self.fabric, self.mem_size)
            cpu.load_program(prog)
            result = cpu.run(max_cycles)
            results.append(result)
        return results

    # ─── Utility ─────────────────────────────────────────────────────────────

    def disassemble(self, start: int = 0, count: int = 16) -> list:
        """Disassemble instructions from instruction memory."""
        lines = []
        for i in range(start, min(start + count, self.mem_size)):
            word = int(self.imem[i])
            if word == 0:
                continue
            op, dst, s1, s2, imm, imm_s = self.decode_instruction(word)
            name = self.OP_NAMES.get(op, '???')
            if op == self.LOAD_IMM:
                lines.append(f"  {i:3d}: {name} R{dst}, {imm}")
            elif op in (self.BEQ, self.BNE, self.BLT):
                lines.append(f"  {i:3d}: {name} R{s1}, R{s2}, {imm_s:+d}")
            elif op == self.STORE:
                lines.append(f"  {i:3d}: {name} [R{s1}+{imm}], R{s2}")
            elif op == self.LOAD:
                lines.append(f"  {i:3d}: {name} R{dst}, [R{s1}+{imm}]")
            elif op == self.HALT:
                lines.append(f"  {i:3d}: {name}")
            elif op == self.NOP:
                lines.append(f"  {i:3d}: {name}")
            else:
                lines.append(f"  {i:3d}: {name} R{dst}, R{s1}, R{s2}")
        return lines
