"""
gltf_extensions.py
------------------

Extension registry for glTF import/export with scope-based handlers.
Each scope has a TypedDict defining the context passed to handlers.
"""

from collections.abc import Callable, Iterable
from typing import Any, Literal, TypeAlias, TypedDict

import numpy as np

from ...constants import log
from ...iteration import IndexedDict
from ...typed import NDArray

# Scopes define where in the glTF load/export process handlers run:
#   material            - after parsing material, can override PBR values
#   texture_source      - when resolving texture image index
#   primitive           - after loading primitive, can add face_attributes
#   primitive_preprocess - before accessor reads, can modify accessors in-place
#   primitive_export    - during mesh export, can compress/modify primitive data
Scope: TypeAlias = Literal[
    "material", "texture_source", "primitive", "primitive_preprocess", "primitive_export"
]


# ----------------------------------------------------------------------
# TypedDict contexts for each scope
# ----------------------------------------------------------------------
#
# These TypedDicts define the MINIMUM fields passed to handlers for each scope.
# Additional fields may be added in future versions for new functionality.
#
# FOR FORWARD COMPATIBILITY: Handlers should access only the fields they need
# and ignore unknown fields. The context is passed as a plain dict at runtime,
# so handlers can safely use dict.get() for optional access or simply not
# reference fields they don't need.
#
# Example handler pattern:
#
#     def my_handler(context: MaterialContext) -> dict | None:
#         # Access only what you need - additional fields won't break this
#         data = context["data"]
#         images = context["images"]
#         return {"baseColorFactor": [1, 0, 0, 1]}
#
# ----------------------------------------------------------------------


class MaterialContext(TypedDict):
    """Context for material scope handlers."""

    data: dict[str, Any]
    parse_textures: Callable[..., dict[str, Any]]
    images: list


class TextureSourceContext(TypedDict):
    """Context for texture_source scope handlers."""

    data: dict[str, Any]


class PrimitiveContext(TypedDict):
    """Context for primitive scope handlers (post-load)."""

    data: dict[str, Any]
    primitive: dict
    mesh_kwargs: dict
    accessors: list


class PrimitivePreprocessContext(TypedDict):
    """Context for primitive_preprocess scope handlers (pre-load)."""

    data: dict[str, Any]
    primitive: dict
    accessors: list
    views: list


class PrimitiveExportContext(TypedDict):
    """Context for primitive_export scope handlers (during export)."""

    mesh: Any
    name: str
    tree: dict
    # a `bufferView` is the position of an entry in here, so the order matters
    buffer_items: IndexedDict
    primitive: dict
    # the arrays the primitive's accessors were built from, in the dtype they
    # would have been written with, keyed by accessor index: a handler storing
    # them itself never unpacks bytes back into numpy, and the exporter stores
    # them after all if no handler claims them
    arrays: dict[int, NDArray]


# Handler type alias - handlers receive a context dict
Handler: TypeAlias = Callable[[Any], Any]

# callback to parse material dict and resolve texture references
# signature: (*, data: dict) -> dict
ParseTextures: TypeAlias = Callable[..., dict[str, Any]]

# Registry: {scope: {extension_name: handler}}
_handlers: dict[str, dict[str, Handler]] = {}


def _deep_merge(target: dict, source: dict) -> None:
    """
    Recursively merge source dict into target dict.

    Parameters
    ----------
    target
      Dict to merge into (modified in place)
    source
      Dict to merge from
    """
    for key, value in source.items():
        if isinstance(value, dict) and key in target and isinstance(target[key], dict):
            # Both are dicts - recurse
            _deep_merge(target[key], value)
        else:
            # Overwrite or set new key
            target[key] = value


def register_handler(name: str, scope: Scope) -> Callable[[Handler], Handler]:
    """
    Decorator to register a handler for a glTF extension.

    Parameters
    ----------
    name
      Extension name, e.g. "KHR_materials_pbrSpecularGlossiness".
    scope
      Handler scope, e.g. "material", "texture_source", "primitive".

    Returns
    -------
    decorator
      Function that registers the handler and returns it unchanged.

    Example
    -------
    >>> @register_handler("MY_extension", scope="material")
    ... def my_handler(context: MaterialContext) -> dict | None:
    ...     data = context["data"]
    ...     images = context["images"]
    ...     return {"baseColorFactor": [1, 0, 0, 1]}
    """
    if scope not in _handlers:
        _handlers[scope] = {}

    def decorator(func: Handler) -> Handler:
        _handlers[scope][name] = func
        return func

    return decorator


