# Mesh extraction: positions, normals, UVs, tangents, vertex colors,
# triangulation, per-face-vertex dedup, multi-material primitive split.
# Uses maya.api.OpenMaya (API 2.0).

import math

from maya import cmds
import maya.api.OpenMaya as om

from . import scene_scan
from . import material_utils
from .. import settings


def _rnd(v):
    return round(float(v), settings.VERTEX_KEY_PRECISION)


def _dag_path(name):
    sel = om.MSelectionList()
    sel.add(name)
    return sel.getDagPath(0)


def _short(name):
    return name.split("|")[-1].split(":")[-1]


def _face_index_from_component(comp_str):
    # 'pCube1.f[12]' -> 12; returns None on failure.
    try:
        inside = comp_str.split(".f[", 1)[1]
        idx = inside.split("]", 1)[0]
        return int(idx)
    except Exception:
        return None


def _member_matches_shape(member_base, mesh_shape):
    try:
        shape_paths = cmds.ls(mesh_shape, long=True) or []
        member_paths = cmds.ls(member_base, long=True) or []
    except Exception:
        return False
    if not shape_paths or not member_paths:
        return False
    shape_path = shape_paths[0]
    shape_parent = cmds.listRelatives(
        shape_path, parent=True, fullPath=True
    ) or []
    valid_paths = set([shape_path] + shape_parent)
    for member_path in member_paths:
        if member_path in valid_paths:
            return True
        member_shapes = cmds.listRelatives(
            member_path, shapes=True, fullPath=True, type="mesh"
        ) or []
        if shape_path in member_shapes:
            return True
    return False


def face_to_sg_map(mesh_shape):
    # Returns a dict face_index -> shading_engine name.
    sgs = material_utils.shading_engines_for_shape(mesh_shape)
    if not sgs:
        return {}, []
    try:
        total = cmds.polyEvaluate(mesh_shape, face=True)
    except Exception:
        return {}, sgs
    if not total:
        return {}, sgs

    mapping = {}
    for sg in sgs:
        members = cmds.sets(sg, query=True) or []
        # Restrict to members of this shape only.
        try:
            expanded = cmds.ls(members, long=True, flatten=True) or []
        except Exception:
            expanded = members
        for e in expanded:
            if not e:
                continue
            base = e.split(".")[0]
            if not _member_matches_shape(base, mesh_shape):
                continue
            if ".f[" in e:
                idx = _face_index_from_component(e)
                if idx is None:
                    continue
                if idx not in mapping:
                    mapping[idx] = sg
            else:
                # Whole-mesh assignment -- fill remaining unmapped faces.
                for i in range(total):
                    if i not in mapping:
                        mapping[i] = sg
    # Any face still unmapped goes to the first SG so we still produce geometry.
    if mapping:
        fallback_sg = sgs[0]
        for i in range(total):
            if i not in mapping:
                mapping[i] = fallback_sg
    return mapping, sgs


def _gather_face_vertex_attrs(dag, mesh_fn, opts):
    # Pre-fetch face-vertex level attributes from the mesh.
    points = mesh_fn.getPoints(om.MSpace.kObject)

    normals = None
    if opts.get("include_normals", True):
        try:
            normals = mesh_fn.getNormals(om.MSpace.kObject)
        except Exception:
            normals = None

    u_array, v_array = None, None
    has_uvs = False
    if opts.get("include_uvs", True):
        try:
            uv_set = mesh_fn.currentUVSetName()
            if uv_set and mesh_fn.numUVs(uv_set) > 0:
                u_array, v_array = mesh_fn.getUVs(uv_set)
                has_uvs = True
        except Exception:
            has_uvs = False

    # Tangents and binormals (per face-vertex). Required for normal maps.
    tangents = None
    binormals = None
    if opts.get("include_tangents", True) and has_uvs:
        try:
            tangents = mesh_fn.getTangents(om.MSpace.kObject)
            binormals = mesh_fn.getBinormals(om.MSpace.kObject)
        except Exception:
            tangents = None
            binormals = None

    color_array = None
    has_colors = False
    if opts.get("include_vertex_colors", False):
        try:
            color_sets = mesh_fn.getColorSetNames() or []
            if color_sets:
                color_array = mesh_fn.getFaceVertexColors(colorSet=color_sets[0])
                has_colors = True
        except Exception:
            has_colors = False

    return {
        "points": points,
        "normals": normals,
        "u_array": u_array,
        "v_array": v_array,
        "has_uvs": has_uvs,
        "tangents": tangents,
        "binormals": binormals,
        "color_array": color_array,
        "has_colors": has_colors,
    }


