"""Default scene component workflow: Hunyuan 50k -> Studio low-poly -> Studio UV -> Blender bake -> conditional UV-preserving reduction.

Step 01 uses Tools/hunyuan_3d_generate.py, with explicit user upload authorization.
Run steps 02-05 in the LLM3DGraphics virtualenv; no Unity serialized assets are written here.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import subprocess
import sys
from pathlib import Path

sys.dont_write_bytecode = True
ROOT = Path(__file__).resolve().parents[1]
WORKBENCH = Path('D:/UD/AI/XModel/LLM3DGraphics')
BLENDER = Path('D:/Tools/Blender/blender-5.0.0-windows-x64/blender.exe')
STEPS = ('hunyuan_web_highpoly', 'studio_lowpoly', 'studio_uv', 'blender_bake',
         'blender_decimate_preserve_uv_if_over_130_percent')


def read(path: Path) -> dict:
    return json.loads(path.read_text(encoding='utf-8-sig'))


def save(path: Path, data: dict) -> None:
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + '\n', encoding='utf-8')


def sha(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def update(**values) -> None:
    print(json.dumps(values, ensure_ascii=True), flush=True)


def main(default_manifest: str | None = None, argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--manifest', default=default_manifest, required=default_manifest is None)
    parser.add_argument('--part', required=True, help='Component id/type, or ground')
    parser.add_argument('--output-dir', help='Production directory under Tools/output')
    parser.add_argument('--dry-run', action='store_true', help='Print the model plan without network, Blender or writes')
    parser.add_argument('--refine-bake', action='store_true', help='After step 05, rebake onto its retained UV to repair projection seams')
    parser.add_argument('--revise-budget', action='store_true', help='Archive the old receipt when only the reviewed triangle budget changes')
    args = parser.parse_args(argv)
    manifest_path = (ROOT / args.manifest).resolve()
    manifest = read(manifest_path)
    items = [manifest.get('ground', {}), *manifest.get('components', [])]
    part = next((p for p in items if p.get('id', p.get('type')) == args.part), None)
    if part is None:
        parser.error(f'No component {args.part!r} in {manifest_path}')
    import re
    scene_id = manifest.get('sceneId', '')
    if not all(re.fullmatch(r'[A-Za-z][A-Za-z0-9_]*', value) for value in (scene_id, args.part)):
        parser.error('sceneId and component id must be English identifiers')
    if part.get('generationStatus') == 'reused':
        update(part=args.part, status='reused', prefab=part.get('prefab') or part.get('existingPrefab'))
        return 0
    homecamp = scene_id == 'HomeCamp'
    default_out = f'Tools/output/homecamp_studio/{args.part}' if homecamp else f'Tools/output/scene_models/{scene_id}/{args.part}'
    out = (ROOT / (args.output_dir or part.get('modelOutputDir') or default_out)).resolve()
    if not out.is_relative_to((ROOT / 'Tools/output').resolve()):
        parser.error('--output-dir must stay inside Tools/output')
    source = ROOT / (part.get('highpoly') or part.get('highPoly') or
                     (f'Tools/output/homecamp_3d/{args.part}/{args.part}.glb' if homecamp else
                      f'Tools/output/scene_models/{scene_id}/{args.part}/highpoly/{args.part}.glb'))
    source = source.resolve()
    image_path = ROOT / part['designImage']
    target = int(part['targetTriangles'])
    size = int(manifest.get('modelPipeline', {}).get('textureSize', 1024))
    if target <= 0 or size <= 0:
        parser.error('targetTriangles and textureSize must be positive')
    if args.dry_run:
        update(part=args.part, steps=STEPS, highpoly=str(source), outputDir=str(out),
               targetTriangles=target, textureSize=size, decimationComparison='actual * 10 > target * 13',
               preserveStudioUV=True, highpolyCommand=['python', 'Tools/hunyuan_3d_generate.py',
               '--image', str(ROOT / part.get('modelReferenceImage', part['designImage'])),
               '--faces', '50k', '--output', str(source)])
        return 0
    authorization = manifest.get('uploadAuthorization')
    if authorization and authorization.get('status') != 'authorized_by_user':
        raise RuntimeError('The manifest upload authorization is not ready')
    expected_image = next((p for p in manifest.get('designImageAudit', []) if p['id'] == args.part), None)
    if expected_image and sha(image_path) != expected_image['sha256']:
        raise RuntimeError('Design input changed after authorization; review the new input first')
    if not source.is_file():
        raise FileNotFoundError(f'Complete Hunyuan 50k web high mesh first (see --dry-run): {source}')
    sys.path.insert(0, str(ROOT / 'Tools'))
    from hunyuan_3d import auth
    from hunyuan_3d.glb_check import assert_glb_textures_ok
    assert_glb_textures_ok(source)
    sys.path.insert(0, str(WORKBENCH))
    from backend import hunyuan_studio as studio
    from backend.retopology import _run_blender_step, compress_bake_maps
    import trimesh
    sys.path.insert(0, str(ROOT / 'Tools/model_tools'))
    from mt_run import run_blender_script
    from png_compress import compress_dir_pngs

    out.mkdir(parents=True, exist_ok=True)
    receipt_path = out / 'production.json'
    receipt = read(receipt_path) if receipt_path.exists() else {}
    settings = dict(source=source.relative_to(ROOT).as_posix(), sourceSha256=sha(source),
                    targetTriangles=target, textureSize=size)
    if receipt and any(receipt.get(k) != v for k, v in settings.items()):
        if (args.revise_budget and receipt.get('sourceSha256') == settings['sourceSha256']
                and receipt.get('textureSize') == size):
            save(out / ('production_target_' + str(receipt['targetTriangles']) + '.json'), receipt)
            receipt['previousTargetTriangles'] = receipt['targetTriangles']
            receipt['targetTriangles'] = target
            receipt['status'] = 'budget_revised_for_surface_and_uv_integrity'
        else:
            raise RuntimeError('Source/settings changed; keep the old receipt and use a new production directory')
    if receipt.get('status') == 'ready_for_unity' and not args.refine_bake:
        for path, digest in receipt['outputs'].items():
            if sha(ROOT / path) != digest:
                raise RuntimeError(f'Completed output changed: {path}')
        update(part=args.part, status='ready_for_unity', cached=True)
        return 0
    receipt.update(settings, part=args.part, status='studio', designImage=part['designImage'],
                   designSha256=sha(image_path), modelPipeline=list(STEPS))
    receipt.setdefault('steps', {})['01'] = 'highpoly_validated'
    save(receipt_path, receipt)
    shutil.copy2(source, out / 'model.glb')
    studio.is_logged_in = lambda client_id=None: auth.is_logged_in()
    studio.load_storage_state = lambda client_id=None: (read(auth.storage_state_path()), None)
    uv = out / 'model_lowpoly_uv.glb'
    if uv.is_file():
        studio.validate_uv_glb(uv)
    else:
        meta_path = out / 'hunyuan_studio.json'
        meta = read(meta_path) if meta_path.exists() else {}
        studio.process_existing(source, out, update, target_faces=target,
                                resume_low_id=meta.get('lowpoly_works_id'), resume_uv_id=meta.get('uv_works_id'))
    receipt.update(status='baking', studio=read(out / 'hunyuan_studio.json'))
    receipt['steps'].update({'02': 'studio_lowpoly_completed', '03': 'studio_uv_completed'})
    save(receipt_path, receipt)

    def blender(script: str, script_args: list[str], log_name: str) -> None:
        with (out / log_name).open('w', encoding='utf-8') as log:
            subprocess.run([str(BLENDER), '-b', '-t', '2', '--factory-startup', '--python-exit-code', '1', '--python',
                            str(ROOT / 'Tools/model_tools' / script), '--', *script_args],
                           check=True, stdout=log, stderr=subprocess.STDOUT, timeout=240)

    aligned = out / 'model_lowpoly_uv_aligned.glb'
    if not aligned.is_file():
        blender('bl_align_studio_uv.py', [str(out)], 'alignment.log')
    if not receipt.get('bakeAligned'):
        _run_blender_step(out, str(BLENDER), 'blender_lowpoly_bake.py',
                          (out / 'model.glb', out / 'model_baked.glb', 1, size, aligned),
                          'bake.log', marker='Lowpoly bake saved to')
        compress_bake_maps(out)
        receipt['bakeAligned'] = True
        receipt['steps']['04'] = 'blender_bake_completed_on_studio_uv'
        save(receipt_path, receipt)
    receipt['steps']['04'] = 'blender_bake_completed_on_studio_uv'
    shutil.copy2(out / 'model_baked.glb', out / 'model_ex.glb')
    mesh = trimesh.load(out / 'model_ex.glb', force='mesh', process=False)
    reduction = dict(triggered=False, faces_before=len(mesh.faces), faces_after=len(mesh.faces), target_faces=target, threshold_percent=30)
    if len(mesh.faces) * 10 > target * 13:
        update(stage='Blender welded geometry / preserved corner UV reduction', part=args.part)
        blender('bl_decimate_homecamp_uv.py', [str(out / 'model_ex.glb'), str(out / 'model_decimated_welded.glb'), str(target)], 'decimate_welded.log')
        reduction = read(out / 'model_decimated_welded.json')
        shutil.copy2(out / 'model_decimated_welded.glb', out / 'model_ex.glb')
        mesh = studio.validate_uv_glb(out / 'model_ex.glb')
    if len(mesh.faces) * 10 > target * 13:
        raise RuntimeError('Final mesh still exceeds the target by more than 30 percent')
    receipt['steps']['05'] = 'uv_preserving_decimation_completed' if reduction['triggered'] else 'skipped_within_130_percent'
    receipt.update(triangles=len(mesh.faces), vertices=len(mesh.vertices), decimation=reduction,
                   bakeAlignment=read(out / 'bake_alignment.json'))
    save(receipt_path, receipt)
    if homecamp:
        blender('bl_repair_homecamp_surface.py', [str(out / 'model_ex.glb'), str(out / 'model_surface_repaired.glb'), args.part], 'surface_repair.log')
        receipt['surfaceRepair'] = read(out / 'model_surface_repaired.json')
        shutil.copy2(out / 'model_surface_repaired.glb', out / 'model_ex.glb')
        receipt['triangles'] = receipt['surfaceRepair']['triangles']
    if args.refine_bake:
        # Collapse retains charts, but coarse triangles can move the sampled
        # texels. Correct the projection on the final mesh, without re-unwrap.
        original_final = out / 'model_decimated_before_rebake.glb'
        shutil.copy2(out / 'model_ex.glb', original_final)
        _run_blender_step(out, str(BLENDER), 'blender_lowpoly_bake.py',
                          (out / 'model.glb', out / 'model_refined.glb', 1, size, original_final),
                          'refine_bake.log', marker='Lowpoly bake saved to')
        compress_bake_maps(out)
        refined = studio.validate_uv_glb(out / 'model_refined.glb')
        if len(refined.faces) != receipt['triangles']:
            raise RuntimeError('Refinement changed the final mesh topology')
        shutil.copy2(out / 'model_refined.glb', out / 'model_ex.glb')
        receipt['projectionRefinement'] = 'Same source high -> final reduced low; UV preserved, no unwrap'
    blender('bl_prepare_studio_trial.py', [str(out), args.part], 'static_export.log')
    receipt['xpbr'] = run_blender_script(str(BLENDER), 'bl_convert.py',
        [str(out / 'static' / (args.part + '.fbx')), str(out / 'xpbr'), str(size)], blender_args=('-t', '2'))
    compress_dir_pngs(str(out / 'xpbr'), suffixes=('_base.png', '_mix.png'))
    # Only these three source files are imported. Unity owns .mat and .meta generation.
    outputs = [out / 'xpbr' / (args.part + suffix) for suffix in ('.fbx', '_base.png', '_mix.png')]
    receipt.update(status='ready_for_unity', staticExport=read(out / 'static_export.json'),
                   outputs={p.relative_to(ROOT).as_posix(): sha(p) for p in outputs})
    save(receipt_path, receipt)
    update(part=args.part, status=receipt['status'], triangles=receipt['triangles'])
    return 0


if __name__ == '__main__':
    raise SystemExit(main())