import hashlib
import json
from pathlib import Path


BASE_MODEL = 'Qwen/Qwen-Image-2.1'
VARIANTS = tuple(f'{steps}step-{precision}' for steps in (2, 4, 8) for precision in ('bf16', 'fp8'))


def resolve_model(model, variant, source):
    if model:
        folder = Path(model).expanduser().resolve()
    elif source == 'huggingface':
        from huggingface_hub import snapshot_download

        folder = Path(snapshot_download(
            'Wuli-art/Qwen-Image-2.1-Turbo', allow_patterns=[f'{variant}/*'],
        )) / variant
    else:
        from modelscope import snapshot_download

        folder = Path(snapshot_download(
            'Wuli-Art/Qwen-Image-2.1-Turbo', allow_patterns=[f'{variant}/*'],
        )) / variant
    if not (folder / 'sampling_config.json').is_file():
        raise ValueError('Select a DiT directory containing sampling_config.json')
    return folder


def read_config(folder):
    sampling = json.loads((folder / 'sampling_config.json').read_text())
    transformer = json.loads((folder / 'config.json').read_text())
    sigmas = sampling['sample_sigmas']
    if len(sigmas) not in (2, 4, 8) or sigmas[0] != 1.0:
        raise ValueError('Unexpected saved sampling schedule')
    if any(first <= second for first, second in zip(sigmas, sigmas[1:])):
        raise ValueError('Sampling sigmas must strictly decrease')
    return sampling, transformer


def resolve_base(model=BASE_MODEL, source='huggingface'):
    path = Path(model).expanduser()
    if path.is_dir():
        return path.resolve()
    patterns = ['model_index.json', 'processor/*', 'scheduler/*', 'text_encoder/*', 'vae/*']
    if source == 'huggingface':
        from huggingface_hub import snapshot_download

        return Path(snapshot_download(model, allow_patterns=patterns))
    from modelscope import snapshot_download

    return Path(snapshot_download(model, allow_patterns=patterns))


def prepare_pipeline_directory(dit_folder, base_folder):
    sampling, transformer = read_config(dit_folder)
    key = hashlib.sha256((str(dit_folder.resolve()) + str(base_folder.resolve())).encode()).hexdigest()[:20]
    destination = Path.home() / '.cache' / 'wuli-qwen-image21' / key
    destination.mkdir(parents=True, exist_ok=True)
    for component in ['processor', 'text_encoder', 'vae']:
        target = base_folder / component
        if not target.is_dir():
            raise ValueError(f'Missing base component: {component}')
        link = destination / component
        if not link.exists():
            link.symlink_to(target.resolve(), target_is_directory=True)
    link = destination / 'transformer'
    if not link.exists():
        link.symlink_to(dit_folder.resolve(), target_is_directory=True)
    scheduler = destination / 'scheduler'
    scheduler.mkdir(exist_ok=True)
    scheduler_config = json.loads((dit_folder / 'scheduler_config.json').read_text())
    (scheduler / 'scheduler_config.json').write_text(json.dumps(scheduler_config, indent=2) + '\n')
    pipeline = json.loads((base_folder / 'model_index.json').read_text())
    pipeline['sample_sigmas'] = sampling['sample_sigmas']
    (destination / 'model_index.json').write_text(json.dumps(pipeline, indent=2) + '\n')
    return destination