from __future__ import annotations

from pathlib import Path

import numpy as np
import wgpu
from PIL import Image


class TextureCache:
    def __init__(self, device, textures_dir: Path):
        self.device = device
        self.textures_dir = textures_dir
        self._views: dict[str, "wgpu.GPUTextureView"] = {}

    def get_view(self, relative_path: str) -> "wgpu.GPUTextureView":
        if relative_path not in self._views:
            self._views[relative_path] = self._load(relative_path)
        return self._views[relative_path]

    def _load(self, relative_path: str) -> "wgpu.GPUTextureView":
        image = Image.open(self.textures_dir / relative_path).convert("RGBA")
        data = np.asarray(image, dtype=np.uint8)
        width, height = image.width, image.height

        texture = self.device.create_texture(
            size=(width, height, 1),
            format=wgpu.TextureFormat.rgba8unorm,
            usage=wgpu.TextureUsage.TEXTURE_BINDING | wgpu.TextureUsage.COPY_DST,
        )
        self.device.queue.write_texture(
            {"texture": texture},
            data.tobytes(),
            {"bytes_per_row": width * 4, "rows_per_image": height},
            (width, height, 1),
        )
        return texture.create_view()