"""Blender 内运行：FBX -> 纯模型 FBX + XPbr 基础/混合贴图。
调用：blender -b --factory-startup --python bl_convert.py -- <src.fbx> <outdir> <max_size>
"""
import os
import re
import sys

import numpy as np

# blender --python 不会把脚本目录加入 sys.path，手动加入以复用 unity_assets（纯 stdlib）
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from unity_assets import safe_name

AO_NAME_RE = re.compile(r"(ao|occlusion|ambient)", re.I)


def encode_normal_rg(n):
    """切线法线 (N,3) -> 球面编码 (N,2)，范围 [0,1]。
    shader 侧 XPbr_Base.hlsl DecodeSphereMap 先 *2-1 再解码。"""
    n = n / np.maximum(np.linalg.norm(n, axis=-1, keepdims=True), 1e-8)
    z = np.clip(n[..., 2], -0.99, 1.0)
    denom = np.sqrt(2.0 * (1.0 + z))
    enc = n[..., :2] / denom[..., None]
    return np.clip(enc * 0.5 + 0.5, 0.0, 1.0)


def decode_sphere_map(enc01):
    """shader DecodeSphereMap 的 numpy 等价（含 *2-1），仅测试用。"""
    e = enc01 * 2.0 - 1.0
    l = 1.0 - e[..., 0] ** 2 - e[..., 1] ** 2
    l = np.clip(l, 0.0, 1.0)
    xy = e * np.sqrt(l)[..., None]
    return np.concatenate([xy * 2.0, (l * 2.0 - 1.0)[..., None]], axis=-1)


# ---------------- 以下仅在 Blender 内执行 ----------------


IMG_EXTS = ('.png', '.jpg', '.jpeg', '.tga', '.exr', '.bmp', '.tif', '.tiff')


def _looks_like_ao(name):
    return bool(AO_NAME_RE.search(name))


def _image_broken(img):
    """外置贴图文件丢失/加载失败时，Blender 留下 0 尺寸空图。"""
    return img.size[0] == 0 or img.size[1] == 0


def _image_stem(img):
    """贴图文件名（不含扩展名）；外置贴图优先用 filepath，比 img.name 可靠。"""
    path = img.filepath if img.filepath else img.name
    return os.path.splitext(os.path.basename(path))[0]


def _find_ao_on_disk(src_dir, maps, warnings):
    """材质未引用 AO 时，从 FBX 同目录按文件名兜底搜索（烘焙工具常只把 _ao.png 放旁边）。
    多候选时优先与 basecolor 贴图同前缀（xxx_color.png -> xxx_ao.png）。"""
    import bpy
    try:
        files = sorted(os.listdir(src_dir))
    except OSError:
        return None
    cands = [f for f in files if f.lower().endswith(IMG_EXTS)
             and _looks_like_ao(os.path.splitext(f)[0])]
    if not cands:
        return None
    pick = None
    if maps["basecolor"] and len(cands) > 1:
        prefix = _image_stem(maps["basecolor"][0]).rsplit("_", 1)[0]
        matched = [f for f in cands if f.startswith(prefix)]
        if matched:
            pick = matched[0]
    if pick is None:
        if len(cands) == 1:
            pick = cands[0]
        else:
            warnings.append("同目录有多张疑似 AO 贴图无法确定归属，忽略：%s" % ", ".join(cands))
            return None
    try:
        img = bpy.data.images.load(os.path.join(src_dir, pick), check_existing=True)
    except RuntimeError as e:
        warnings.append("AO 贴图 %s 加载失败：%s" % (pick, e))
        return None
    if _image_broken(img):
        warnings.append("AO 贴图 %s 加载失败（0 尺寸）" % pick)
        return None
    return img


def _image_pixels(img):
    """bpy Image -> float32 (h, w, 4)，原始值（不做色彩变换）。"""
    w, h = img.size
    buf = np.empty(w * h * 4, dtype=np.float32)
    img.pixels.foreach_get(buf)
    return buf.reshape(h, w, 4)


def _scaled_pixels(img, tw, th):
    """取缩放到 (tw, th) 的像素；用副本 scale，不污染原图。"""
    if tuple(img.size) == (tw, th):
        return _image_pixels(img)
    cp = img.copy()
    try:
        cp.scale(tw, th)
        return _image_pixels(cp)
    finally:
        import bpy
        bpy.data.images.remove(cp)


