"""Durable, integrity-checked geometry handoffs for explicit texture requests."""
import hashlib
import json
from pathlib import Path
import shutil

from hf_space import ovoxel_results
from hf_space.split_pipeline import CHECKPOINT_FILES, OPTIONAL_CHECKPOINT_FILES, validate_checkpoint

HANDOFF_FILES = ('input.png', 'geometry.npz',
                 *(f'geometry-checkpoint/{name}' for name in CHECKPOINT_FILES))


def _handoff_files(folder):
    metadata = validate_checkpoint(folder / 'geometry-checkpoint')
    return (*HANDOFF_FILES, *(f'geometry-checkpoint/{name}' for name in OPTIONAL_CHECKPOINT_FILES
                             if name in metadata['files']))


def _digest(path):
    digest = hashlib.sha256()
    with path.open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            digest.update(block)
    return digest.hexdigest()


def _file(root, name):
    path = root / name
    if (root.is_symlink() or path.is_symlink() or path.parent.is_symlink()
            or not path.resolve().is_relative_to(root.resolve()) or not path.is_file()
            or path.stat().st_size <= 0):
        raise ValueError('The saved geometry is incomplete or invalid.')
    return path


def persist(source, report):
    """Publish the white preview and its exact texturing input; ready is last."""
    source = Path(source)
    folder = source.parent
    identifier = folder.name
    destination = ovoxel_results.directory(identifier)
    if source.name != 'geometry.npz' or destination.exists():
        raise ValueError('The geometry output destination is invalid.')
    files, metadata = ovoxel_results.local_manifest(folder)
    expected_frames = ovoxel_results._frame_names(metadata)
    if metadata.get('modes') != ['clay'] or metadata.get('default_mode') != 'clay':
        raise ValueError('Geometry must have a clay preview.')
    if set(files) != {f'frames/{name}' for name in expected_frames}:
        raise ValueError('The geometry preview is incomplete.')
    # No user-selected path or archive member is accepted as a handoff file.
    handoff_files = _handoff_files(folder)
    names = (*handoff_files, 'report.json', 'preview.json', 'timings.jsonl',
             *(f'views/{name}' for name in expected_frames))
    for name in names:
        _file(folder, name)
    destination.mkdir(parents=True)
    bindings = {}
    for name in names:
        original = _file(folder, name)
        target = destination / name
        target.parent.mkdir(parents=True, exist_ok=True)
        expected = {'size': original.stat().st_size, 'sha256': _digest(original)}
        shutil.copyfile(original, target)
        if target.stat().st_size != expected['size'] or _digest(target) != expected['sha256']:
            raise ValueError('The saved geometry failed integrity validation.')
        if name in handoff_files:
            bindings[name] = expected
    (destination / 'geometry-manifest.json').write_text(json.dumps({
        'schema_version': 1, 'files': bindings}), encoding='utf-8')
    (destination / 'display-manifest.json').write_text(json.dumps({
        'files': files, 'preview_metadata': metadata}), encoding='utf-8')
    (destination / 'ready.json').write_text(json.dumps({
        'complete': True, 'identifier': identifier, 'output_kind': 'geometry',
        'project_revision': report['project_revision']}), encoding='utf-8')
    return identifier


def restore(identifier, destination):
    """Restore only a verified geometry handoff, without allocating a GPU."""
    folder = ovoxel_results.directory(identifier)
    ready = json.loads(_file(folder, 'ready.json').read_text(encoding='utf-8'))
    if ready.get('complete') is not True or ready.get('output_kind') != 'geometry' or ready.get('identifier') != identifier:
        raise ValueError('The saved geometry is not ready for texturing.')
    manifest = json.loads(_file(folder, 'geometry-manifest.json').read_text(encoding='utf-8'))
    bindings = manifest.get('files')
    handoff_files = _handoff_files(folder)
    if manifest.get('schema_version') != 1 or not isinstance(bindings, dict) or set(bindings) != set(handoff_files):
        raise ValueError('The saved geometry handoff is invalid.')
    destination = Path(destination)
    if destination.is_symlink() or destination.exists() and any(destination.iterdir()):
        raise ValueError('Geometry must be restored into an empty local directory.')
    destination.mkdir(parents=True, exist_ok=True)
    for name in handoff_files:
        source = _file(folder, name)
        expected = bindings[name]
        if (not isinstance(expected, dict) or source.stat().st_size != expected.get('size')
                or _digest(source) != expected.get('sha256')):
            raise ValueError('The saved geometry failed integrity validation.')
        target = destination / name
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(source, target)
        if target.stat().st_size != expected['size'] or _digest(target) != expected['sha256']:
            raise ValueError('The restored geometry failed integrity validation.')
    return destination