# Material extraction across Arnold / Redshift / V-Ray / standardSurface /
# openPBRSurface / MaterialX / Lambert. Pulls full PBR slots when available:
# base color, metallic, roughness, normal, emissive, occlusion, opacity.
# Walks the upstream shading graph to find texture nodes behind
# utility / colorCorrect / bump / normalMap / layered nodes.
#
# Maya's biggest pain versus Blender is that standardSurface, openPBRSurface,
# and MaterialX graphs each name their inputs differently and route textures
# through renderer-specific nodes, so a naive "find a file node on .color"
# misses most modern materials. This module solves that with: a wide
# attribute net (settings.*_ATTRS), a path-bearing-node BFS that is agnostic
# to node type, an existing-file path resolver, and a whole-history fallback
# that guarantees the albedo is never silently dropped.

import os

from maya import cmds

from .. import settings
from . import materialx_document


_MAX_UPSTREAM_DEPTH = 16


def _neutralize_textured_factors(out):
    # When a texture drives a slot, Maya/MaterialX ignore the stored scalar at
    # render time, but glTF multiplies factor x texture. A leftover gray base
    # color (e.g. 0.8) would wrongly darken a textured albedo, so reset the
    # factors to neutral wherever a texture is present.
    if out["textures"].get("base_color"):
        alpha = out["base_color"][3] if len(out["base_color"]) > 3 else 1.0
        out["base_color"] = [1.0, 1.0, 1.0, alpha]
    if out["textures"].get("emissive"):
        out["emissive"] = [1.0, 1.0, 1.0]
    if out["textures"].get("metallic"):
        out["metallic"] = 1.0
    if out["textures"].get("roughness"):
        out["roughness"] = 1.0


def _materialx_xml_for(material_node):
    # MaterialXSurfaceShader stores the evaluated graph on .renderDocument.
    # Fall back to the bind document on a connected materialxStack if needed.
    for attr in ("renderDocument", "documentString", "document"):
        plug = "{0}.{1}".format(material_node, attr)
        try:
            if cmds.objExists(plug):
                val = cmds.getAttr(plug)
                if isinstance(val, str) and "<materialx" in val:
                    return val
        except Exception:
            pass
    try:
        stacks = cmds.listConnections(material_node, type="materialxStack") or []
    except Exception:
        stacks = []
    for stack in stacks:
        for attr in ("documentString", "document", "renderDocument"):
            plug = "{0}.{1}".format(stack, attr)
            try:
                if cmds.objExists(plug):
                    val = cmds.getAttr(plug)
                    if isinstance(val, str) and "<materialx" in val:
                        return val
            except Exception:
                pass
    return None


def _extract_from_materialx(material_node, out):
    # Returns True if the material was resolved from a MaterialX document.
    node_type = _safe_node_type(material_node)
    if node_type not in ("MaterialXSurfaceShader",):
        # Some LookdevX materials are other node types but still carry a
        # renderDocument; treat any node exposing one as MaterialX.
        if not cmds.objExists("{0}.renderDocument".format(material_node)):
            return False

    xml_text = _materialx_xml_for(material_node)
    if not xml_text:
        return False

    data = materialx_document.read_render_document(xml_text)
    if not data.get("ok"):
        return False

    tex = out["textures"]

    def _path(slot):
        entry = data.get(slot) or {}
        t = entry.get("texture")
        return t.get("path") if t else None

    bc = data["base_color"]
    if bc.get("texture"):
        tex["base_color"] = bc["texture"].get("path")
    elif bc.get("value"):
        v = bc["value"]
        if len(v) >= 3:
            out["base_color"] = [v[0], v[1], v[2], 1.0]

    mt = data["metallic"]
    if mt.get("texture"):
        tex["metallic"] = mt["texture"].get("path")
    elif mt.get("value"):
        out["metallic"] = float(mt["value"][0])

    rg = data["roughness"]
    if rg.get("texture"):
        tex["roughness"] = rg["texture"].get("path")
        out["roughness_invert"] = bool(rg.get("invert_hint"))
    elif rg.get("value"):
        out["roughness"] = float(rg["value"][0])

    nm = data["normal"]
    if nm.get("texture"):
        tex["normal"] = nm["texture"].get("path")
        if nm.get("strength") is not None:
            out["normal_scale"] = float(nm["strength"])

    em = data["emissive"]
    if em.get("texture"):
        tex["emissive"] = em["texture"].get("path")
    elif em.get("value"):
        v = em["value"]
        weight = em.get("weight")
        scale = float(weight[0]) if weight else 1.0
        if len(v) >= 3:
            out["emissive"] = [
                max(0.0, min(1.0, v[0] * scale)),
                max(0.0, min(1.0, v[1] * scale)),
                max(0.0, min(1.0, v[2] * scale)),
            ]

    op = data["opacity"]
    if op.get("texture"):
        tex["opacity"] = op["texture"].get("path")

    # MaterialX standard_surface is single-sided by default; leave double_sided
    # to the global UI override.
    return True