def _new_primitive(name, has_normals, has_uvs, has_tangents, has_colors):
    return {
        "name": name,
        "positions": [],
        "normals": [] if has_normals else None,
        "uvs": [] if has_uvs else None,
        "tangents": [] if has_tangents else None,
        "colors": [] if has_colors else None,
        "indices": [],
        "vertex_map": {},
        "triangle_count": 0,
    }


def _finalize_primitive(prim):
    # Strip None lists and the temporary vertex_map.
    out = {
        "name": prim["name"],
        "positions": prim["positions"],
        "indices": prim["indices"],
        "triangle_count": len(prim["indices"]) // 3,
    }
    for key in ("normals", "uvs", "tangents", "colors"):
        v = prim.get(key)
        out[key] = v if v else []
    return out


def extract_mesh(transform_path, opts):
    # Returns (mesh_data, warnings) where mesh_data is:
    # {"name": str, "primitives": [ {primitive dict (per SG)} ... ]}
    warnings = []
    shape = scene_scan.get_mesh_shape(transform_path)
    if not shape:
        return None, ["No mesh shape under '{0}'".format(transform_path)]

    dag = _dag_path(shape)
    mesh_fn = om.MFnMesh(dag)

    attrs = _gather_face_vertex_attrs(dag, mesh_fn, opts)

    split_by_material = opts.get("split_by_material", True)
    face_sg, sgs = face_to_sg_map(shape)
    if not split_by_material:
        face_sg = {}  # treat whole mesh as one primitive

    primitives_by_sg = {}
    mesh_name = _short(transform_path)

    has_normals = attrs["normals"] is not None
    has_uvs = attrs["has_uvs"]
    has_tangents = attrs["tangents"] is not None and attrs["binormals"] is not None
    has_colors = attrs["has_colors"]
    flip_v = opts.get("flip_uv_v", True)

    poly_iter = om.MItMeshPolygon(dag)
    fv_counter = 0  # running face-vertex index across the whole mesh
    source_face_count = 0
    source_triangle_face_count = 0
    source_quad_count = 0
    source_ngon_count = 0

    while not poly_iter.isDone():
        face_index = poly_iter.index()
        sg = face_sg.get(face_index) if face_sg else None
        sg_key = sg if sg else "__default__"
        if sg_key not in primitives_by_sg:
            prim_name = "{0}_{1}".format(mesh_name, _short(sg) if sg else "default")
            prim = _new_primitive(prim_name, has_normals, has_uvs, has_tangents, has_colors)
            prim["material_node"] = material_utils.get_material_for_sg(sg) if sg else None
            primitives_by_sg[sg_key] = prim
        prim = primitives_by_sg[sg_key]

        poly_verts = poly_iter.getVertices()
        local_count = len(poly_verts)
        source_face_count += 1
        if local_count == 3:
            source_triangle_face_count += 1
        elif local_count == 4:
            source_quad_count += 1
        elif local_count > 4:
            source_ngon_count += 1

        local_normals = None
        if has_normals:
            local_normals = []
            for i in range(local_count):
                try:
                    local_normals.append(poly_iter.normalIndex(i))
                except Exception:
                    local_normals.append(-1)

        local_uvs = None
        if has_uvs:
            local_uvs = []
            for i in range(local_count):
                try:
                    local_uvs.append(poly_iter.getUVIndex(i))
                except Exception:
                    local_uvs.append(-1)

        local_tangents = None
        if has_tangents:
            local_tangents = []
            for i in range(local_count):
                try:
                    local_tangents.append(poly_iter.tangentIndex(i))
                except Exception:
                    local_tangents.append(-1)

        local_colors = None
        if has_colors:
            local_colors = []
            for i in range(local_count):
                idx = fv_counter + i
                if 0 <= idx < len(attrs["color_array"]):
                    c = attrs["color_array"][idx]
                    local_colors.append((c.r, c.g, c.b, c.a))
                else:
                    local_colors.append((1.0, 1.0, 1.0, 1.0))

        try:
            _tri_points, tri_verts = poly_iter.getTriangles(om.MSpace.kObject)
        except Exception:
            tri_verts = None

        if not tri_verts or len(tri_verts) % 3 != 0:
            warnings.append(
                "Mesh '{0}': failed to triangulate face {1}.".format(mesh_name, face_index)
            )
            fv_counter += local_count
            poly_iter.next()
            continue

        slot_for_pos = {}
        for slot, pos_idx in enumerate(poly_verts):
            slot_for_pos.setdefault(pos_idx, slot)

        for tri_pos_idx in tri_verts:
            slot = slot_for_pos.get(tri_pos_idx, 0)
            pos_idx = poly_verts[slot]
            p = attrs["points"][pos_idx]
            pos_tuple = (_rnd(p.x), _rnd(p.y), _rnd(p.z))

            nrm_tuple = None
            if local_normals and 0 <= local_normals[slot] < len(attrs["normals"]):
                n = attrs["normals"][local_normals[slot]]
                nrm_tuple = (_rnd(n.x), _rnd(n.y), _rnd(n.z))

            uv_tuple = None
            if local_uvs and 0 <= local_uvs[slot] < len(attrs["u_array"]):
                uv_idx = local_uvs[slot]
                u_val = attrs["u_array"][uv_idx]
                v_val = attrs["v_array"][uv_idx]
                if flip_v:
                    v_val = 1.0 - v_val
                uv_tuple = (_rnd(u_val), _rnd(v_val))

            tan_tuple = None
            if local_tangents and 0 <= local_tangents[slot] < len(attrs["tangents"]):
                t = attrs["tangents"][local_tangents[slot]]
                b = attrs["binormals"][local_tangents[slot]]
                # Handedness w = sign(dot(cross(N, T), B))
                w = 1.0
                if nrm_tuple is not None:
                    nx, ny, nz = nrm_tuple
                    cx = ny * t.z - nz * t.y
                    cy = nz * t.x - nx * t.z
                    cz = nx * t.y - ny * t.x
                    dot = cx * b.x + cy * b.y + cz * b.z
                    w = -1.0 if dot < 0.0 else 1.0
                tan_tuple = (_rnd(t.x), _rnd(t.y), _rnd(t.z), w)

            col_tuple = None
            if local_colors:
                c = local_colors[slot]
                col_tuple = (_rnd(c[0]), _rnd(c[1]), _rnd(c[2]), _rnd(c[3]))

            key = (pos_tuple, nrm_tuple, uv_tuple, tan_tuple, col_tuple)
            v_index = prim["vertex_map"].get(key)
            if v_index is None:
                v_index = len(prim["positions"]) // 3
                prim["vertex_map"][key] = v_index
                prim["positions"].extend(pos_tuple)
                if prim["normals"] is not None:
                    prim["normals"].extend(nrm_tuple if nrm_tuple else (0.0, 1.0, 0.0))
                if prim["uvs"] is not None:
                    prim["uvs"].extend(uv_tuple if uv_tuple else (0.0, 0.0))
                if prim["tangents"] is not None:
                    prim["tangents"].extend(tan_tuple if tan_tuple else (1.0, 0.0, 0.0, 1.0))
                if prim["colors"] is not None:
                    prim["colors"].extend(col_tuple if col_tuple else (1.0, 1.0, 1.0, 1.0))
            prim["indices"].append(v_index)

        fv_counter += local_count
        poly_iter.next()

    primitives = [_finalize_primitive(p) for p in primitives_by_sg.values()]
    # Carry the material node along for the exporter to extract.
    for orig, finalized in zip(primitives_by_sg.values(), primitives):
        finalized["material_node"] = orig.get("material_node")

    return {
        "name": mesh_name,
        "primitives": primitives,
        "source_face_count": source_face_count,
        "source_triangle_face_count": source_triangle_face_count,
        "source_quad_count": source_quad_count,
        "source_ngon_count": source_ngon_count,
    }, warnings
