"""Blender 内执行的渲染工作进程；请通过 Blender --python 调用。

输入只接受服务端创建的 task.json，不执行用户提供的 Python。
文件路径限制属于纵深防御，不替代操作系统对原生 .blend 解析器的隔离。
"""

from __future__ import annotations

import json
import os
import re
import sys
import time
import traceback
from pathlib import Path

import bpy


def atomic_json(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_name(path.name + ".tmp")
    temp.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
    os.replace(temp, path)


def inside(path: Path, root: Path) -> bool:
    try:
        path.resolve().relative_to(root.resolve())
        return True
    except (OSError, ValueError):
        return False


def all_ids():
    """遍历 Blender 的顶层数据块，不遍历文本内容，也不执行驱动器。"""
    seen = set()
    for prop in bpy.data.bl_rna.properties:
        if prop.type != "COLLECTION":
            continue
        for value in getattr(bpy.data, prop.identifier, ()):
            if isinstance(value, bpy.types.ID) and value.as_pointer() not in seen:
                seen.add(value.as_pointer())
                yield value


class AssetGuard:
    def __init__(self, root: Path, blend_path: Path, warnings: list):
        self.root = root.resolve()
        self.blend_path = blend_path.resolve()
        self.warnings = warnings
        self.assets = []
        self._seen = set()
        self.files = {}
        self.directories = {}
        # 只重定位到真正上传的文件；拒绝指向目录外的符号链接。
        for path in self.root.rglob("*"):
            if not inside(path, self.root):
                continue
            mapping = self.directories if path.is_dir() else self.files
            mapping.setdefault(path.name.casefold(), []).append(path.resolve())

    def resolve(self, original: str, owner=None, is_dir=False):
        normalized = original.replace("\\", "/")
        base = self.blend_path.parent
        library = getattr(owner, "library", None)
        if library:
            try:
                libpath = Path(bpy.path.abspath(library.filepath)).resolve()
                if inside(libpath, self.root):
                    base = libpath.parent
            except Exception:
                pass
        if normalized.startswith("//"):
            target = base / normalized[2:]
        elif re.match(r"^[A-Za-z]:/", normalized):
            target = Path("/__unavailable_windows_drive__/") / normalized[3:]
        else:
            target = Path(normalized)
            if not target.is_absolute():
                target = base / target
        target = target.resolve()
        if inside(target, self.root) and (target.is_dir() if is_dir else target.is_file()):
            return target, "ok"
        # UDIM 路径中的 token 不是实际文件，但每张 tile 都必须在上传目录。
        if "<UDIM>" in normalized and inside(target.parent, self.root):
            tiles = list(target.parent.glob(target.name.replace("<UDIM>", "[0-9][0-9][0-9][0-9]")))
            if tiles and all(inside(p, self.root) and p.is_file() for p in tiles):
                return target, "ok"
        # 相对路径被保留时优先匹配尾部路径，适应 Windows 上传至 Linux。
        parts = [p for p in normalized.split("/") if p and p not in (".", "..") and not p.endswith(":")]
        for count in range(min(len(parts), 8), 0, -1):
            candidate = (self.root.joinpath(*parts[-count:])).resolve()
            if inside(candidate, self.root) and (candidate.is_dir() if is_dir else candidate.is_file()):
                return candidate, "remapped"
        mapping = self.directories if is_dir else self.files
        matches = mapping.get(Path(normalized).name.casefold(), [])
        if len(matches) == 1:
            return matches[0], "remapped"
        return None, "missing" if inside(target, self.root) else "blocked"

    def record(self, kind, name, original, resolved, status):
        key = (kind, name, original)
        if key not in self._seen:
            self._seen.add(key)
            self.assets.append({"kind": kind, "name": name, "original": original,
                                "resolved": str(resolved) if resolved else None, "status": status})

    def clean_libraries(self):
        bad_libraries = set()
        for library in list(bpy.data.libraries):
            original = library.filepath
            # 链接库已在 open_mainfile 内加载：不能将越界库的已加载数据保留下来。
            # 不对已加载的库做 basename 重定位，以免把来源判断误认为已重载成功。
            raw = Path(bpy.path.abspath(original)).resolve()
            allowed = inside(raw, self.root) and raw.is_file()
            self.record("library", library.name, original, raw if allowed else None,
                        "ok" if allowed else "blocked")
            if allowed:
                library.filepath = str(raw)
            else:
                bad_libraries.add(library)
        changed = True
        while changed:
            changed = False
            for library in bpy.data.libraries:
                if getattr(library, "parent", None) in bad_libraries and library not in bad_libraries:
                    bad_libraries.add(library)
                    changed = True
        if bad_libraries:
            unsafe = [item for item in all_ids() if getattr(item, "library", None) in bad_libraries]
            bpy.data.batch_remove(ids=unsafe + list(bad_libraries))
            self.warnings.append("已删除上传目录外链接库及其数据块；请将库文件随项目一起上传。")

    def clean_datablock_files(self):
        # 未上传数据块直接删除，避免已经缓存于内存的外部像素或字体进入渲染。
        collections = ("images", "fonts", "sounds", "movieclips", "volumes", "cache_files")
        remove = []
        for kind in collections:
            for item in list(getattr(bpy.data, kind, ())):
                original = getattr(item, "filepath", "")
                if not original or (kind == "fonts" and original == "<builtin>"):
                    continue
                packed = bool(getattr(item, "packed_file", None)) or bool(getattr(item, "packed_files", []))
                if packed:
                    self.record(kind, item.name, original, None, "packed")
                    continue
                # 程序化图像和内部渲染结果不依赖磁盘。
                if kind == "images" and getattr(item, "source", "") in ("GENERATED", "VIEWER"):
                    continue
                path, status = self.resolve(original, item)
                self.record(kind, item.name, original, path, status)
                if path:
                    try:
                        item.filepath = str(path)
                        if kind == "images" and status == "remapped":
                            item.reload()
                    except (AttributeError, RuntimeError) as exc:
                        raise RuntimeError(f"无法安全重定位资源 {item.name}: {exc}") from exc
                else:
                    remove.append(item)
        if remove:
            bpy.data.batch_remove(ids=remove)

    def clean_rna_paths(self):
        # 仅检查明确参与渲染的属性；遍历整个 RNA 会进入 UI 工作区或巨大网格，
        # 还可能访问 Blender 中尚未初始化的指针，因此不能用通用递归扫描。
        visited = set()

        def scan(value, owner):
            if value is None or not hasattr(value, "bl_rna"):
                return
            pointer = value.as_pointer()
            if pointer in visited:
                return
            visited.add(pointer)
            for prop in value.bl_rna.properties:
                if prop.type != "STRING" or prop.subtype not in {"FILE_PATH", "DIR_PATH"} or prop.is_readonly:
                    continue
                original = getattr(value, prop.identifier)
                if not original or original == "<builtin>":
                    continue
                path, status = self.resolve(original, owner, prop.subtype == "DIR_PATH")
                self.record(f"{value.bl_rna.identifier}.{prop.identifier}", owner.name, original, path, status)
                try:
                    setattr(value, prop.identifier, str(path) if path else "")
                except Exception as exc:
                    raise RuntimeError(f"无法清除外部文件路径 {owner.name}/{prop.identifier}") from exc

        for scene in bpy.data.scenes:
            # 输出和烘焙路径不属于输入资产，由可信渲染流程另行指定。
            scene.render.filepath = ""
            if hasattr(scene.render, "bake"):
                scene.render.bake.filepath = ""
            rigid = scene.rigidbody_world
            if rigid:
                scan(rigid.point_cache, scene)

        for obj in bpy.data.objects:
            for modifier in obj.modifiers:
                scan(modifier, obj)
                for field in ("domain_settings", "flow_settings", "effector_settings", "point_cache"):
                    scan(getattr(modifier, field, None), obj)
                for bake in getattr(modifier, "bakes", ()):
                    scan(bake, obj)
                canvas = getattr(modifier, "canvas_settings", None)
                if canvas:
                    for surface in canvas.canvas_surfaces:
                        scan(surface, obj)
                        scan(surface.point_cache, obj)
            for system in obj.particle_systems:
                scan(system.point_cache, obj)
                scan(system.settings, obj)

        for texture in bpy.data.textures:
            scan(texture, texture)
            for field in ("voxel_data", "point_density"):
                scan(getattr(texture, field, None), texture)

        for cache_file in bpy.data.cache_files:
            # Alembic 缓存除了主文件，也可能叠加多份外部 layer 文件。
            for layer in getattr(cache_file, "layers", ()):
                scan(layer, cache_file)

        trees = list(bpy.data.node_groups)
        for collection in (bpy.data.materials, bpy.data.worlds, bpy.data.lights):
            for owner in collection:
                tree = getattr(owner, "node_tree", None)
                if tree:
                    trees.append(tree)
        for tree in trees:
            # 合成器和视频序列已经关闭，文件输出节点永远不能参与执行。
            if tree.bl_idname == "CompositorNodeTree":
                continue
            scan(tree, tree)
            for node in tree.nodes:
                scan(node, tree)

    def clean_nodes(self):
        trees = list(bpy.data.node_groups)
        for item in all_ids():
            tree = getattr(item, "node_tree", None)
            if tree:
                trees.append(tree)
        for tree in trees:
            for node in list(tree.nodes):
                kind = node.bl_idname
                # 几何节点的导入路径可能由动态字符串计算而来，不能靠静态路径检查保证安全。
                if kind == "ShaderNodeScript" or (kind.startswith("GeometryNode") and ("Import" in kind or "File" in kind)):
                    self.warnings.append(f"已禁用外部文件/脚本节点：{tree.name}/{node.name} ({kind})。")
                    tree.nodes.remove(node)

    def run(self):
        self.clean_libraries()
        self.clean_datablock_files()
        self.clean_nodes()
        self.clean_rna_paths()
        return self.assets


def harden_scenes(warnings):
    bpy.context.preferences.filepaths.use_scripts_auto_execute = False
    warnings.append("自动 Python 脚本、脚本节点和 OSL 已禁用；依赖 Python 表达式的驱动器或插件不会按原项目执行。")
    for scene in bpy.data.scenes:
        if scene.render.use_compositing:
            warnings.append(f"场景 {scene.name} 的合成器已禁用；输出为相机直接渲染结果。")
        if scene.render.use_sequencer and scene.sequence_editor:
            warnings.append(f"场景 {scene.name} 的视频序列编辑器已禁用。")
        scene.render.use_compositing = False
        scene.render.use_sequencer = False
        scene.render.use_stamp = False
        if hasattr(scene.render, "use_persistent_data"):
            scene.render.use_persistent_data = False
        if hasattr(scene, "cycles"):
            scene.cycles.shading_system = False
        # 禁止模拟缓存自动写入来源路径；模拟应事先烘焙并上传缓存。
        for obj in scene.objects:
            for modifier in obj.modifiers:
                if modifier.type == "FLUID" and modifier.fluid_type == "DOMAIN":
                    settings = modifier.domain_settings
                    if settings and hasattr(settings, "cache_type"):
                        settings.cache_type = "MODULAR"


def select_scene(config):
    name = config.get("scene")
    scene = bpy.data.scenes.get(name) if name else bpy.context.scene
    if scene is None:
        raise ValueError(f"找不到场景：{name}")
    if bpy.context.window:
        bpy.context.window.scene = scene
    camera_name = config.get("camera")
    if camera_name:
        camera = scene.objects.get(camera_name)
        if camera is None or camera.type != "CAMERA":
            raise ValueError(f"场景 {scene.name} 中没有相机：{camera_name}")
        scene.camera = camera
    return scene


def inspect_scene(assets):
    scenes = []
    for scene in bpy.data.scenes:
        scenes.append({"name": scene.name, "camera": scene.camera.name if scene.camera else None,
                       "cameras": [obj.name for obj in scene.objects if obj.type == "CAMERA"],
                       "frame_start": scene.frame_start, "frame_end": scene.frame_end,
                       "frame_step": scene.frame_step, "fps": scene.render.fps,
                       "fps_base": scene.render.fps_base,
                       "width": scene.render.resolution_x, "height": scene.render.resolution_y,
                       "resolution_percentage": scene.render.resolution_percentage,
                       "engine": scene.render.engine,
                       "view_layers": [layer.name for layer in scene.view_layers]})
    return {"scenes": scenes, "cameras": [obj.name for obj in bpy.data.objects if obj.type == "CAMERA"],
            "assets": assets,
            "missing_assets": [a for a in assets if a["status"] in {"missing", "blocked"}]}


def configure_render(scene, config):
    engine = config.get("engine", "EEVEE")
    if engine == "CYCLES":
        scene.render.engine = "CYCLES"
    else:
        engines = scene.render.bl_rna.properties["engine"].enum_items.keys()
        scene.render.engine = "BLENDER_EEVEE_NEXT" if "BLENDER_EEVEE_NEXT" in engines else "BLENDER_EEVEE"
    for key, default in (("resolution_x", 960), ("resolution_y", 540)):
        source = "width" if key == "resolution_x" else "height"
        setattr(scene.render, key, int(config.get(source, default)))
    scene.render.resolution_percentage = int(config.get("resolution_percentage", 100))
    scene.render.film_transparent = bool(config.get("transparent", False))
    scene.render.fps = int(config.get("fps", 24))
    scene.render.fps_base = 1.0
    scene.render.use_file_extension = True
    scene.render.use_overwrite = True
    scene.render.use_placeholder = False
    scene.render.image_settings.file_format = config.get("format", "PNG")
    scene.render.image_settings.color_mode = "RGBA"
    scene.render.image_settings.color_depth = str(config.get("color_depth", 8))
    if engine == "CYCLES":
        scene.cycles.samples = int(config.get("samples", 64))
        scene.cycles.use_denoising = bool(config.get("denoise", True))
        scene.cycles.use_adaptive_sampling = True
        scene.cycles.adaptive_threshold = float(config.get("adaptive_threshold", 0.01))
        scene.cycles.seed = int(config.get("seed", 0))
        scene.cycles.shading_system = False
    else:
        scene.eevee.taa_render_samples = int(config.get("samples", 64))
    if hasattr(scene.render, "use_motion_blur"):
        scene.render.use_motion_blur = bool(config.get("motion_blur", False))
    elif hasattr(scene.cycles, "use_motion_blur"):
        scene.cycles.use_motion_blur = bool(config.get("motion_blur", False))
    layer_name = config.get("view_layer")
    if layer_name:
        if scene.view_layers.get(layer_name) is None:
            raise ValueError(f"找不到视图层：{layer_name}")
        for layer in scene.view_layers:
            layer.use = layer.name == layer_name
    if not scene.camera:
        raise ValueError("场景没有活动相机，请在 JSON 中指定 camera 或先在 Blender 中添加相机。")


def device_candidates(requested, engine="CYCLES"):
    if engine == "EEVEE":
        return ("CPU" if requested == "CPU" else "EEVEE_GPU",)
    if requested == "AUTO":
        return ("OPTIX", "CUDA")
    if requested not in {"OPTIX", "CUDA", "CPU"}:
        raise ValueError(f"不支持的设备类型：{requested}")
    return (requested,)


def use_device(scene, backend):
    if scene.render.engine != "CYCLES":
        import gpu
        if hasattr(gpu, "init"):
            gpu.init()
        renderer = gpu.platform.renderer_get()
        vendor = gpu.platform.vendor_get()
        software = any(tag in renderer.lower() for tag in ("llvmpipe", "softpipe", "swrast", "software", "lavapipe"))
        if backend != "CPU" and software:
            raise RuntimeError(f"GPU + Eevee 检测到软件渲染器：{renderer}。服务器未提供可用的硬件图形接口；请检查 GPU 分区和 OpenGL/Vulkan 驱动。不会回退至 CPU。")
        if backend == "CPU" and not software:
            raise RuntimeError("CPU Eevee 未确认 Mesa 软件渲染器；拒绝误用硬件 GPU。")
        return [{"name": renderer, "vendor": vendor, "type": "CPU" if software else "GPU", "use": True}]
    addon = bpy.context.preferences.addons.get("cycles")
    if addon is None:
        raise RuntimeError("此 Blender 安装不包含 Cycles。")
    prefs = addon.preferences
    if backend == "CPU":
        scene.cycles.device = "CPU"
        return [{"name": "CPU", "type": "CPU", "use": True}]
    prefs.compute_device_type = backend
    if hasattr(prefs, "refresh_devices"):
        prefs.refresh_devices()
    else:
        prefs.get_devices()
    devices = []
    for device in prefs.devices:
        # 绝不启用混合 CPU 渲染；ZeroGPU 验证失败时必须明确报错。
        device.use = device.type == backend
        if device.use:
            devices.append({"name": device.name, "type": device.type, "id": device.id, "use": True})
    if not devices:
        raise RuntimeError(f"Cycles 未找到可用的 {backend} GPU；拒绝静默回退至 CPU。")
    scene.cycles.device = "GPU"
    return devices


def choose_device(scene, requested):
    failures = []
    for backend in device_candidates(requested):
        try:
            return backend, use_device(scene, backend)
        except Exception as exc:
            failures.append(f"{backend}: {exc}")
    raise RuntimeError("GPU 设备初始化失败。" + " | ".join(failures))


def make_probe_scene():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    scene = bpy.context.scene
    bpy.ops.mesh.primitive_cube_add()
    obj = bpy.context.object
    obj.rotation_euler = (0.2, 0.3, 0.0)
    bpy.ops.object.camera_add(location=(3, -4, 3))
    camera = bpy.context.object
    camera.rotation_euler = (-camera.location).to_track_quat("-Z", "Y").to_euler()
    scene.camera = camera
    bpy.ops.object.light_add(type="AREA", location=(2, -2, 4))
    bpy.context.object.data.energy = 500
    bpy.context.object.data.size = 3
    return scene


def render_one(scene, output_dir, frame, config):
    extension = ".exr" if config.get("format", "PNG") == "OPEN_EXR" else ".png"
    final = output_dir / f"frame_{frame:06d}{extension}"
    temporary = output_dir / f".frame_{frame:06d}.partial{extension}"
    scene.frame_set(frame)
    scene.render.filepath = str(temporary)
    started = time.monotonic()
    bpy.ops.render.render(write_still=True, scene=scene.name)
    if not temporary.is_file() or temporary.stat().st_size == 0:
        raise RuntimeError(f"Blender 没有生成第 {frame} 帧输出。")
    os.replace(temporary, final)
    record = {"frame": frame, "path": str(final), "seconds": round(time.monotonic() - started, 3),
              "bytes": final.stat().st_size}
    # 每帧刷盘：GPU 租约中途结束时，父进程仍能恢复已经完成的帧。
    with (output_dir / "completed.jsonl").open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, ensure_ascii=False) + "\n")
        handle.flush()
        os.fsync(handle.fileno())
    print("RENDER_FRAME_COMPLETED " + json.dumps(record), flush=True)
    return record


