"""Load an explicit released world-model variant; never substitute adapters."""
import hashlib
import json
from pathlib import Path, PurePosixPath


def safe_path(root, relative):
    rel = PurePosixPath(relative)
    if rel.is_absolute() or not rel.parts or '..' in rel.parts or '\\' in relative:
        raise ValueError(f'Invalid relative path: {relative!r}')
    root = Path(root).resolve()
    path = (root / relative).resolve()
    if path != root and root not in path.parents:
        raise ValueError('Path escapes bundle root')
    return path


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


def verify_bundle(root):
    root = Path(root)
    manifest = json.loads((root / 'manifest.json').read_text())
    if manifest.get('schema_version') != 1 or not manifest.get('files'):
        raise ValueError('Unsupported or empty manifest')
    seen = set()
    for item in manifest['files']:
        name = item['path']
        if name in seen:
            raise ValueError(f'Duplicate file: {name}')
        seen.add(name)
        path = safe_path(root, name)
        if not path.is_file() or path.stat().st_size != item['bytes']:
            raise ValueError(f'Missing file or wrong size: {name}')
        if file_sha256(path) != item['sha256']:
            raise ValueError(f'Checksum mismatch: {name}')
    return manifest


def load_world_model(root, dataset, variant='main', device='cpu', verify=True):
    """Load the predictor only. Codec/actions/T5 preprocessing are separate.

    FP32 and SDPA match the audited loaders. Dataset/variant selection is
    explicit; missing text adapters are an error, never silently zero-filled.
    """
    import torch
    from safetensors.torch import load_file
    from transformers import AutoConfig, AutoModelForCausalLM
    from .text_conditioning import TextWorldModel

    root = Path(root)
    if verify:
        verify_bundle(root)
    settings = json.loads(safe_path(root, f'{dataset}/futureworlds_config.json').read_text())
    if settings['dataset'] != dataset or variant not in settings['variants']:
        raise ValueError('Dataset or variant does not match bundle configuration')
    entry = settings['variants'][variant]
    folder = safe_path(root, entry['path'])
    cfg = AutoConfig.from_pretrained(folder, local_files_only=True, trust_remote_code=False)
    core = AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.float32, attn_implementation='sdpa', trust_remote_code=False)
    core.load_state_dict(load_file(str(folder / 'model.safetensors')), strict=True)
    if settings['text_conditioning']['enabled']:
        tc = settings['text_conditioning']
        model = TextWorldModel(core, text_width=tc['width'], layers=tc['layers'], inner=tc['inner'])
        adapters = folder / 'text_adapters.safetensors'
        if not adapters.is_file():
            raise FileNotFoundError(f'Required trained text adapters missing: {adapters}')
        model.adapters.load_state_dict(load_file(str(adapters)), strict=True)
    else:
        model = core
    return model.to(device).eval()