from __future__ import annotations

from pathlib import Path

import numpy as np
import wgpu

SHADER_PATH = Path(__file__).parent / "shaders" / "textured_quad.wgsl"

# position(vec2) + uv(vec2), one unit quad shared by every instance
QUAD_VERTICES = np.array(
    [
        [0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 1.0, 0.0],
        [0.0, 1.0, 0.0, 1.0],
        [1.0, 1.0, 1.0, 1.0],
    ],
    dtype=np.float32,
)
QUAD_INDICES = np.array([0, 1, 2, 2, 1, 3], dtype=np.uint16)

CAMERA_UNIFORM_SIZE = 16  # scale(vec2<f32>) + offset(vec2<f32>)


class QuadPipeline:
    def __init__(self, device, texture_format: str):
        self.device = device

        shader = device.create_shader_module(code=SHADER_PATH.read_text())

        self.camera_buffer = device.create_buffer(
            size=CAMERA_UNIFORM_SIZE,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST,
        )
        self.camera_bind_group_layout = device.create_bind_group_layout(
            entries=[
                {
                    "binding": 0,
                    "visibility": wgpu.ShaderStage.VERTEX,
                    "buffer": {"type": wgpu.BufferBindingType.uniform},
                },
            ]
        )
        self.camera_bind_group = self._make_camera_bind_group(self.camera_buffer)

        self.texture_bind_group_layout = device.create_bind_group_layout(
            entries=[
                {
                    "binding": 0,
                    "visibility": wgpu.ShaderStage.FRAGMENT,
                    "texture": {
                        "sample_type": wgpu.TextureSampleType.float,
                        "view_dimension": wgpu.TextureViewDimension.d2,
                    },
                },
                {
                    "binding": 1,
                    "visibility": wgpu.ShaderStage.FRAGMENT,
                    "sampler": {"type": wgpu.SamplerBindingType.filtering},
                },
            ]
        )

        pipeline_layout = device.create_pipeline_layout(
            bind_group_layouts=[self.camera_bind_group_layout, self.texture_bind_group_layout]
        )

        self.pipeline = device.create_render_pipeline(
            layout=pipeline_layout,
            vertex={
                "module": shader,
                "entry_point": "vs_main",
                "buffers": [
                    {
                        "array_stride": 4 * 4,
                        "step_mode": wgpu.VertexStepMode.vertex,
                        "attributes": [
                            {"format": wgpu.VertexFormat.float32x2, "offset": 0, "shader_location": 0},
                            {"format": wgpu.VertexFormat.float32x2, "offset": 8, "shader_location": 1},
                        ],
                    },
                    {
                        "array_stride": 4 * 4,
                        "step_mode": wgpu.VertexStepMode.instance,
                        "attributes": [
                            {"format": wgpu.VertexFormat.float32x2, "offset": 0, "shader_location": 2},
                            {"format": wgpu.VertexFormat.float32x2, "offset": 8, "shader_location": 3},
                        ],
                    },
                ],
            },
            primitive={"topology": wgpu.PrimitiveTopology.triangle_list, "cull_mode": wgpu.CullMode.none},
            fragment={
                "module": shader,
                "entry_point": "fs_main",
                "targets": [
                    {
                        "format": texture_format,
                        "blend": {
                            "color": {
                                "src_factor": wgpu.BlendFactor.src_alpha,
                                "dst_factor": wgpu.BlendFactor.one_minus_src_alpha,
                                "operation": wgpu.BlendOperation.add,
                            },
                            "alpha": {
                                "src_factor": wgpu.BlendFactor.one,
                                "dst_factor": wgpu.BlendFactor.one_minus_src_alpha,
                                "operation": wgpu.BlendOperation.add,
                            },
                        },
                    }
                ],
            },
        )

        self.quad_vertex_buffer = device.create_buffer(
            size=QUAD_VERTICES.nbytes, usage=wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST
        )
        device.queue.write_buffer(self.quad_vertex_buffer, 0, QUAD_VERTICES.tobytes())

        self.quad_index_buffer = device.create_buffer(
            size=QUAD_INDICES.nbytes, usage=wgpu.BufferUsage.INDEX | wgpu.BufferUsage.COPY_DST
        )
        device.queue.write_buffer(self.quad_index_buffer, 0, QUAD_INDICES.tobytes())
        self.quad_index_count = len(QUAD_INDICES)

        self.sampler = device.create_sampler(mag_filter=wgpu.FilterMode.nearest, min_filter=wgpu.FilterMode.nearest)

    def make_texture_bind_group(self, texture_view):
        return self.device.create_bind_group(
            layout=self.texture_bind_group_layout,
            entries=[
                {"binding": 0, "resource": texture_view},
                {"binding": 1, "resource": self.sampler},
            ],
        )

    def _make_camera_bind_group(self, buffer):
        return self.device.create_bind_group(
            layout=self.camera_bind_group_layout,
            entries=[{"binding": 0, "resource": {"buffer": buffer, "offset": 0, "size": CAMERA_UNIFORM_SIZE}}],
        )

    def create_camera_binding(self):
        """A second, independent (buffer, bind_group) pair sharing this pipeline's camera
        layout - e.g. Renderer's screen-space UI camera, which must never share a buffer with
        the world camera: both get written once per frame, and a render pass only sees whatever
        a uniform buffer holds at submit time, so two live transforms need two buffers.
        """
        buffer = self.device.create_buffer(
            size=CAMERA_UNIFORM_SIZE, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST
        )
        return buffer, self._make_camera_bind_group(buffer)

    def write_camera(self, buffer, scale: tuple[float, float], offset: tuple[float, float]) -> None:
        data = np.array([scale[0], scale[1], offset[0], offset[1]], dtype=np.float32)
        self.device.queue.write_buffer(buffer, 0, data.tobytes())

    def update_camera(self, scale: tuple[float, float], offset: tuple[float, float]) -> None:
        self.write_camera(self.camera_buffer, scale, offset)

    def make_instance_buffer(self, instances: np.ndarray):
        buffer = self.device.create_buffer(
            size=instances.nbytes, usage=wgpu.BufferUsage.VERTEX | wgpu.BufferUsage.COPY_DST
        )
        self.device.queue.write_buffer(buffer, 0, instances.tobytes())
        return buffer