"""Resolve the original released UMT5 assets without loading an archive in RAM."""
import os
from pathlib import Path
import shutil
import zipfile

from filelock import FileLock
from huggingface_hub import hf_hub_download

REPO = 'AlayaLab/FloodDiffusion2'
REVISION = '2402033aff49850b31c1c5be9179c2fd7cb476cb'
PREFIX = 'deps/t5_umt5-xxl-enc-bf16/'
FILES = {
    'models_t5_umt5-xxl-enc-bf16.pth': (11361920418, 0x201263e2),
    'google/umt5-xxl/special_tokens_map.json': (6623, 0x52b4332b),
    'google/umt5-xxl/spiece.model': (4548313, 0x3d61aeda),
    'google/umt5-xxl/tokenizer.json': (16837417, 0xdf89c52f),
    'google/umt5-xxl/tokenizer_config.json': (61728, 0x882574af),
}


def resolve_t5():
    override = os.getenv('FLOOD2_T5_DIR')
    if override:
        directory = Path(override).resolve()
        for name, (size, _) in FILES.items():
            if (directory/name).stat().st_size != size:
                raise ValueError(f'Unexpected UMT5 asset size: {name}')
    else:
        directory = Path.home() / '.cache/flood2-live/t5'
        directory.mkdir(parents=True, exist_ok=True)
        with FileLock(str(directory/'.extract.lock')):
            if not all((directory/name).is_file() and (directory/name).stat().st_size == size
                       for name,(size,_) in FILES.items()):
                download_dir = directory/'download'
                archive = hf_hub_download(REPO, 'deps.zip', revision=REVISION,
                                          local_dir=download_dir)
                with zipfile.ZipFile(archive) as source:
                    for name, (size, crc) in FILES.items():
                        target = directory/name
                        if target.is_file() and target.stat().st_size == size:
                            continue
                        info = source.getinfo(PREFIX+name)
                        if info.file_size != size or info.CRC != crc:
                            raise ValueError(f'Unexpected UMT5 archive member: {name}')
                        target.parent.mkdir(parents=True, exist_ok=True)
                        temporary = target.with_suffix(target.suffix+'.part')
                        with source.open(info) as src, temporary.open('wb') as dst:
                            shutil.copyfileobj(src, dst, length=8*1024*1024)
                        if temporary.stat().st_size != size:
                            raise ValueError(f'Incomplete UMT5 extraction: {name}')
                        temporary.replace(target)
                # This dedicated temporary archive is no longer needed after
                # extraction; preserve room for ZeroGPU's tensor offload file.
                archive_path = Path(archive).resolve()
                if archive_path.parent == download_dir.resolve():
                    archive_path.unlink()
    return directory/'models_t5_umt5-xxl-enc-bf16.pth', directory/'google/umt5-xxl'