# MaterialX document reader.
#
# Maya's MaterialXSurfaceShader / LookdevX nodes do NOT expose their textures
# as Maya DG 'file' nodes. The whole shading graph is serialized as a MaterialX
# XML string on the shader's 'renderDocument' (and the bind document on the
# materialxStack). Any DG-traversal exporter therefore finds zero textures and
# the material exports flat -- this is the root cause of "MaterialX materials
# show no textures" in glTF/GLB.
#
# This module parses that XML (standard library only, no Maya dependency) and
# resolves the standard_surface / open_pbr_surface inputs down to the concrete
# <image>/<tiledimage> file paths, carrying per-texture colorspace and a few
# conversion hints (normal strength, probable glossiness). The result is a
# portable, glTF-friendly channel map that the material pipeline can turn into
# baseColor / metallic-roughness / normal / emissive textures.
#
# No f-strings. ASCII only.

import os
import xml.etree.ElementTree as ET


# MaterialX nodes that are pure pass-through / type or colorspace converters:
# we walk straight through them to reach the upstream image node.
_PASSTHROUGH_NODES = (
    "adsk_converter",
    "convert",
    "normalize",
    "remap",
    "range",
    "clamp",
    "smoothstep",
    "contrast",
    "saturate",
    "hsvadjust",
    "luminance",
    "extract",
    "separate",
    "swizzle",
    "channels",
)

# Nodes that carry an actual texture file path.
_IMAGE_NODES = ("image", "tiledimage", "tiledImage")

# MaterialX surface-shader categories we understand.
_SURFACE_NODES = (
    "standard_surface",
    "open_pbr_surface",
    "openpbr_surface",
    "UsdPreviewSurface",
    "gltf_pbr",
)

# Strip an MaterialX namespace prefix from a tag (e.g. '{uri}node' -> 'node').
def _localname(tag):
    if "}" in tag:
        return tag.split("}", 1)[1]
    return tag