def unregistered(extensions: Iterable[str], scope: Scope) -> set:
    """
    Find extension names with no registered handler for a scope.

    Parameters
    ----------
    extensions
      Extension names, i.e. the keys of a glTF "extensions" dict.
    scope
      Handler scope to check against.

    Returns
    -------
    missing
      Extension names with no handler registered for the scope.
    """
    return set(extensions) - _handlers.get(scope, {}).keys()


def handle_extensions(
    *,
    extensions: dict[str, Any] | None,
    scope: Scope,
    failed: set | None = None,
    **kwargs,
) -> Any:
    """
    Process extensions dict for a given scope, calling registered handlers.

    Parameters
    ----------
    extensions
      The "extensions" dict from a glTF element, or None.
    scope
      Handler scope to invoke.
    failed
      If passed the name of any extension whose handler raised is added here.
    **kwargs
      Scope-specific arguments that will be combined with extension data
      into a typed context dict. Required kwargs by scope:
        - material: parse_textures, images
        - texture_source: (none)
        - primitive: primitive, mesh_kwargs, accessors
        - primitive_preprocess: primitive, accessors, views
        - primitive_export: mesh, name, tree, buffer_items, primitive, arrays

    Returns
    -------
    results
      Dict of {extension_name: result} for most scopes.
      For scopes ending in "_source", returns first non-None result.
      For "primitive" scope, automatically merges results into mesh_kwargs.
    """
    if not extensions or scope not in _handlers:
        return {} if not scope.endswith("_source") else None

    results = {}
    for ext_name, data in extensions.items():
        if ext_name not in _handlers[scope]:
            continue
        try:
            # Build context dict with data + all kwargs
            context = {"data": data, **kwargs}
            if (result := _handlers[scope][ext_name](context)) is not None:
                results[ext_name] = result
        except Exception as e:
            if failed is not None:
                failed.add(ext_name)
            log.warning(f"failed to process extension {ext_name}: {e}")

    # for _source scopes return first result, otherwise return all results
    if scope.endswith("_source"):
        return next(iter(results.values()), None)

    # for primitive scope, automatically merge results into mesh_kwargs
    if scope == "primitive" and "mesh_kwargs" in kwargs:
        mesh_kwargs = kwargs["mesh_kwargs"]
        for ext_result in results.values():
            if not isinstance(ext_result, dict):
                continue
            # merge extension results, recursively merging nested dicts
            for key, value in ext_result.items():
                if isinstance(value, dict):
                    if key not in mesh_kwargs:
                        mesh_kwargs[key] = {}
                    _deep_merge(mesh_kwargs[key], value)
                else:
                    mesh_kwargs[key] = value

    return results


# ----------------------------------------------------------------------
# Built-in handlers
# ----------------------------------------------------------------------


@register_handler("KHR_materials_pbrSpecularGlossiness", scope="material")
def _specular_glossiness(context: MaterialContext) -> dict[str, Any] | None:
    """
    Convert specular-glossiness material to PBR metallic-roughness.

    Parameters
    ----------
    context
      MaterialContext with extension data, parse_textures function, and images.

    Returns
    -------
    pbr_dict
      PBR metallic-roughness parameters, or None on failure.
    """
    try:
        from ...visual.gloss import specular_to_pbr

        return specular_to_pbr(**context["parse_textures"](data=context["data"]))
    except Exception:
        log.debug("failed to convert specular-glossiness", exc_info=True)
        return None


@register_handler("EXT_texture_webp", scope="texture_source")
def _texture_webp_source(context: TextureSourceContext) -> int | None:
    """
    Return image source index from EXT_texture_webp.

    Parameters
    ----------
    context
      TextureSourceContext with extension data.

    Returns
    -------
    source_index
      Index into glTF images array, or None if not present.
    """
    return context["data"].get("source")


# the optional attributes draco can absorb, as
# (glTF name, `DracoPy.encode` keyword, `DracoPy.AttributeType`)
_draco_optional = (
    ("COLOR_0", "colors", "COLOR"),
    ("TEXCOORD_0", "tex_coord", "TEX_COORD"),
    ("NORMAL", "normals", "NORMAL"),
)

# how hard draco tries, which trades export time for size and is not lossy
_DRACO_COMPRESSION = 6

# bits draco quantizes positions onto, which is what blender emits: a vertex
# moves up to half a step, i.e. `mesh.extents.max() * 2**-(bits + 1)`
_DRACO_QUANTIZATION = 14


