"""Run the LittleBear fixed-skeleton skin-only job in the existing DGXA container."""
from pathlib import Path
import hashlib
import json
import subprocess

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "Tools/output/littlebear_studio/rig_mixamo41"
JOB_NAME = "littlebear_mixamo41_20261026_a1"
REMOTE = "/home/spark/SkinTokens/tmp_jobs/" + JOB_NAME
CONTAINER = "/workspace/SkinTokens/tmp_jobs/" + JOB_NAME


def run(args, capture=False):
    print("RUN", args[0], flush=True)
    return subprocess.run(args, check=True, capture_output=capture, text=capture,
                          encoding="utf-8" if capture else None)


def sha(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main():
    source = OUT / "source.glb"
    if not source.is_file():
        raise SystemExit("Prepare the fixed Mixamo skeleton and source.glb first")
    script = OUT / "run_tokenrig.sh"
    script.write_text("\n".join([
        "#!/usr/bin/env bash", "set -euo pipefail",
        "export PATH=/opt/py310/bin:$PATH",
        "export PYTHONPATH=/workspace/SkinTokens:/workspace/blender_dev/build_bpy/bin",
        "export BLENDER_SYSTEM_SCRIPTS=/workspace/blender_dev/blender/scripts",
        "export BLENDER_SYSTEM_DATAFILES=/workspace/blender_dev/blender/release/datafiles",
        "export PYTHONUNBUFFERED=1",
        "python3 demo.py --input " + CONTAINER + "/source.glb --output " + CONTAINER + "/rigged_raw.glb --use_skeleton --use_transfer --skeleton_template original --rig_cls vroid",
        "" ]), encoding="utf-8", newline="\n")
    ssh = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=8", "DGXA"]
    run(ssh + ["mkdir -p " + REMOTE])
    existing = run(ssh + ["if test -s " + REMOTE + "/rigged_raw.glb; then sha256sum " + REMOTE + "/source.glb; fi"], capture=True).stdout.split()
    if existing and existing[0] != sha(source):
        raise RuntimeError("Existing job belongs to a different input; use a new job name")
    if not existing:
        run(["scp", str(source), str(script), "DGXA:" + REMOTE + "/"])
    remote_hash = run(ssh + ["sha256sum " + REMOTE + "/source.glb"], capture=True).stdout.split()[0]
    if remote_hash != sha(source):
        raise RuntimeError("Remote input checksum mismatch")
    if not existing:
        run(ssh + ["docker exec -w /workspace/SkinTokens tokenrig bash " + CONTAINER + "/run_tokenrig.sh > " + REMOTE + "/tokenrig.log 2>&1"])
    for filename in ["rigged_raw.glb", "tokenrig.log"]:
        run(["scp", "DGXA:" + REMOTE + "/" + filename, str(OUT / filename)])
    remote_output = run(ssh + ["sha256sum " + REMOTE + "/rigged_raw.glb"], capture=True).stdout.split()[0]
    if remote_output != sha(OUT / "rigged_raw.glb"):
        raise RuntimeError("Remote output checksum mismatch")
    (OUT / "remote_receipt.json").write_text(json.dumps(dict(job=JOB_NAME,
        input_sha256=remote_hash, output_sha256=remote_output,
        flags=["--use_skeleton", "--use_transfer", "--skeleton_template", "original", "--rig_cls", "vroid"],
        reference_weights_used=False), indent=2), encoding="utf-8")
    print("TOKENRIG_WEIGHTS_READY", flush=True)


if __name__ == "__main__":
    main()