import hashlib
import importlib.metadata
import json
import os
import platform
import sys
import time
import uuid
from pathlib import Path

ROOT = Path("/data")
UPSTREAM = Path("/opt/space/upstream")
MODEL_ROOT = ROOT / "models/make-it-animatable"
WEIGHTS = {"bw.pth":424693698, "bw_normal.pth":429008352,
           "joints.pth":525833771, "joints_coarse.pth":424706373,
           "pose.pth":525848433}

def configure_storage():
    if not ROOT.is_dir():
        raise RuntimeError("/data is not mounted. Mount the existing persistent bucket at /data.")
    for name in ("cache/huggingface","cache/torch","cache/xdg","cache/pip",
                 "models","inputs","outputs","temp","logs"):
        (ROOT/name).mkdir(parents=True,exist_ok=True)
    os.environ.update(HF_HOME=str(ROOT/"cache/huggingface"),
                      HF_HUB_CACHE=str(ROOT/"cache/huggingface/hub"),
                      TORCH_HOME=str(ROOT/"cache/torch"),
                      XDG_CACHE_HOME=str(ROOT/"cache/xdg"),
                      PIP_CACHE_DIR=str(ROOT/"cache/pip"),
                      TMPDIR=str(ROOT/"temp"),
                      GRADIO_TEMP_DIR=str(ROOT/"temp/gradio"))
    probe = ROOT/"temp"/("write-check-"+uuid.uuid4().hex)
    with probe.open("x") as stream:
        stream.write("ok")
    if probe.read_text() != "ok":
        raise RuntimeError("/data read/write verification failed")
    probe.unlink()
    import tempfile
    tempfile.tempdir = str(ROOT/"temp")
    missing = [str(MODEL_ROOT/"output/best/new"/name) for name,size in WEIGHTS.items()
               if not (MODEL_ROOT/"output/best/new"/name).is_file()
               or (MODEL_ROOT/"output/best/new"/name).stat().st_size != size]
    bones = MODEL_ROOT/"data/Mixamo/bones.fbx"
    if not bones.is_file() or bones.stat().st_size == 0:
        missing.append(str(bones))
    if missing:
        raise RuntimeError("Required cached assets missing or incomplete; startup stopped without "
                           "overwriting or redownloading assets: "+", ".join(missing))
    for name in ("data","output"):
        target = UPSTREAM/name
        expected = MODEL_ROOT/name
        if target.is_symlink():
            if target.resolve() != expected.resolve():
                raise RuntimeError("Unexpected asset link: "+str(target))
        elif target.exists():
            raise RuntimeError("Asset path must be an image-managed symlink: "+str(target))
        else:
            target.symlink_to(expected, target_is_directory=True)
    print("STORAGE_READY /data writable; cached weights reused; downloads=0; installs=0",flush=True)

def startup_checks():
    import torch
    import torch_cluster
    import pytorch3d
    from pytorch3d.ops import sample_farthest_points
    versions = {}
    for name in ("torch","torchvision","torch-cluster","torch-geometric",
                 "torch-scatter","torch-sparse","torch-spline-conv","pytorch3d","numpy","gradio"):
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = "not required / not installed"
    report = dict(python=platform.python_version(),packages=versions,cuda=torch.version.cuda,
                  cuda_available=torch.cuda.is_available(),
                  gpu=torch.cuda.get_device_name(0) if torch.cuda.is_available() else None)
    print("RUNTIME_VERSIONS "+json.dumps(report),flush=True)
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU unavailable. Dedicated GPU required; CPU inference is disabled.")
    if torch.__version__ != "2.1.2+cu121" or torch.version.cuda != "12.1":
        raise RuntimeError("Runtime Torch/CUDA differs from the pinned build")
    x = torch.rand((128,3),device="cuda")
    out = torch_cluster.fps(x,ratio=0.25)
    _, indices = sample_farthest_points(x[None],K=32)
    y = torch.ones((32,32),device="cuda") @ torch.ones((32,32),device="cuda")
    torch.cuda.synchronize()
    if len(out) != 32 or indices.shape != (1,32) or y[0,0].item() != 32:
        raise RuntimeError("CUDA binary smoke test failed")
    report["cuda_execution"] = "PASS: matmul, torch_cluster.fps, pytorch3d FPS"
    persistence = {}
    for name in ("cache/huggingface","cache/torch","cache/xdg","cache/pip","inputs","outputs"):
        marker = ROOT/name/".persistence-check"
        existed = marker.is_file()
        if not existed:
            with marker.open("x") as stream:
                stream.write(uuid.uuid4().hex)
        persistence[name] = {"marker":marker.read_text(),"survived_previous_start":existed}
    report["persistence"] = persistence
    report["cached_weights"] = WEIGHTS
    report["startup_time"] = time.time()
    report["startup_id"] = uuid.uuid4().hex
    report["outputs_present"] = sorted(str(p.relative_to(ROOT)) for p in (ROOT/"outputs").glob("*"))
    previous = ROOT/"logs/last-startup.json"
    if previous.is_file():
        report["previous_startup"] = json.loads(previous.read_text()).get("startup_id")
    (ROOT/"logs"/(report["startup_id"]+".json")).write_text(json.dumps(report,indent=2))
    previous.write_text(json.dumps(report,indent=2))
    print("CUDA_EXECUTION_PASS "+json.dumps(report),flush=True)
    return report