class MaterialXDoc(object):

    def __init__(self, xml_text):
        self.ok = False
        self.error = None
        self.colorspace = None          # document working color space
        self._nodes_by_name = {}        # name -> element
        self._surface = None
        try:
            # The render document may have a leading XML declaration; ET handles it.
            root = ET.fromstring(xml_text)
        except Exception as exc:
            self.error = "MaterialX XML parse failed: {0}".format(exc)
            return
        self._root = root
        self.colorspace = root.get("colorspace")
        for el in root.iter():
            name = el.get("name")
            if name:
                self._nodes_by_name[name] = el
        self._surface = self._find_surface()
        self.ok = self._surface is not None
        if not self.ok and self.error is None:
            self.error = "No standard_surface / open_pbr_surface node found."

    # ---------- discovery ----------

    def _find_surface(self):
        # Prefer the shader referenced by a surfacematerial node; otherwise the
        # first surface-shader-typed node in the document.
        for el in self._root.iter():
            if _localname(el.tag) == "surfacematerial":
                ss = self._input_nodename(el, "surfaceshader")
                if ss and ss in self._nodes_by_name:
                    return self._nodes_by_name[ss]
        for el in self._root.iter():
            cat = _localname(el.tag)
            if cat in _SURFACE_NODES:
                return el
            if el.get("type") == "surfaceshader" and el.get("nodedef"):
                return el
        return None

    def surface_category(self):
        if self._surface is None:
            return None
        return _localname(self._surface.tag)

    # ---------- input helpers ----------

    def _find_input(self, element, input_name):
        for child in element:
            if _localname(child.tag) != "input":
                continue
            if child.get("name") == input_name:
                return child
        return None

    def _input_nodename(self, element, input_name):
        inp = self._find_input(element, input_name)
        if inp is None:
            return None
        # MaterialX uses 'nodename' for node-graph connections and
        # 'output'/'nodegraph' for graph outputs; we honor nodename + nodegraph.
        return inp.get("nodename") or inp.get("nodegraph")

    def _input_value(self, element, input_name):
        inp = self._find_input(element, input_name)
        if inp is None:
            return None
        return inp.get("value")

    # ---------- texture resolution ----------

    def _resolve_image(self, node_name, _depth=0):
        # Walk from a node down through pass-through converters to the first
        # image node. Returns dict {path, colorspace, strength, invert_hint}
        # or None.
        if not node_name or _depth > 24:
            return None
        el = self._nodes_by_name.get(node_name)
        if el is None:
            return None
        cat = _localname(el.tag)

        if cat in _IMAGE_NODES:
            file_inp = self._find_input(el, "file")
            if file_inp is None:
                return None
            path = file_inp.get("value")
            if not path:
                return None
            return {
                "path": path,
                "colorspace": file_inp.get("colorspace") or el.get("colorspace"),
                "strength": None,
                "node": node_name,
            }

        # normal_map node: capture strength, then follow its 'input'.
        if cat in ("normal_map", "normalmap"):
            strength = self._input_value(el, "strength")
            up = self._input_nodename(el, "input") or self._input_nodename(el, "in")
            res = self._resolve_image(up, _depth + 1)
            if res is not None and strength is not None:
                try:
                    res["strength"] = float(strength)
                except Exception:
                    pass
            return res

        # Pass-through / converter nodes: follow the first connected input.
        if cat in _PASSTHROUGH_NODES or el.get("Autodesk-internalConverter"):
            for child in el:
                if _localname(child.tag) != "input":
                    continue
                up = child.get("nodename") or child.get("nodegraph")
                if up:
                    res = self._resolve_image(up, _depth + 1)
                    if res is not None:
                        return res
            return None

        # Unknown node with a single upstream input: try to follow it so we do
        # not give up on graphs that use math/mix nodes.
        for child in el:
            if _localname(child.tag) != "input":
                continue
            up = child.get("nodename") or child.get("nodegraph")
            if up:
                res = self._resolve_image(up, _depth + 1)
                if res is not None:
                    return res
        return None

    def _channel_texture(self, input_name):
        nodename = self._input_nodename(self._surface, input_name)
        if not nodename:
            return None
        return self._resolve_image(nodename)

    def _channel_constant(self, input_name):
        raw = self._input_value(self._surface, input_name)
        if raw is None:
            return None
        parts = raw.replace(",", " ").split()
        try:
            nums = [float(p) for p in parts]
        except Exception:
            return None
        return nums if nums else None

    # ---------- public API ----------

    def extract(self):
        # Returns a portable channel map describing the material in glTF terms.
        result = {
            "ok": self.ok,
            "error": self.error,
            "surface": self.surface_category(),
            "doc_colorspace": self.colorspace,
            "base_color": {"texture": None, "value": None},
            "metallic": {"texture": None, "value": None},
            "roughness": {"texture": None, "value": None, "invert_hint": False},
            "normal": {"texture": None, "strength": None},
            "emissive": {"texture": None, "value": None, "weight": None},
            "opacity": {"texture": None, "value": None},
            "transmission": {"value": None},
            "coat": {"value": None},
            "specular": {"value": None},
        }
        if not self.ok:
            return result

        # base color
        bc = self._channel_texture("base_color")
        result["base_color"]["texture"] = bc
        result["base_color"]["value"] = self._channel_constant("base_color")

        # metalness
        mt = self._channel_texture("metalness")
        result["metallic"]["texture"] = mt
        result["metallic"]["value"] = self._channel_constant("metalness")

        # specular roughness (detect glossiness by source filename)
        rg = self._channel_texture("specular_roughness")
        result["roughness"]["texture"] = rg
        result["roughness"]["value"] = self._channel_constant("specular_roughness")
        if rg and rg.get("path"):
            stem = os.path.basename(rg["path"]).lower()
            if ("gloss" in stem or "smooth" in stem) and "rough" not in stem:
                result["roughness"]["invert_hint"] = True

        # normal
        nm = self._channel_texture("normal")
        if nm:
            result["normal"]["texture"] = nm
            result["normal"]["strength"] = nm.get("strength")

        # emission
        em = self._channel_texture("emission_color")
        result["emissive"]["texture"] = em
        result["emissive"]["value"] = self._channel_constant("emission_color")
        result["emissive"]["weight"] = self._channel_constant("emission")

        # opacity
        op = self._channel_texture("opacity")
        result["opacity"]["texture"] = op
        result["opacity"]["value"] = self._channel_constant("opacity")

        # advanced scalar hints (for Standard/Advanced profiles)
        result["transmission"]["value"] = self._channel_constant("transmission")
        result["coat"]["value"] = self._channel_constant("coat")
        result["specular"]["value"] = self._channel_constant("specular")

        return result


def read_render_document(xml_text):
    # Convenience wrapper: parse + extract in one call.
    doc = MaterialXDoc(xml_text)
    return doc.extract()


def colorspace_is_srgb(colorspace):
    if not colorspace:
        return False
    cs = colorspace.lower()
    return "srgb" in cs or "rec709" in cs or "rec 709" in cs


def colorspace_is_data(colorspace):
    # 'none', 'raw', 'linear', 'scene-linear' are treated as non-color data.
    if not colorspace:
        return True
    cs = colorspace.lower()
    if "srgb" in cs:
        return False
    return cs in ("none", "raw", "linear") or "linear" in cs or "data" in cs