def _trace_image(socket, warnings, what):
    """沿输入连线找贴图。返回 (bpy.types.Image | None, channel | None)。
    穿透 Normal Map / Separate Color 节点；channel 为 0..3 或 None（用 R/灰度）。"""
    if not socket.is_linked:
        return None, None
    node = socket.links[0].from_node
    out_name = socket.links[0].from_socket.name
    if node.type == 'TEX_IMAGE':
        return node.image, None
    if node.type == 'NORMAL_MAP':
        return _trace_image(node.inputs['Color'], warnings, what)
    if node.type in ('SEPARATE_COLOR', 'SEPRGB'):
        ch = {'Red': 0, 'Green': 1, 'Blue': 2, 'R': 0, 'G': 1, 'B': 2}.get(out_name, 0)
        img, _ = _trace_image(node.inputs[0], warnings, what)
        return img, ch
    warnings.append("%s: 无法识别的节点 %s，改用常量" % (what, node.type))
    return None, None


def _material_maps(mat, warnings, src_dir):
    """从材质 Principled BSDF 提取贴图与常量。"""
    maps = {"basecolor": None, "alpha": None, "normal": None,
            "metallic": None, "roughness": None, "ao": None}
    consts = {"basecolor": (0.5, 0.5, 0.5, 1.0), "metallic": 0.0, "roughness": 0.5}
    if not mat.use_nodes:
        warnings.append("材质未用节点，全部用默认常量")
        return maps, consts
    bsdf = next((n for n in mat.node_tree.nodes if n.type == 'BSDF_PRINCIPLED'), None)
    if bsdf is None:
        warnings.append("未找到 Principled BSDF，全部用默认常量")
        return maps, consts

    def grab(input_name, key):
        sock = bsdf.inputs.get(input_name)
        if sock is None:
            return
        img, ch = _trace_image(sock, warnings, key)
        if img is not None:
            if _image_broken(img):
                warnings.append("%s 贴图 %s 加载失败（外置文件缺失？），改用常量"
                                % (key, img.name))
            else:
                maps[key] = (img, ch)

    grab('Base Color', 'basecolor')
    grab('Alpha', 'alpha')
    grab('Normal', 'normal')
    grab('Metallic', 'metallic')
    grab('Roughness', 'roughness')
    consts["basecolor"] = tuple(bsdf.inputs['Base Color'].default_value)
    consts["metallic"] = float(bsdf.inputs['Metallic'].default_value)
    consts["roughness"] = float(bsdf.inputs['Roughness'].default_value)

    # AO：材质节点树内按名匹配（未连到 BSDF 也算）
    used = {m[0] for m in maps.values() if m}
    for n in mat.node_tree.nodes:
        if n.type == 'TEX_IMAGE' and n.image and n.image not in used \
                and not _image_broken(n.image) and _looks_like_ao(n.image.name):
            maps["ao"] = (n.image, None)
            break
    if maps["ao"] is None:
        img = _find_ao_on_disk(src_dir, maps, warnings)
        if img is not None:
            maps["ao"] = (img, None)
    return maps, consts


def _target_size(maps, max_size):
    ws = [m[0].size[0] for m in maps.values() if m]
    hs = [m[0].size[1] for m in maps.values() if m]
    if not ws:
        return 4, 4
    return min(max(ws), max_size), min(max(hs), max_size)


def _channel(pixels, ch):
    return pixels[..., ch if ch is not None else 0]


def _save_png(path, rgba):
    """float32 (h,w,4) -> PNG。写原始值（Non-Color，避免色彩管理改写数据）。"""
    import bpy
    h, w = rgba.shape[:2]
    img = bpy.data.images.new("mt_tmp", width=w, height=h, alpha=True)
    img.colorspace_settings.name = 'Non-Color'
    img.pixels.foreach_set(np.ascontiguousarray(rgba, dtype=np.float32).ravel())
    img.filepath_raw = path
    img.file_format = 'PNG'
    img.alpha_mode = 'CHANNEL_PACKED'
    img.save()
    bpy.data.images.remove(img)


