"""GPU buffer adapters used by ModernGL renders.

These classes only translate numpy array layouts to GPU buffers and back. They
do not own scene semantics, camera state, or frontend behavior.
"""

import numpy as np
import moderngl as mgl

from ..core import System


class Field3DBuffer:
    """Upload a vector field with one 3-float record per lattice site."""

    def __init__(self, mglctx: mgl.Context = None, system: System = None):
        self._ctx = mglctx
        self._system = system
        self.vbo = self._ctx.buffer(reserve=12 * self._system.number_of_spins)

    def upload(self, state: np.ndarray):
        """Write a ``(..., 3)`` state array into the underlying VBO."""
        assert isinstance(state, np.ndarray)
        self.vbo.write(state.reshape((-1, 3)).astype("f4").tobytes())

    def release(self) -> None:
        """Release the underlying ModernGL buffer."""
        vbo = getattr(self, "vbo", None)
        if vbo is None:
            return
        try:
            vbo.release()
        except Exception:
            pass
        self.vbo = None


class Field1DBuffer:
    """Upload one float per lattice site."""

    def __init__(self, mglctx: mgl.Context = None, system: System = None):
        self._ctx = mglctx
        self._system = system
        self.vbo = self._ctx.buffer(reserve=4 * self._system.number_of_spins)

    def upload(self, state: np.ndarray):
        """Write a scalar field into the underlying VBO."""
        assert isinstance(state, np.ndarray)
        self.vbo.write(state.flatten().astype("f4").tobytes())

    def release(self) -> None:
        """Release the underlying ModernGL buffer."""
        vbo = getattr(self, "vbo", None)
        if vbo is None:
            return
        try:
            vbo.release()
        except Exception:
            pass
        self.vbo = None


class Float3DBuffer:
    """Fixed-size buffer for arbitrary 3-float records."""

    def __init__(self, mglctx: mgl.Context = None, count: int = None):
        self._ctx = mglctx
        self._count = count
        self.vbo = self._ctx.buffer(reserve=12 * self._count)

    def upload(self, data: np.ndarray):
        """Write exactly ``count`` 3-float records."""
        assert isinstance(data, np.ndarray)
        data = data.reshape((-1, 3))
        assert data.shape[0] == self._count
        self.vbo.write(data.astype("f4").tobytes())

    def download(self):
        """Read the buffer back as a ``(count, 3)`` float array."""
        return np.frombuffer(self.vbo.read(), dtype=np.float32).reshape(
            (-1, 3)
        )  # , count=self._count)

    def release(self) -> None:
        """Release the underlying ModernGL buffer."""
        vbo = getattr(self, "vbo", None)
        if vbo is None:
            return
        try:
            vbo.release()
        except Exception:
            pass
        self.vbo = None


class IntBuffer:
    """Fixed-size integer index or mask buffer."""

    def __init__(self, mglctx: mgl.Context = None, count=None):
        self._ctx = mglctx
        self._count = count
        self.ibo = self._ctx.buffer(reserve=4 * self._count)

    @property
    def count(self):
        return self._count

    def upload(self, data: np.ndarray):
        """Write exactly ``count`` int32 values."""
        assert isinstance(data, np.ndarray)
        data = data.flatten()
        assert data.shape == (self._count,)
        self.ibo.write(data.astype("i4").tobytes())

    def download(self):
        """Read the buffer back as an ``int32`` array."""
        return np.frombuffer(self.ibo.read(), dtype=np.int32)  # , count=self._count)

    def release(self) -> None:
        """Release the underlying ModernGL buffer."""
        ibo = getattr(self, "ibo", None)
        if ibo is None:
            return
        try:
            ibo.release()
        except Exception:
            pass
        self.ibo = None