def _safe(node, attr, default=None):
    try:
        return cmds.getAttr("{0}.{1}".format(node, attr))
    except Exception:
        return default


def _safe_node_type(node):
    try:
        return cmds.nodeType(node)
    except Exception:
        return ""


def _color_to_rgba(node, attr):
    val = _safe(node, attr)
    if val is None:
        return None
    if isinstance(val, list) and val and isinstance(val[0], (list, tuple)):
        triple = val[0]
    elif isinstance(val, (list, tuple)) and len(val) >= 3:
        triple = val
    else:
        return None
    try:
        return [float(triple[0]), float(triple[1]), float(triple[2]), 1.0]
    except Exception:
        return None


def _scalar(node, attr, default=0.0):
    v = _safe(node, attr)
    if v is None:
        return default
    if isinstance(v, list):
        if not v:
            return default
        first = v[0]
        if isinstance(first, (list, tuple)):
            try:
                return float(first[0])
            except Exception:
                return default
        try:
            return float(first)
        except Exception:
            return default
    try:
        return float(v)
    except Exception:
        return default


def shading_engines_for_shape(mesh_shape):
    if not mesh_shape:
        return []
    sgs = cmds.listConnections(mesh_shape, type="shadingEngine") or []
    seen = set()
    out = []
    for sg in sgs:
        if sg in seen:
            continue
        seen.add(sg)
        out.append(sg)
    return out


def get_material_for_sg(sg):
    # Try surfaceShader first, then renderer-specific plugs.
    candidates = ("surfaceShader", "aiSurfaceShader", "rsSurfaceShader", "vrayMaterial")
    for plug_attr in candidates:
        plug = "{0}.{1}".format(sg, plug_attr)
        try:
            if not cmds.objExists(plug):
                continue
        except Exception:
            continue
        srcs = cmds.listConnections(plug, source=True, destination=False) or []
        for sh in srcs:
            return sh
    return None


def get_material_for_mesh(mesh_shape):
    for sg in shading_engines_for_shape(mesh_shape):
        mat = get_material_for_sg(sg)
        if mat:
            return mat
    return None


def _candidate_paths_from_node(node):
    # Collect every non-empty texture-path string on this node.
    out = []
    if not node:
        return out
    for attr in settings.TEXTURE_PATH_ATTRS:
        plug = "{0}.{1}".format(node, attr)
        try:
            if not cmds.objExists(plug):
                continue
        except Exception:
            continue
        path = _safe(node, attr)
        if isinstance(path, str) and path.strip():
            out.append(path.strip())
    return out


