"""Generate one UniMate motion from an existing rig, preserving caption/facing.

Project adapter only: does not modify the external LLM3DGraphics checkout.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
from pathlib import Path
import sys
import uuid


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--rig', type=Path, required=True)
    parser.add_argument('--out', type=Path, required=True)
    parser.add_argument('--face-right', required=True)
    parser.add_argument('--face-left', required=True)
    parser.add_argument('--motion', choices=('Idle','Walk','Attack'), default='Walk')
    parser.add_argument('--duration', type=float, default=2)
    parser.add_argument('--device', choices=('cuda','cpu'), default='cuda',
                        help='CPU fallback preserves other DGXA services when GPU memory is full')
    parser.add_argument('--prompt')
    args = parser.parse_args()
    args.prompt = args.prompt or {
        'Idle': 'a small bipedal character stands upright and idles with gentle breathing',
        'Walk': 'a small bipedal character walks forward upright with alternating short steps and gentle arm swings',
        'Attack': 'a small bipedal character performs one forward body slam attack and recovers to standing',
    }[args.motion]
    root = Path(os.environ.get('LLM3D_ROOT', r'D:\UD\AI\XModel\LLM3DGraphics'))
    sys.path.insert(0, str(root))
    import backend.unimate_runner as runner
    args.out.mkdir(parents=True, exist_ok=True)
    stamp=uuid.uuid4().hex
    working=args.out/'generations'/stamp
    working.mkdir(parents=True,exist_ok=False)
    rig = args.rig.resolve()
    if not rig.is_file():
        raise SystemExit(f'Missing rig: {rig}')
    # The upstream English normalizer replaces all walk captions with "a dog";
    # the Chinese old caption contained 跳 and became a jump. Keep this caption.
    runner.normalize_unimate_prompt = lambda caption: caption.strip()
    original_cfg = runner._cfg
    helper_name = 'aigame_' + args.motion.lower() + '_' + uuid.uuid4().hex[:8] + '_helpers'
    runner._cfg = lambda: {**original_cfg(), 'unimate_helpers_dir': original_cfg()['unimate_jobs_dir'] + '/' + helper_name}
    original_scp = runner._scp_to
    def upload(local, remote, host, timeout=600):
        if remote.endswith('/run_job.sh'):
            text = Path(local).read_text(encoding='utf-8')
            marker = 'face_r, face_l = _detect_face_pair(arm)'
            if marker not in text:
                raise RuntimeError('Upstream facing adapter changed; inspect before generating.')
            text = text.replace(marker, f'face_r, face_l = {args.face_right!r}, {args.face_left!r}')
            config_marker='cfg["dataset"]["dataset_list"] = ["objaverse"]'
            if config_marker not in text:
                raise RuntimeError('Upstream inference configuration changed; inspect device adapter.')
            text=text.replace(config_marker,config_marker+'\ncfg["sampling"]["device"] = '+repr(args.device))
            local = working / 'run_motion_job.sh'
            local.write_text(text, encoding='utf-8', newline='\n')
        return original_scp(local, remote, host, timeout)
    runner._scp_to = upload
    original_ssh = runner._ssh
    def offline(host, command, timeout=120):
        if 'run_job.sh' in command:
            command = 'export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1; ' + command
        return original_ssh(host, command, timeout)
    runner._ssh = offline
    provenance = dict(rig=str(rig), rig_sha256=hashlib.sha256(rig.read_bytes()).hexdigest(),
                      prompt=args.prompt, face_right=args.face_right, face_left=args.face_left,
                      motion=args.motion,duration=args.duration,device=args.device)
    (working / 'motion_request.json').write_text(json.dumps(provenance, indent=2), encoding='utf-8')
    print(json.dumps(provenance), flush=True)
    result = runner.submit_unimate_job(rig_path=rig, prompt=args.prompt,
        directory=working.resolve(), log_path=working.resolve() / 'pipeline.log',
        job_id='aigame_' + args.motion.lower() + '_' + uuid.uuid4().hex, duration_sec=args.duration)
    provenance['source_fbx_sha256']=hashlib.sha256((working/'animated.fbx').read_bytes()).hexdigest()
    provenance['status']='generated'
    provenance['remote_job']=result['remote_job']
    # Promote artifacts only after successful inference; commit receipt last.
    # A failed generation must never relabel an older cached FBX as a new result.
    for file in list(working.iterdir()):
        if file.is_file() and file.name!='motion_request.json':
            os.replace(file,args.out/file.name)
    receipt=working/'motion_request.json'
    receipt.write_text(json.dumps(provenance,indent=2),encoding='utf-8')
    os.replace(receipt,args.out/'motion_request.json')
    if args.motion=='Walk':
        (args.out/'walk_request.json').write_text(json.dumps(provenance,indent=2),encoding='utf-8')
    print(json.dumps(result), flush=True)


if __name__ == '__main__':
    main()