@register_handler("KHR_draco_mesh_compression", scope="primitive_preprocess")
def draco_decode(context: PrimitivePreprocessContext) -> None:
    """
    Replace a primitive's placeholder accessors with decompressed draco data.

    The accessors of a draco-compressed primitive have no `bufferView`, so the
    loader filled them with zeros before calling us. All of the geometry is in
    a single opaque blob, and the extension carries the indirection we need to
    unpack it: a mapping of glTF attribute name to draco attribute id.

    Parameters
    ----------
    context
      PrimitivePreprocessContext, whose `accessors` we mutate in-place.
    """
    import DracoPy

    data = context["data"]
    accessors = context["accessors"]
    attributes = context["primitive"].get("attributes", {})

    # one blob holds every compressed attribute for this primitive
    decoded = DracoPy.decode(context["views"][data["bufferView"]])

    # the extension stores name -> draco id and we look up the other way
    names = {ident: name for name, ident in data["attributes"].items()}

    # overwrite the zero placeholders in-place by accessor index
    for attr in decoded.attributes:
        name = names.get(attr["unique_id"])
        if name in attributes:
            accessors[attributes[name]] = attr["data"]

    # indices aren't an attribute so they're not in the extension mapping
    indices = context["primitive"].get("indices")
    faces = getattr(decoded, "faces", None)
    if indices is not None and faces is not None:
        accessors[indices] = faces


@register_handler("KHR_draco_mesh_compression", scope="primitive_export")
def draco_encode(context: PrimitiveExportContext) -> bool | None:
    """
    Compress a primitive's geometry into a single draco buffer.

    Every array in `arrays` is absorbed, so the exporter left their accessors
    with no `bufferView` rather than storing the same data twice. Returning
    `None` makes it store them after all, so a failure here exports the
    primitive uncompressed rather than pointing at accessors full of zeros.

    Parameters
    ----------
    context
      PrimitiveExportContext, whose `primitive` and `buffer_items` we mutate.

    Returns
    -------
    compressed
      True if the geometry is now inside a draco buffer, None if not.
    """
    import DracoPy

    from . import _buffer_append

    primitive = context["primitive"]
    attributes = primitive["attributes"]
    # the arrays the accessors were built from keyed by accessor index, so we
    # never unpack bytes back into numpy and can only claim what was recorded
    arrays = context["arrays"]

    # the optional attributes this primitive actually has: anything the
    # exporter stored itself, like a custom `_ATTRIBUTE`, is not in `arrays`
    absorb, optional = [], {}
    for name, keyword, kind in _draco_optional:
        value = arrays.get(attributes.get(name))
        if value is None:
            continue
        absorb.append((name, kind))
        # DracoPy asserts on float64 for UV and normals, colors stay uint8
        optional[keyword] = value if value.dtype == np.uint8 else value.astype(np.float64)

    position = arrays[attributes["POSITION"]]
    indices = arrays[primitive["indices"]]

    # `preserve_order` is load-bearing: without it draco rewelds and permutes
    # vertices, which would invalidate the count/min/max already written into
    # every accessor and silently misindex any attribute we didn't compress
    buffer = DracoPy.encode(
        points=position,
        faces=indices,
        preserve_order=True,
        quantization_bits=_DRACO_QUANTIZATION,
        compression_level=_DRACO_COMPRESSION,
        **optional,
    )

    # decode what we just encoded, which does two things we can't get otherwise:
    # it reports the attribute ids draco assigned, and it proves the round trip
    # preserved our counts before we let the exporter skip storing the source.
    # DO NOT remove this check: it is all that stands between a lossy encoder
    # and silently corrupt geometry in the exported file.
    check = DracoPy.decode(buffer)
    if len(check.points) != len(position) or len(check.faces) != len(indices):
        log.warning("draco round trip changed the vertex count, not compressing")
        return None

    # draco identifies attributes by kind, the extension identifies them by id.
    # an attribute draco silently dropped is missing from `ids`, and the
    # `KeyError` leaves this primitive uncompressed rather than pointing at
    # an accessor the exporter is about to stop storing
    ids = {attr["attribute_type"]: attr["unique_id"] for attr in check.attributes}
    absorbed = {
        name: ids[getattr(DracoPy.AttributeType, kind)]
        for name, kind in [("POSITION", "POSITION"), *absorb]
    }

    primitive.setdefault("extensions", {})["KHR_draco_mesh_compression"] = {
        # identical geometry encodes identically so this dedupes for free
        "bufferView": _buffer_append(context["buffer_items"], buffer),
        "attributes": absorbed,
    }
    return True