def _texture_path_from_node(node):
    # Prefer a path that resolves to a real file on disk. This avoids
    # returning a UDIM/frame pattern (e.g. wood.<UDIM>.exr) from
    # computedFileTextureNamePattern when a concrete fileTextureName exists,
    # and lets env-var/workspace-resolved paths win over raw ones.
    candidates = _candidate_paths_from_node(node)
    if not candidates:
        return None
    for path in candidates:
        try:
            if os.path.isfile(os.path.normpath(path)):
                return path
        except Exception:
            pass
    # None exist on disk (UDIM pattern, missing file). Return the first so the
    # caller can warn rather than silently drop the slot.
    return candidates[0]


def _bfs_upstream_texture(material_node, seed_attrs):
    # Breadth-first walk upstream from the requested material attributes.
    # Texture nodes vary by renderer, so path-bearing attributes are more
    # reliable than a fixed node-type allowlist.
    if not material_node:
        return None
    visited = set()
    queue = []
    seeded = False
    for attr in seed_attrs:
        plug = "{0}.{1}".format(material_node, attr)
        try:
            if not cmds.objExists(plug):
                continue
        except Exception:
            continue
        srcs = cmds.listConnections(plug, source=True, destination=False) or []
        for s in srcs:
            queue.append((s, 0))
            seeded = True
    if not seeded:
        return None
    while queue:
        node, depth = queue.pop(0)
        if not node or node in visited:
            continue
        if depth > _MAX_UPSTREAM_DEPTH:
            continue
        visited.add(node)
        if _texture_path_from_node(node):
            return node
        # bump2d, aiNormalMap, etc. -- walk through
        upstream = cmds.listConnections(node, source=True, destination=False) or []
        for u in upstream:
            queue.append((u, depth + 1))
    return None


def _texture_path(material_node, attrs):
    return _texture_path_from_node(
        _bfs_upstream_texture(material_node, attrs)
    )


def _chain_has_invert(material_node, seed_attrs):
    # True if any aiColorCorrect / colorCorrect / reverse node upstream of the
    # given material attributes has its invert flag enabled.
    visited = set()
    queue = []
    for attr in seed_attrs:
        plug = "{0}.{1}".format(material_node, attr)
        try:
            if not cmds.objExists(plug):
                continue
        except Exception:
            continue
        for s in (cmds.listConnections(plug, source=True, destination=False) or []):
            queue.append((s, 0))
    while queue:
        node, depth = queue.pop(0)
        if not node or node in visited or depth > _MAX_UPSTREAM_DEPTH:
            continue
        visited.add(node)
        ntype = _safe_node_type(node)
        if ntype == "reverse":
            return True
        if ntype in ("aiColorCorrect", "colorCorrect"):
            for a in ("invert", "invertX", "alphaIsLuminance"):
                if a == "invert" and _safe(node, "invert"):
                    return True
        for u in (cmds.listConnections(node, source=True, destination=False) or []):
            queue.append((u, depth + 1))
    return False


def _uv_transform_from_node(node):
    # Read UV tiling / offset so it can be written as KHR_texture_transform.
    # Arnold aiImage uses sscale/tscale/soffset/toffset; Maya file nodes use a
    # place2dTexture with repeatU/V + offsetU/V. Returns {scale,offset} or None.
    if not node:
        return None
    ntype = _safe_node_type(node)
    s = t = 1.0
    so = to = 0.0
    if ntype == "aiImage":
        s = _scalar(node, "sscale", 1.0)
        t = _scalar(node, "tscale", 1.0)
        so = _scalar(node, "soffset", 0.0)
        to = _scalar(node, "toffset", 0.0)
    else:
        p2ds = cmds.listConnections(node, type="place2dTexture") or []
        if p2ds:
            p = p2ds[0]
            s = _scalar(p, "repeatU", 1.0)
            t = _scalar(p, "repeatV", 1.0)
            so = _scalar(p, "offsetU", 0.0)
            to = _scalar(p, "offsetV", 0.0)
    if (abs(s - 1.0) < 1e-6 and abs(t - 1.0) < 1e-6
            and abs(so) < 1e-6 and abs(to) < 1e-6):
        return None
    return {"scale": [s, t], "offset": [so, to]}


