"""Reuse curated examples or validated inference results before requesting a GPU."""
import hashlib
import json
import os
from pathlib import Path
import shutil
import time
import uuid

from hf_space import example_assets

ROOT = Path(__file__).resolve().parents[1]
EXAMPLE_IMAGE = ROOT/'assets/pitcher_plant.png'
CACHE = Path('/data/generated/pbr-results-v1')
EXAMPLE_NAMES = example_assets.EXAMPLE_NAMES


def revision():
    return os.environ.get('SUPRATEX_MODEL_REVISION', '499db28a11b614873a2cfbcfd1990024ba0f1d34')


def image_digest(image):
    rgba = image.convert('RGBA')
    digest = hashlib.sha256(f'RGBA:{rgba.width}:{rgba.height}:'.encode())
    digest.update(rgba.tobytes())
    return digest.hexdigest()


def cache_key(image, seed):
    request = {'input_rgba_sha256': image_digest(image), 'seed': int(seed),
               'project_revision': revision(), 'profile': 'torch28-sdpa-pbr4096-v1',
               'project_repository': os.environ.get('SUPRATEX_MODEL_REPO', 'supra3d/supra-tex')}
    return hashlib.sha256(json.dumps(request, sort_keys=True).encode()).hexdigest(), request


def _lookup_saved(key, request):
    try:
        record = json.loads((CACHE/(key+'.json')).read_text())
        model = CACHE/(key+'.glb')
        if record.get('request') == request and example_assets.valid_model(model, record['sha256'], record['glb_bytes']):
            return str(model)
    except (OSError, ValueError, KeyError):
        pass
    return None


def default_preview():
    return example_assets.lookup_name('pitcher_plant')


def lookup(image, seed):
    if image is None:
        return None
    name = example_assets.identify(image, seed)
    if name:
        # Never fall back to a superseded generated example.
        return example_assets.lookup_name(name)
    key, request = cache_key(image, seed)
    return _lookup_saved(key, request)


def prepare_examples():
    example_assets.prepare()


def examples_status():
    return example_assets.status()


def store(image, seed, source, report):
    key, request = cache_key(image, seed)
    CACHE.mkdir(parents=True, exist_ok=True)
    model = CACHE/(key+'.glb')
    temporary = CACHE/(key+'.'+uuid.uuid4().hex+'.part')
    try:
        shutil.copyfile(source, temporary)
        if not example_assets.valid_model(temporary, report['sha256'], report['glb_bytes']):
            raise ValueError('PBR result cache checksum mismatch')
        temporary.replace(model)
    finally:
        temporary.unlink(missing_ok=True)
    record = {'request': request, 'sha256': report['sha256'], 'glb_bytes': report['glb_bytes'],
              'source_job': report.get('job_directory'), 'created_unix': time.time()}
    temporary = CACHE/(key+'.'+uuid.uuid4().hex+'.json.part')
    try:
        temporary.write_text(json.dumps(record, indent=2))
        temporary.replace(CACHE/(key+'.json'))
    finally:
        temporary.unlink(missing_ok=True)
    return str(model)