#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""comfyui-AICG3D 修复脚本（2026-09-27 实测有效）

修三个已在实机验证的问题：

  1. __init__.py 语法错误 —— 仓库 main 分支 2026-09-25 那次提交
     "Update README ... and professional console logs" 把 5 行 print 提出了 try 块，
     导致插件在 ComfyUI 启动时直接 SyntaxError，整套节点无法加载
     （GitHub issue #6 即此问题，社区只能回退到 b376ea8）。
  2. h3easy/nodes.py 顶部把真实存在的 h3easy/sampling_strategies.py 用
     "Mock for missing sampling_strategies.py" 顶掉了（一共 3 处），
     而该文件与上游字节一致、就在仓库里 → SelfLift 采样策略实际失效，
     连上 SelfLift 节点时 sample_with_sampling_plan() 返回 None，渲染节点输出空 latent。
  3. （可选 --with-gh）把合并进来却永远注册不上的 4 个 Goohai 节点
     （MiniMaxH3IntegrationGH / ...AdapterGH / DualClockT8GH / AVDecodeT8GH）
     补进注册表。实验性：本插件 __init__.py 会走 V1 分支并 return，
     ComfyUI 0.37 的 comfy_entrypoint(V3) 分支因此永远不被调用。

用法：
    python3 apply_fixes.py <comfyui-AICG3D 目录> [--with-gh] [--dry-run]
"""
from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

MOCK_RE = re.compile(
    r"#\s*-*\s*(?:AICG3D\s+)?Mock for missing sampling_strategies\.py\s*-*\s*\n"
    r'SAMPLING_PLAN_TYPE = "default"\n'
    r'SELFLIFT_KIND = "none"\n'
    r"class MiniMaxH3SamplingPlan:\n"
    r"    def __init__\(self, \*args, \*\*kwargs\): pass\n"
    r"def sample_with_sampling_plan\(\*args, \*\*kwargs\):\n"
    r"    return None\n"
    r"(?:#\s*-+[^\n]*\n)?"
)

REAL_IMPORT = """from h3easy.sampling_strategies import (
    SAMPLING_PLAN_TYPE,
    SELFLIFT_KIND,
    MiniMaxH3SamplingPlan,
    sample_with_plan as sample_with_sampling_plan,
)
"""

GH_NAMES = (
    "MiniMaxH3IntegrationGH",
    "MiniMaxH3IntegrationAdapterGH",
    "MiniMaxH3DualClockT8GH",
    "MiniMaxH3AVDecodeT8GH",
)

GH_PATCH = '''
    # --- 补注册被漏掉的 Goohai 节点（apply_fixes.py --with-gh）---
    try:
        from h3goohai import nodes as _gh_nodes
        for _gh_name in %r:
            _gh_cls = getattr(_gh_nodes, _gh_name, None)
            if _gh_cls is None:
                continue
            _gh_schema = _gh_cls.GET_SCHEMA()
            mapping[_gh_schema.node_id] = _gh_cls
            if _gh_schema.display_name:
                display[_gh_schema.node_id] = _gh_schema.display_name
    except Exception as _gh_err:
        print("[AICG3D] Goohai 节点补注册跳过:", _gh_err)
''' % (GH_NAMES,)


def fix_init(root: Path, dry: bool) -> str:
    p = root / "__init__.py"
    if not p.exists():
        return "  [1] 跳过：找不到 __init__.py"
    lines = p.read_text(encoding="utf-8").splitlines(True)
    in_try = False
    fixed = []
    n = 0
    for line in lines:
        stripped = line.lstrip()
        if stripped.startswith("try:"):
            in_try = True
        elif in_try and stripped.startswith(("except", "finally")):
            in_try = False
        if in_try and line.startswith("print("):
            line = "    " + line
            n += 1
        fixed.append(line)
    if n == 0:
        return "  [1] __init__.py 已正常（未发现顶格的 print），无需修复"
    if not dry:
        p.write_text("".join(fixed), encoding="utf-8")
    return f"  [1] __init__.py 修复 {n} 行缩进{'（试运行未写入）' if dry else ''}"


def fix_mock(root: Path, dry: bool) -> str:
    p = root / "h3easy" / "nodes.py"
    real = root / "h3easy" / "sampling_strategies.py"
    if not p.exists():
        return "  [2] 跳过：找不到 h3easy/nodes.py"
    src = p.read_text(encoding="utf-8")
    found = MOCK_RE.findall(src)
    if len(found) == 0:
        return "  [2] 未发现 mock 块（可能已修过或用的是修复版）"
    if not real.exists():
        return "  [2] 警告：仓库内没有 h3easy/sampling_strategies.py，保持 mock 不动"
    new = MOCK_RE.sub("", src)
    if "from h3easy.sampling_strategies import" not in new:
        new = new.replace("import torch\n", REAL_IMPORT + "import torch\n", 1)
    if not dry:
        p.write_text(new, encoding="utf-8")
    return f"  [2] 移除 {len(found)} 处 mock，接入真实 sampling_strategies{'（试运行未写入）' if dry else ''}"


def fix_gh(root: Path, dry: bool) -> str:
    p = root / "aicg3d" / "registry.py"
    if not p.exists():
        return "  [3] 跳过：找不到 aicg3d/registry.py"
    src = p.read_text(encoding="utf-8")
    if "h3goohai" in src:
        return "  [3] registry.py 已包含 Goohai 补注册，跳过"
    anchor = "    for node_id, node_class in mapping.items():"
    if anchor not in src:
        return "  [3] 找不到插入锚点，跳过（请手工处理）"
    if not dry:
        p.write_text(src.replace(anchor, GH_PATCH + anchor, 1), encoding="utf-8")
    return f"  [3] 已在 registry.py 插入 Goohai 4 节点补注册（实验性）{'（试运行未写入）' if dry else ''}"


def main() -> int:
    ap = argparse.ArgumentParser(description="comfyui-AICG3D 修复脚本")
    ap.add_argument("plugin_dir", type=Path, help="插件目录，例如 ComfyUI/custom_nodes/comfyui-AICG3D")
    ap.add_argument("--with-gh", action="store_true", help="同时补注册 4 个 Goohai 节点（实验性）")
    ap.add_argument("--dry-run", action="store_true", help="只报告不写入")
    a = ap.parse_args()

    root: Path = a.plugin_dir.resolve()
    if not (root / "__init__.py").exists():
        print(f"× {root} 看起来不是 comfyui-AICG3D 目录")
        return 2
    print(f"修复目标：{root}{'（试运行）' if a.dry_run else ''}")
    print(fix_init(root, a.dry_run))
    print(fix_mock(root, a.dry_run))
    if a.with_gh:
        print(fix_gh(root, a.dry_run))

    import py_compile

    bad = []
    for f in [root / "__init__.py", root / "h3easy" / "nodes.py", root / "aicg3d" / "registry.py"]:
        if not f.exists():
            continue
        try:
            py_compile.compile(str(f), doraise=True, cfile="/tmp/_chk.pyc")
        except Exception as e:
            bad.append(f"{f.name}: {str(e).splitlines()[0]}")
    print("语法检查：" + ("全部通过 ✔" if not bad else "仍有错误 -> " + "; ".join(bad)))
    print("下一步：重启 ComfyUI，并在浏览器里 Ctrl+F5 强刷。")
    return 0 if not bad else 1


if __name__ == "__main__":
    sys.exit(main())