def _history_base_color_fallback(material_node, already_found):
    # Last resort for opaque MaterialX / LookdevX / third-party graphs whose
    # input attribute names we do not recognize: walk the entire upstream
    # history and take the first texture node that is not already claimed by
    # another slot. This guarantees the albedo is exported even when the
    # material node hides it behind a non-standard plug.
    if not material_node:
        return None
    try:
        history = cmds.listHistory(material_node, pruneDagObjects=True) or []
    except Exception:
        return None
    claimed = set(p for p in already_found if p)
    for node in history:
        if node == material_node:
            continue
        path = _texture_path_from_node(node)
        if not path:
            continue
        if path in claimed:
            continue
        # Skip obvious non-color maps when their node name hints at it.
        lowered = node.lower()
        if any(tag in lowered for tag in (
                "normal", "bump", "rough", "metal", "displace", "_ao", "occlusion")):
            continue
        return path
    return None


def _read_color_value(material_node, attrs):
    for attr in attrs:
        plug = "{0}.{1}".format(material_node, attr)
        try:
            if not cmds.objExists(plug):
                continue
        except Exception:
            continue
        c = _color_to_rgba(material_node, attr)
        if c is not None:
            return c
    return None


def _read_scalar_value(material_node, attrs, default):
    for attr in attrs:
        plug = "{0}.{1}".format(material_node, attr)
        try:
            if not cmds.objExists(plug):
                continue
        except Exception:
            continue
        return _scalar(material_node, attr, default)
    return default


def _read_opacity(material_node):
    for attr in settings.OPACITY_ATTRS:
        plug = "{0}.{1}".format(material_node, attr)
        try:
            if not cmds.objExists(plug):
                continue
        except Exception:
            continue
        value = _scalar(material_node, attr, 1.0)
        if attr == "transparency":
            value = 1.0 - value
        return max(0.0, min(1.0, value))
    return 1.0


