"""Studio low-poly trial for any character (generalized from greathorn_studio_trial.py), steps 02-05 of the official flow:

02 Hunyuan 3D Studio low-poly topology -> 03 Studio UV unwrapping -> align to the exact source high -> 04 Blender bake
(diffuse/normal/AO/metallic/roughness) onto the Studio UVs -> 05 UV-preserving decimation only when over target x1.3.

The original prefab, controller, clips and rig are hash-snapshotted and must stay unchanged; the Studio candidate is
packaged and rebound to the original skeleton by studio_lowpoly_finish.py.

Run with the 3D workbench environment (it provides the Studio integration):
    D:/UD/AI/XModel/LLM3DGraphics/.venv/Scripts/python.exe Tools/studio_lowpoly_trial.py --actor RockBall --kind Animal \
        --source Tools/output/rockball_3d/rockball.glb --rig Tools/output/rockball_dgx/high_rig/rigged.glb --prefab RockBall_Rigged.prefab
"""
import sys
sys.dont_write_bytecode = True
import argparse
import hashlib
import json
import shutil
import subprocess
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
STUDIO = Path('D:/UD/AI/XModel/LLM3DGraphics')
BLENDER = 'D:/Tools/Blender/blender-5.0.0-windows-x64/blender.exe'
sys.path.insert(0, str(ROOT / 'Tools'))
from hunyuan_3d import auth
sys.path.insert(0, str(STUDIO))
from backend import hunyuan_studio as studio
from backend.retopology import _run_blender_step, compress_bake_maps, limit_studio_faces


def sha(path): return hashlib.sha256(path.read_bytes()).hexdigest()
def update(**values): print(json.dumps(values, ensure_ascii=True), flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--actor', required=True)
    ap.add_argument('--kind', default='Animal', choices=('Animal', 'Enemy'))
    ap.add_argument('--source', type=Path, required=True, help='the exact high model the original rig was made from')
    ap.add_argument('--rig', type=Path, required=True)
    ap.add_argument('--prefab', required=True, help='original animated prefab file name inside the actor folder')
    ap.add_argument('--tris', type=int, default=2000)
    ap.add_argument('--size', type=int, default=1024)
    a = ap.parse_args()
    out = ROOT / 'Tools/output/studio_trial' / a.actor
    source = (ROOT / a.source).resolve()
    out.mkdir(parents=True, exist_ok=True)
    receipt = out / 'trial.json'
    current = json.loads(receipt.read_text(encoding='utf-8')) if receipt.exists() else {}
    digest = sha(source)
    if current and current.get('sourceSha256') != digest: raise ValueError('Trial source changed')
    original = ROOT / f'AIClient/Assets/Bundles/Characters/{a.kind}/{a.actor}'
    originals = [original / a.prefab, original / f'{a.actor}_Animator.controller', *[original / 'Clips' / f'{n}.anim' for n in ('Idle', 'Walk', 'Attack')], ROOT / a.rig]
    snapshot = {str(p.relative_to(ROOT)): sha(p) for p in originals}
    if current.get('originals') and current['originals'] != snapshot: raise ValueError('Original skeleton/animations changed during trial')
    current.update(actor=a.actor, kind=a.kind, prefab=a.prefab, rig=str(a.rig).replace('\\', '/'), source=str(source.relative_to(ROOT)).replace('\\', '/'),
                   sourceSha256=digest, targetTriangles=a.tris, textureSize=a.size, originals=snapshot, status='studio')
    receipt.write_text(json.dumps(current, indent=2), encoding='utf-8')
    shutil.copy2(source, out / 'model.glb')
    # Reuse this project's authorised Hunyuan session through the installed Studio integration; credentials stay local.
    studio.is_logged_in = lambda client_id=None: auth.is_logged_in()
    studio.load_storage_state = lambda client_id=None: (json.loads(auth.storage_state_path().read_text(encoding='utf-8')), None)
    uv = out / 'model_lowpoly_uv.glb'
    if uv.exists(): studio.validate_uv_glb(uv)
    else:
        meta = json.loads((out / 'hunyuan_studio.json').read_text(encoding='utf-8')) if (out / 'hunyuan_studio.json').exists() else {}
        studio.process_existing(source, out, update, target_faces=a.tris, resume_low_id=meta.get('lowpoly_works_id'), resume_uv_id=meta.get('uv_works_id'))
    current.update(status='baking', studio=json.loads((out / 'hunyuan_studio.json').read_text(encoding='utf-8')))
    receipt.write_text(json.dumps(current, indent=2), encoding='utf-8')
    aligned = out / 'model_lowpoly_uv_aligned.glb'
    if not aligned.exists():
        # Studio centres/normalises its mesh; align it to the source high before any ray projection.
        with (out / 'alignment.log').open('w', encoding='utf-8') as log:
            subprocess.run([BLENDER, '-b', '--factory-startup', '--python', str(ROOT / 'Tools/model_tools/bl_align_studio_uv.py'), '--', str(out)], stdout=log, stderr=subprocess.STDOUT, check=True, timeout=120)
    if not current.get('bakeAligned'):
        _run_blender_step(out, BLENDER, 'blender_lowpoly_bake.py', (out / 'model.glb', out / 'model_baked.glb', 1, a.size, aligned), 'bake.log', marker='Lowpoly bake saved to')
        compress_bake_maps(out)
        current['bakeAligned'] = True
    shutil.copy2(out / 'model_baked.glb', out / 'model_ex.glb')
    mesh, decimation = limit_studio_faces(out, a.tris, BLENDER, update)
    current.update(status='mesh_ready', triangles=len(mesh.faces), vertices=len(mesh.vertices), decimation=decimation,
                   bakeAlignment=json.loads((out / 'bake_alignment.json').read_text(encoding='utf-8')), finalSha256=sha(out / 'model_ex.glb'))
    receipt.write_text(json.dumps(current, indent=2), encoding='utf-8')
    update(status='mesh_ready', triangles=len(mesh.faces), decimation=decimation)


if __name__ == '__main__': main()