"""
A global object shared by all renderers.
"""

import os
import wgpu

from ....resources import Resource, Buffer
from ....utils.trackable import Trackable
from ....utils import array_from_shadertype
from ....utils.text import glyph_atlas

from .utils import gpu_caches


# Definition uniform struct with standard info related to transforms,
# provided to each shader as uniform at slot 0.
# todo: a combined transform would be nice too, for performance
# todo: same for ndc_to_world transform (combined inv transforms)
stdinfo_uniform_type = dict(
    cam_transform="4x4xf4",
    cam_transform_inv="4x4xf4",
    projection_transform="4x4xf4",
    projection_transform_inv="4x4xf4",
    ndc_offset="4xf4",
    physical_size="2xf4",
    logical_size="2xf4",
)


def _preconfigure_wgpu_device():
    # PyGfx needs 'float32-filterable' because vanilla WGPU does not support
    # interpolating samplers for float32 textures, which is sad for e.g. volume
    # rendering. Technically we can check if the feature is available and fall back
    # (in some way) when it's not available, but this easily results in complex
    # code. I think we can probably get away with requiring only a few features that
    # are available on the main target platforms.
    wgpu.preconfigure_default_device("pygfx", required_features={"float32-filterable"})

    if adapter_name := os.environ.get("PYGFX_WGPU_ADAPTER_NAME"):
        # Similar to https://github.com/gfx-rs/wgpu?tab=readme-ov-file#environment-variables
        adapters = wgpu.gpu.enumerate_adapters_sync()
        selected_adapters = [a for a in adapters if adapter_name in a.summary]
        if not selected_adapters:
            raise ValueError(
                f"Adapter with name '{adapter_name}' (set by PYGFX_WGPU_ADAPTER_NAME env var) not found."
            )
        else:
            wgpu.preconfigure_default_device("pygfx", adapter=selected_adapters[0])


_preconfigure_wgpu_device()


class Shared(Trackable):
    """An object to store global data to share between multiple wgpu
    renderers. Each renderer updates the data and then passes this down
    to the pipeline containers.

    The renderer instantiates and stores the singleton shared object.
    """

    _instance = None

    def __init__(self, *, canvas=None):
        super().__init__()

        # Set this instance as the global one
        assert Shared._instance is None
        Shared._instance = self

        # Select device
        self._device = wgpu.get_default_device()
        self._adapter = self._device.adapter

        self._create_diagnostics()

        # Create a uniform buffer for std info
        # Stored on _store so if we'd ever swap it out for another buffer,
        # the pipeline automatically update.
        self._store.uniform_buffer = Buffer(
            array_from_shadertype(stdinfo_uniform_type), force_contiguous=True
        )
        self._store.uniform_buffer._wgpu_usage |= wgpu.BufferUsage.UNIFORM

        # Init glyph atlas texture
        self._store.glyph_atlas_texture = None
        self._store.glyph_atlas_info_buffer = None
        self.pre_render_hook()

    def pre_render_hook(self):
        """Called by the renderer on the beginning of a draw."""
        tex = glyph_atlas.texture
        if tex is not self._store["glyph_atlas_texture"]:
            self._store.glyph_atlas_texture = tex
        buffer = glyph_atlas.info_buffer
        if buffer is not self._store["glyph_atlas_info_buffer"]:
            self._store.glyph_atlas_info_buffer = buffer

    @classmethod
    def get_instance(cls):
        return cls._instance

    @property
    def adapter(self):
        """The shared WGPU adapter object."""
        return self._adapter

    @property
    def device(self):
        """The shared WGPU device object."""
        return self._device

    @property
    def uniform_buffer(self):
        """The shared uniform buffer in which the renderer puts
        information about the canvas and camera.
        """
        return self._store.uniform_buffer

    @property
    def glyph_atlas_texture(self):
        """The shared glyph atlas (as a texture view and sampler) for objects that want to render text."""
        return self._store.glyph_atlas_texture

    @property
    def glyph_atlas_info_buffer(self):
        """A buffer containing per-glyph metadata (rects and more)."""
        return self._store.glyph_atlas_info_buffer

    def _create_diagnostics(self):
        # Since diagnostics objects, are shown in order of creation,
        # its good to instantiate our diagnostics objects right after loading
        # the wgpu_native backend from wgpu.
        PyGfxAdapterInfoDiagnostics("pygfx_adapter_info")
        PyGfxFeaturesDiagnostics("pygfx_features")
        PyGfxLimitsDiagnostics("pygfx_limits")
        PygfxCacheDiagnostics("pygfx_caches")
        PygfxResourceDiagnostics("pygfx_resources")


def select_power_preference(power_preference):
    """Select whether a powerful or battery-friendly GPU is selected.

    This is an alias to ``wgpu.preconfigure_default_device(.., power_preference=...)``.

    Accepts a value from ``wgpu.PowerPreference``: "high-performance" or "low-power".
    This function must be called before before the first ``Renderer`` is created.
    """
    wgpu.preconfigure_default_device("pygfx", power_preference=power_preference)