def extract_material(material_node):
    # Returns a dict describing a Maya material with full PBR slot population
    # where possible. Texture entries are file-node paths on disk (callers will
    # copy them and assign URIs).
    out = {
        "name": settings.DEFAULT_MATERIAL_NAME,
        "base_color": list(settings.DEFAULT_BASE_COLOR),
        "metallic": 0.0,
        "roughness": 1.0,
        "emissive": [0.0, 0.0, 0.0],
        "double_sided": False,
        "alpha_mode": "OPAQUE",
        "normal_scale": None,
        "roughness_invert": False,
        "source_type": None,
        "uv_transform": None,
        "transmission": 0.0,
        "ior": 1.5,
        "textures": {
            "base_color": None,
            "metallic": None,
            "roughness": None,
            "normal": None,
            "emissive": None,
            "occlusion": None,
            "opacity": None,
        },
    }
    if not material_node or not cmds.objExists(material_node):
        return out

    out["name"] = material_node
    out["source_type"] = _safe_node_type(material_node)

    # MaterialXSurfaceShader / LookdevX materials keep their whole graph in an
    # embedded MaterialX XML document, not as Maya DG nodes, so the DG-based
    # extraction below would find nothing. Resolve them from the document.
    if _extract_from_materialx(material_node, out):
        _neutralize_textured_factors(out)
        return out

    # Scalar / color factors
    base_color = _read_color_value(material_node, settings.BASE_COLOR_ATTRS)
    if base_color is not None:
        out["base_color"] = base_color

    out["metallic"] = _read_scalar_value(material_node, settings.METALLIC_ATTRS, 0.0)
    out["roughness"] = _read_scalar_value(material_node, settings.ROUGHNESS_ATTRS, 1.0)

    emissive = _read_color_value(material_node, settings.EMISSIVE_ATTRS)
    if emissive is not None:
        emission_weight = _read_scalar_value(
            material_node, settings.EMISSION_WEIGHT_ATTRS, 1.0
        )
        # glTF core emissiveFactor must stay within [0, 1]. openPBRSurface
        # emissionLuminance is in nits and can be large, so clamp here rather
        # than emit an out-of-range factor.
        out["emissive"] = [
            max(0.0, min(1.0, float(channel) * emission_weight))
            for channel in emissive[:3]
        ]

    opacity_texture = _texture_path(material_node, settings.OPACITY_ATTRS)
    if opacity_texture:
        out["base_color"][3] = 1.0
        out["alpha_mode"] = "MASK"
        out["alpha_cutoff"] = 0.5
        out["textures"]["opacity"] = opacity_texture
        # Cutout foliage (alpha mask) is almost always double-sided; without
        # this the back-facing leaf cards get culled and the plant looks empty.
        out["double_sided"] = True
    else:
        opacity = _read_opacity(material_node)
        if opacity < 0.999:
            out["base_color"][3] = max(0.0, min(1.0, opacity))
            out["alpha_mode"] = "BLEND"

    # Textures
    out["textures"]["base_color"] = _texture_path(
        material_node, settings.BASE_COLOR_ATTRS
    )
    out["textures"]["normal"] = _texture_path(
        material_node, settings.NORMAL_ATTRS
    )
    out["textures"]["emissive"] = _texture_path(
        material_node, settings.EMISSIVE_ATTRS
    )
    out["textures"]["occlusion"] = _texture_path(
        material_node, settings.OCCLUSION_ATTRS
    )
    out["textures"]["roughness"] = _texture_path(
        material_node, settings.ROUGHNESS_ATTRS
    )
    out["textures"]["metallic"] = _texture_path(
        material_node, settings.METALLIC_ATTRS
    )
    if out["textures"]["roughness"]:
        out["roughness"] = 1.0
    if out["textures"]["metallic"]:
        out["metallic"] = 1.0

    # UV tiling / offset (KHR_texture_transform), read from the base-color node.
    out["uv_transform"] = _uv_transform_from_node(
        _bfs_upstream_texture(material_node, settings.BASE_COLOR_ATTRS)
    )

    # Glass / transmission -> KHR_materials_transmission + KHR_materials_ior.
    transmission = _read_scalar_value(material_node, settings.TRANSMISSION_ATTRS, 0.0)
    if transmission > 1e-4:
        out["transmission"] = max(0.0, min(1.0, transmission))
        out["ior"] = _read_scalar_value(material_node, settings.IOR_ATTRS, 1.5)
        # A clear transmissive surface should not be tinted dark by a stray base
        # color factor when it has no albedo texture.
        if not out["textures"]["base_color"]:
            out["base_color"] = [1.0, 1.0, 1.0, out["base_color"][3]]

    # Honor an aiColorCorrect 'invert' in the roughness chain so the user can
    # flip a smoothness/roughness map from Maya without us guessing.
    if out["textures"]["roughness"] and not out["roughness_invert"]:
        if _chain_has_invert(material_node, settings.ROUGHNESS_ATTRS):
            out["roughness_invert"] = True

    # Albedo safety net: if the recognized base-color attributes yielded no
    # texture (common with MaterialX / LookdevX graphs), scan the whole
    # upstream history for a color map that no other slot already claimed.
    if not out["textures"]["base_color"]:
        already = [
            out["textures"]["normal"],
            out["textures"]["roughness"],
            out["textures"]["metallic"],
            out["textures"]["emissive"],
            out["textures"]["occlusion"],
            out["textures"]["opacity"],
        ]
        fallback = _history_base_color_fallback(material_node, already)
        if fallback:
            out["textures"]["base_color"] = fallback

    _neutralize_textured_factors(out)

    # Two-sidedness hint
    for attr in ("doubleSided", "double_sided"):
        v = _safe(material_node, attr)
        if v is not None:
            try:
                out["double_sided"] = bool(v)
                break
            except Exception:
                pass

    return out
