"""
SoundBlock GPU Kernel — WebGPU-accelerated mantissa packing.

Uses the same wgpu-py backend as the colony force kernel.
Runs mantissa_pack.wgsl compute shader on the GPU.

All 16 files × 17,067 bins = 273,072 parallel operations.
Each invocation packs 6 bytes into one bin. Zero interdependence.

Usage:
    kernel = SoundBlockGPUKernel()
    kernel.pack_all(files, container_re, container_im)
    kernel.unpack(file_name, container_re, container_im) -> bytes
"""

import numpy as np
import struct
import time
from pathlib import Path

try:
    import wgpu
    _WGPU_AVAILABLE = True
except ImportError:
    _WGPU_AVAILABLE = False


def gpu_available() -> bool:
    if not _WGPU_AVAILABLE:
        return False
    try:
        device = wgpu.utils.get_default_device()
        return device is not None
    except Exception:
        return False


_SHADER_DIR = Path(__file__).parent


class SoundBlockGPUKernel:
    """GPU-accelerated mantissa packing for Sound Block storage."""

    def __init__(self):
        if not _WGPU_AVAILABLE:
            raise ImportError("wgpu package required: pip install wgpu")

        self._device = wgpu.utils.get_default_device()
        self._shader_src = (_SHADER_DIR / "mantissa_pack.wgsl").read_text()
        self._create_pipeline()

    def _create_pipeline(self):
        d = self._device
        module = d.create_shader_module(code=self._shader_src)

        self._bgl = d.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
            {"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
            {"binding": 3, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])

        pl = d.create_pipeline_layout(bind_group_layouts=[self._bgl])
        self._pipeline = d.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "main"})

    def pack(self, data: bytes, start_bin: int, re_buf, im_buf):
        """
        Pack data bytes into frequency bins on the GPU.

        Args:
            data: raw bytes to pack
            start_bin: starting frequency bin index
            re_buf: GPU buffer for real part (storage)
            im_buf: GPU buffer for imaginary part (storage)

        Returns:
            (elapsed_ms, bins_written)
        """
        d = self._device
        byte_count = len(data)
        bin_count = (byte_count + 5) // 6  # ceil(byte_count / 6)

        # Pad data to multiple of 4 for u32 packing
        padded_len = ((byte_count + 3) // 4) * 4
        padded_data = data + b'\x00' * (padded_len - byte_count)

        # Pack data as big-endian u32 array
        data_u32 = np.frombuffer(padded_data, dtype='>u4').astype(np.uint32)

        # Create params uniform: bin_count, start_bin, byte_count, direction=0
        params = struct.pack('IIII', bin_count, start_bin, byte_count, 0)

        # Create GPU buffers
        params_buf = d.create_buffer_with_data(data=params,
            usage=wgpu.BufferUsage.UNIFORM)
        data_buf = d.create_buffer_with_data(data=data_u32.tobytes(),
            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC)

        # Bind group
        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": params_buf}},
            {"binding": 1, "resource": {"buffer": re_buf}},
            {"binding": 2, "resource": {"buffer": im_buf}},
            {"binding": 3, "resource": {"buffer": data_buf}},
        ])

        # Dispatch
        workgroups = (bin_count + 255) // 256
        t0 = time.perf_counter()
        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._pipeline)
        compute_pass.set_bind_group(0, bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])

        # Wait for GPU to finish
        d._poll_wait()
        elapsed_ms = (time.perf_counter() - t0) * 1000

        return elapsed_ms, bin_count

    def unpack(self, start_bin: int, byte_count: int, re_buf, im_buf) -> tuple:
        """
        Unpack data bytes from frequency bins on the GPU.

        Returns:
            (data: bytes, elapsed_ms)
        """
        d = self._device
        bin_count = (byte_count + 5) // 6

        # Allocate output data buffer
        padded_len = ((byte_count + 3) // 4) * 4
        data_u32 = np.zeros(padded_len // 4, dtype=np.uint32)

        params = struct.pack('IIII', bin_count, start_bin, byte_count, 1)
        params_buf = d.create_buffer_with_data(data=params,
            usage=wgpu.BufferUsage.UNIFORM)
        data_buf = d.create_buffer_with_data(data=data_u32.tobytes(),
            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC)

        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": params_buf}},
            {"binding": 1, "resource": {"buffer": re_buf}},
            {"binding": 2, "resource": {"buffer": im_buf}},
            {"binding": 3, "resource": {"buffer": data_buf}},
        ])

        workgroups = (bin_count + 255) // 256
        t0 = time.perf_counter()
        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._pipeline)
        compute_pass.set_bind_group(0, bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])
        d._poll_wait()
        elapsed_ms = (time.perf_counter() - t0) * 1000

        # Read back
        result = d.queue.read_buffer(data_buf)
        result_u32 = np.frombuffer(result, dtype=np.uint32)
        # Convert big-endian u32 back to bytes
        result_bytes = result_u32.astype('>u4').tobytes()[:byte_count]

        return result_bytes, elapsed_ms

    def create_container_buffers(self, fft_size: int):
        """Create GPU-resident re and im buffers for a Sound Block container."""
        d = self._device
        buf_size = fft_size * 4  # f32 = 4 bytes (GPU uses f32, not f64)

        re_buf = d.create_buffer(size=buf_size,
            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST)
        im_buf = d.create_buffer(size=buf_size,
            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST)

        return re_buf, im_buf, fft_size

    def pack_all(self, files: list, fft_size: int) -> dict:
        """
        Pack all files into one GPU-resident Sound Block container.

        Args:
            files: list of (name, data_bytes) tuples
            fft_size: size of the FFT container

        Returns:
            dict with timing, allocations, container buffers
        """
        re_buf, im_buf, _ = self.create_container_buffers(fft_size)
        allocations = []
        next_bin = 1
        total_bytes = 0
        timings = []

        t0 = time.perf_counter()

        for name, data in files:
            bins_needed = (len(data) + 5) // 6
            elapsed_ms, bins_written = self.pack(data, next_bin, re_buf, im_buf)
            alloc = {
                'name': name,
                'start_bin': next_bin,
                'byte_count': len(data),
                'bins_used': bins_written,
                'end_bin': next_bin + bins_written - 1,
                'encode_ms': elapsed_ms,
            }
            allocations.append(alloc)
            timings.append(elapsed_ms)
            next_bin += bins_written
            total_bytes += len(data)

        total_ms = (time.perf_counter() - t0) * 1000

        return {
            'allocations': allocations,
            'total_ms': total_ms,
            'total_bytes': total_bytes,
            'timings': timings,
            're_buf': re_buf,
            'im_buf': im_buf,
            'fft_size': fft_size,
            'throughput_mbps': (total_bytes / (1024 * 1024)) / (total_ms / 1000) if total_ms > 0 else 0,
        }

    def unpack_all(self, allocations: list, re_buf, im_buf) -> dict:
        """
        Unpack all files from a GPU-resident container.

        Returns:
            dict with recovered data, timings
        """
        results = []
        t0 = time.perf_counter()

        for alloc in allocations:
            data, elapsed_ms = self.unpack(
                alloc['start_bin'], alloc['byte_count'], re_buf, im_buf)
            results.append({
                'name': alloc['name'],
                'data': data,
                'decode_ms': elapsed_ms,
            })

        total_ms = (time.perf_counter() - t0) * 1000
        total_bytes = sum(len(r['data']) for r in results)

        return {
            'results': results,
            'total_ms': total_ms,
            'total_bytes': total_bytes,
            'throughput_mbps': (total_bytes / (1024 * 1024)) / (total_ms / 1000) if total_ms > 0 else 0,
        }
