"""Resolve original or size-bounded Wan components without importing torch.

DiffSynth accepts one string or a list of shard paths per ModelConfig. Do not
pass an index JSON to its binary loader or treat each shard as a separate model.
WAN_MODEL_ROOT may point to an independently downloaded model bundle.
"""
import json
import os
from pathlib import Path
from project_paths import project_path


def wan_base_root(root=None):
    return str(Path(root or os.environ.get('WAN_MODEL_ROOT') or
                    project_path('models/Wan2.1-VACE-1.3B')).expanduser().resolve())


def indexed_shards(index_path):
    index_path = Path(index_path)
    index = json.loads(index_path.read_text())
    mapping = index.get('weight_map')
    if not isinstance(mapping, dict) or not mapping:
        raise ValueError('empty or invalid model weight index: ' + str(index_path))
    names = set()
    for name in mapping.values():
        if not isinstance(name, str) or Path(name).name != name or not name.endswith('.safetensors'):
            raise ValueError('unsafe shard path in model index')
        names.add(name)
    paths = [index_path.parent / name for name in sorted(names)]
    for path in paths:
        if not path.is_file():
            raise FileNotFoundError('missing model shard: ' + str(path))
    return [str(path) for path in paths]


def component_path(root, original, indices, strict=False):
    original_path = root / original
    if original_path.is_file():
        return str(original_path)
    for name in indices:
        index_path = root / name
        if index_path.is_file():
            return indexed_shards(index_path)
    if strict:
        raise FileNotFoundError('missing original or indexed model component: ' + str(original_path))
    # Preserve historical command defaults when no weights are present (e.g. dry-run).
    return str(original_path)


def wan_model_paths(root=None, strict=False):
    root = Path(wan_base_root(root))
    return [
        component_path(root, 'diffusion_pytorch_model.safetensors', (
            'transformer/diffusion_pytorch_model.safetensors.index.json',
            'diffusion_pytorch_model.safetensors.index.json',
        ), strict),
        component_path(root, 'models_t5_umt5-xxl-enc-bf16.pth', (
            'text_encoder/model.safetensors.index.json',
            'text_encoder/models_t5_umt5-xxl-enc-bf16.safetensors.index.json',
            'models_t5_umt5-xxl-enc-bf16.safetensors.index.json',
        ), strict),
        component_path(root, 'Wan2.1_VAE.pth', (), strict),
    ]


def wan_tokenizer_path(root=None, strict=False):
    path = Path(wan_base_root(root)) / 'google/umt5-xxl'
    if strict and not (path / 'tokenizer.json').is_file() and not (path / 'spiece.model').is_file():
        raise FileNotFoundError('missing Wan umT5 tokenizer: ' + str(path))
    return str(path)