def main():
    if "--" not in sys.argv or len(sys.argv[sys.argv.index("--") + 1:]) != 1:
        raise ValueError("用法：blender --background --factory-startup --disable-autoexec --python render_worker.py -- task.json")
    task_path = Path(sys.argv[sys.argv.index("--") + 1]).resolve(strict=True)
    task = json.loads(task_path.read_text(encoding="utf-8"))
    result_path = Path(task["result_path"]).resolve()
    output_dir = Path(task["output_dir"]).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    config = task.get("config", {})
    result = {"ok": False, "action": task["action"], "blender_version": bpy.app.version_string,
              "engine": config.get("engine", "EEVEE"), "warnings": [], "devices": [], "backend": None, "frames": [], "timings": {}}
    start = time.monotonic()
    try:
        action = task["action"]
        if action == "probe":
            scene = make_probe_scene()
            probe_config = {**config, "width": 32, "height": 32, "resolution_percentage": 100,
                            "samples": 1, "denoise": False, "format": "PNG", "color_depth": "8"}
            harden_scenes(result["warnings"])
            configure_render(scene, probe_config)
            failures = []
            for backend in device_candidates(config.get("device", "CPU"), config.get("engine", "EEVEE")):
                try:
                    result["devices"] = use_device(scene, backend)
                    result["backend"] = backend
                    result["frames"] = [render_one(scene, output_dir, 1, probe_config)]
                    result["probe_render_succeeded"] = True
                    break
                except Exception as exc:
                    failures.append(f"{backend}: {exc}")
            else:
                raise RuntimeError("实际渲染探针失败。" + " | ".join(failures))
            result["warnings"].extend(failures)
        elif action in {"inspect", "render"}:
            root = Path(task["project_root"]).resolve(strict=True)
            blend_path = Path(task["blend_path"]).resolve(strict=True)
            if not inside(blend_path, root) or blend_path.suffix.lower() != ".blend":
                raise ValueError(".blend 文件必须位于上传项目目录内。")
            bpy.context.preferences.filepaths.use_scripts_auto_execute = False
            bpy.ops.wm.open_mainfile(filepath=str(blend_path), load_ui=False, use_scripts=False)
            harden_scenes(result["warnings"])
            assets = AssetGuard(root, blend_path, result["warnings"]).run()
            result.update(inspect_scene(assets))
            result["timings"]["load_and_inspect_seconds"] = round(time.monotonic() - start, 3)
            if action == "render":
                if config.get("strict_assets", True) and result["missing_assets"]:
                    missing = ", ".join(a["name"] for a in result["missing_assets"][:8])
                    raise ValueError(f"存在缺失或越界资源：{missing}。请上传完整资源，或将 strict_assets 设为 false 后接受缺失效果。")
                scene = select_scene(config)
                configure_render(scene, config)
                frames = task.get("frames") or [int(config.get("frame", config.get("frame_start", scene.frame_current)))]
                for frame in frames:
                    if not isinstance(frame, int) or isinstance(frame, bool) or abs(frame) > 1_048_574:
                        raise ValueError("帧编号必须是 Blender 支持范围内的整数。")
                # AUTO 必须通过实际首帧验证；OPTIX 初始化成功也可能在加载内核时失败。
                failures = []
                for backend in device_candidates(config.get("device", "CPU"), config.get("engine", "EEVEE")):
                    try:
                        result["devices"] = use_device(scene, backend)
                        result["backend"] = backend
                        result["frames"].append(render_one(scene, output_dir, frames[0], config))
                        break
                    except Exception as exc:
                        failures.append(f"{backend}: {exc}")
                else:
                    raise RuntimeError("所有指定后端的首帧实际渲染均失败；没有回退至 CPU。" + " | ".join(failures))
                result["warnings"].extend(failures)
                atomic_json(result_path, result)
                for frame in frames[1:]:
                    result["frames"].append(render_one(scene, output_dir, frame, config))
                    result["timings"]["elapsed_seconds"] = round(time.monotonic() - start, 3)
                    atomic_json(result_path, result)
        else:
            raise ValueError(f"未知任务：{action}")
        result["ok"] = True
    except Exception as exc:
        result["error"] = str(exc)
        result["error_type"] = type(exc).__name__
        traceback.print_exc()
    finally:
        result["timings"]["elapsed_seconds"] = round(time.monotonic() - start, 3)
        atomic_json(result_path, result)
    if not result["ok"]:
        raise RuntimeError(result.get("error", "渲染失败"))


if __name__ == "__main__":
    main()