def select_adapter(adapter):
    """Select a specific adapter / GPU.

    This is an alias to ``wgpu.preconfigure_default_device(.., adapter=...)``.

    Select an adapter as obtained via ``wgpu.gpu.enumerate_adapters_sync()``, which
    can be useful in multi-gpu environments.

    For example::

        adapters = wgpu.gpu.enumerate_adapters_sync()
        adapters_tesla = [a for a in adapters if "Tesla" in a.summary]
        adapters_discrete = [a for a in adapters if "DiscreteGPU" in a.summary]
        pygfx.renderers.wgpu.select_adapter(adapters_discrete[0])

    Note that using this function reduces the portability of your code, because
    it's highly specific for your current machine/environment.

    The order of the adapters returned by ``wgpu.gpu.enumerate_adapters_sync()`` is
    such that Vulkan adapters go first, then Metal, then D3D12, then OpenGL.
    Within each category, the order as provided by the particular backend is
    maintained. Note that the same device may be present via multiple backends
    (e.g. vulkan/opengl).

    We cannot make guarantees about whether the order of the adapters matches
    the order as reported by e.g. ``nvidia-smi``. We have found that on a Linux
    multi-gpu cluster, the order does match, but we cannot promise that this is
    always the case. If you want to make sure, do some testing by allocating big
    buffers and checking memory usage using ``nvidia-smi``.

    Example to allocate and check GPU mem usage::

        import subprocess

        import wgpu
        import torch

        def allocate_gpu_mem_with_wgpu(idx):
            a = wgpu.gpu.enumerate_adapters_sync()[idx]
            d = a.request_device_sync()
            b = d.create_buffer(size=10*2**20, usage=wgpu.BufferUsage.COPY_DST)
            return b

        def allocate_gpu_mem_with_torch(idx):
            d = torch.device(f"cuda:{idx}")
            return torch.ones([2000, 10], dtype=torch.float32, device=d)

        def show_mem_usage():
            print(subprocess.run(["nvidia-smi"]))

    See https://github.com/pygfx/wgpu-py/issues/482 for more details.
    """
    wgpu.preconfigure_default_device("pygfx", adapter=adapter)


def enable_wgpu_features(*features):
    """Enable specific features (as strings) on the wgpu device.

    This is an alias to ``wgpu.preconfigure_default_device(.., required_features=...)``.

    WARNING: enabling features means that your code may not work on all
    devices. The point of wgpu is that it can make a promise that a
    visualization works and looks the same on any device. Using features
    breaks that promise, and may cause your code to not work on e.g.
    mobile devices or certain operating systems.

    Features can also be turned off by prefixing with "!". This e.g. allows
    using a subset of Pygfx on old hardware by using
    ``enable_wgpu_features("!float32-filterable")``. Use with care; some parts
    of Pygfx will not function without the default features enabled.

    This function must be called before before the first ``Renderer`` is created.
    It can be called multiple times to enable more features. Note that feature names
    are invariant to use of dashes versus underscores.

    For more information on features:

    * ``wgpu.FeatureName`` for all possible official features.
    * ``renderer.device.adapter.features`` for the features available on the current system.
    * ``renderer.device.features`` for the currently enabled features.
    * https://gpuweb.github.io/gpuweb/#gpufeaturename for the official webgpu features (excl. native features).
    * https://docs.rs/wgpu/latest/wgpu/struct.Features.html for the features and their limitations in wgpu-core.
    """
    wgpu.preconfigure_default_device("pygfx", required_features=features)


def set_wgpu_limits(**limits):
    """Set specific limits (as key-value pairs) on the wgpu device.

    This is an alias to ``wgpu.preconfigure_default_device(.., required_limits=...)``.

    WARNING: setting high limits may make your code less portable across devices.

    This function must be called before before the first ``Renderer`` is created.
    It can be called multiple times to override or enable more limits. Note that
    limit names are invariant to use of dashes versus underscores.

    For more information on limits:

    * ``renderer.device.adapter.limits`` for the (max) limits available on the current system.
    * ``renderer.device.limits`` for the currently set limits.
    * https://gpuweb.github.io/gpuweb/#limits for the official webgpu limits.
    """
    wgpu.preconfigure_default_device("pygfx", required_limits=limits)


def get_shared():
    """Get the globally shared instance.

    Creates it if it does not yet exist. This should not be called at the import
    time of any module. Use this to get the global device:
    ``get_shared().device``.
    """
    if Shared._instance is None:
        Shared()

    return Shared._instance


class PyGfxAdapterInfoDiagnostics(wgpu.DiagnosticsBase):
    def get_dict(self):
        shared = get_shared()
        adapter = shared.adapter
        return adapter.info


class PyGfxFeaturesDiagnostics(wgpu.DiagnosticsBase):
    def get_dict(self):
        shared = get_shared()
        adapter = shared.adapter
        device = shared.device

        feature_names = list(wgpu.FeatureName)
        feature_names += sorted(adapter.features.difference(wgpu.FeatureName))
        result = {}
        for key in feature_names:
            result[key] = {
                "adapter": key in adapter.features,
                "device": key in device.features,
            }
        return result


class PyGfxLimitsDiagnostics(wgpu.DiagnosticsBase):
    def get_dict(self):
        shared = get_shared()
        adapter = shared.adapter
        device = shared.device

        result = {}
        for key in adapter.limits.keys():
            result[key] = {
                "adapter": adapter.limits[key],
                "device": device.limits.get(key, False),
            }
        return result


class PygfxCacheDiagnostics(wgpu.DiagnosticsBase):
    def get_dict(self):
        result = {}
        for cache_name, stats in gpu_caches.get_stats().items():
            count, hits, misses = stats
            result[cache_name] = {"count": count, "hits": hits, "misses": misses}
        return result


class PygfxResourceDiagnostics(wgpu.DiagnosticsBase):
    def get_dict(self):
        return Resource._resource_counts


def print_wgpu_report():
    """Print a report on the internal status of WGPU. Can be useful
    in debugging, and for providing details when making a bug report.
    """
    wgpu.diagnostics.print_report()