def _convert_material(mat, model_name, outdir, max_size, used_snames, src_dir):
    warnings = []
    maps, consts = _material_maps(mat, warnings, src_dir)
    tw, th = _target_size(maps, max_size)

    # --- 基础贴图 ---
    base = np.empty((th, tw, 4), dtype=np.float32)
    if maps["basecolor"]:
        base[..., :3] = _scaled_pixels(maps["basecolor"][0], tw, th)[..., :3]
    else:
        base[..., :3] = np.array(consts["basecolor"][:3], dtype=np.float32)
    if maps["alpha"]:
        img, ch = maps["alpha"]
        base[..., 3] = _channel(_scaled_pixels(img, tw, th), ch)
        alpha_mode = "transparency"
    elif maps["basecolor"] and maps["basecolor"][0].channels == 4 and \
            _image_has_alpha(maps["basecolor"][0]):
        base[..., 3] = _scaled_pixels(maps["basecolor"][0], tw, th)[..., 3]
        alpha_mode = "transparency"
    elif maps["ao"]:
        img, ch = maps["ao"]
        base[..., 3] = _channel(_scaled_pixels(img, tw, th), ch)
        alpha_mode = "ao"
    else:
        base[..., 3] = 1.0
        alpha_mode = "white"

    # --- 混合贴图 ---
    mix = np.empty((th, tw, 4), dtype=np.float32)
    if maps["normal"]:
        npx = _scaled_pixels(maps["normal"][0], tw, th)[..., :3] * 2.0 - 1.0
        mix[..., 0:2] = encode_normal_rg(npx)
    else:
        mix[..., 0:2] = 0.5
        warnings.append("无法线贴图，RG 填平面法线")
    if maps["metallic"]:
        img, ch = maps["metallic"]
        mix[..., 2] = _channel(_scaled_pixels(img, tw, th), ch)
    else:
        mix[..., 2] = consts["metallic"]
    if maps["roughness"]:
        img, ch = maps["roughness"]
        mix[..., 3] = _channel(_scaled_pixels(img, tw, th), ch)
    else:
        mix[..., 3] = consts["roughness"]

    # 中文等非 ASCII 材质名会同样退化（如都变 "mat"），加序号防止互相覆盖
    sname = base_sname = safe_name(mat.name)
    n = 1
    while sname in used_snames:
        n += 1
        sname = "%s_%d" % (base_sname, n)
    used_snames.add(sname)
    if sname != base_sname:
        warnings.append("材质名 %s 清洗后与其它材质重名，输出改用 %s" % (mat.name, sname))
    base_png = os.path.abspath(os.path.join(outdir, "%s_%s_base.png" % (model_name, sname)))
    mix_png = os.path.abspath(os.path.join(outdir, "%s_%s_mix.png" % (model_name, sname)))
    _save_png(base_png, base)
    _save_png(mix_png, mix)

    return {
        "name": mat.name, "safe_name": sname,
        "base_png": os.path.basename(base_png), "mix_png": os.path.basename(mix_png),
        "alpha_mode": alpha_mode, "size": [tw, th],
        "sources": {k: (v[0].name if v else None) for k, v in maps.items()},
        "warnings": warnings,
    }


def _image_has_alpha(img):
    """采样判断 alpha 是否非全白（全白视为无透明信息）。"""
    px = _image_pixels(img)
    a = px[..., 3]
    return bool((a < 0.995).any())


def _strip_textures():
    import bpy
    for mat in bpy.data.materials:
        if not mat.use_nodes:
            continue
        for n in list(mat.node_tree.nodes):
            if n.type in ('TEX_IMAGE', 'NORMAL_MAP', 'SEPARATE_COLOR', 'SEPRGB'):
                mat.node_tree.nodes.remove(n)


def main():
    import bpy, json
    argv = sys.argv[sys.argv.index("--") + 1:]
    src, outdir, max_size = argv[0], argv[1], int(argv[2])
    os.makedirs(outdir, exist_ok=True)
    model_name = os.path.splitext(os.path.basename(src))[0]

    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.fbx(filepath=src)

    used_mats = []
    for obj in bpy.data.objects:
        if obj.type != 'MESH':
            continue
        for slot in obj.material_slots:
            if slot.material and slot.material not in used_mats:
                used_mats.append(slot.material)

    src_dir = os.path.dirname(os.path.abspath(src))
    used_snames = set()
    results = [_convert_material(m, model_name, outdir, max_size, used_snames, src_dir)
               for m in used_mats]

    _strip_textures()
    fbx_out = os.path.join(outdir, model_name + ".fbx")
    bpy.ops.export_scene.fbx(filepath=fbx_out, path_mode='STRIP', embed_textures=False)

    print("MT_SUMMARY " + json.dumps(
        {"model": model_name, "fbx": os.path.basename(fbx_out), "materials": results},
        ensure_ascii=False))


if __name__ == "__main__" and "--" in sys.argv:
    main()