"""CPU-only release validation, including relocatable frozen codec round trips."""
from pathlib import Path
import argparse
import gc
import hashlib
import json
import os
import sys

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))


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


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--codecs', action='store_true', help='Also load and test both pinned frozen codecs on CPU')
    args = parser.parse_args()
    # This command is explicitly CPU-only, including optional backend probes.
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    os.environ.setdefault('TORCHINDUCTOR_TRITON_DISABLE_DEVICE_DETECTION', '1')
    # Full-precision codecs do not use TorchAO. Some optional TorchAO builds
    # query CUDA during import even when devices are hidden; exclude it here.
    sys.modules['torchao'] = None
    import numpy as np
    import torch
    import quasar
    torch.set_num_threads(2)
    assert Path(quasar.__file__).resolve().is_relative_to(ROOT)
    release = json.loads((ROOT / 'release.json').read_text())
    manifest = json.loads((ROOT / 'files.sha256.json').read_text())
    for name, item in manifest.items():
        file = ROOT / name
        assert file.is_file() and file.stat().st_size == item['bytes'] and sha(file) == item['sha256'], name
    from quasar.training.checkpoint import load_quasar
    from quasar.tokenizers.text import TextTokenizer
    for folder, expected in release['weight_hashes'].items():
        path = ROOT / folder
        assert sha(path / 'model.safetensors') == expected
        model, vocab, meta = load_quasar(path, 'cpu')
        assert model.num_params() == 90247794
        assert all(bool(torch.isfinite(p).all()) for p in model.parameters())
        model.eval()
        tokenizer = TextTokenizer.load(path / 'tokenizers/text_tokenizer.json')
        tokens = [vocab.special('bos'), vocab.special('task:chat_text'), vocab.special('input'), vocab.special('text')]
        tokens += tokenizer.encode('Hello.') + [vocab.special('/text'), vocab.special('output'), vocab.special('text')]
        with torch.inference_mode():
            logits = model(torch.tensor([tokens]))
        assert torch.isfinite(logits).all() and logits.shape[-1] == 54144
        audio = meta['data']['preparation']['audio_codec']
        assert (path / audio['path'] / 'model.safetensors').resolve() == (ROOT / 'codecs/audio/model.safetensors').resolve()
        vq = json.loads((path / 'tokenizers/vq/vq_config.json').read_text())
        assert (path / 'tokenizers/vq' / vq['base_path']).resolve() == (ROOT / 'codecs/vision').resolve()
        print(folder, 'finite CPU forward / original weight hash / parameter count / codec paths passed', flush=True)
        del model, logits
        gc.collect()
    if args.codecs:
        from quasar.inference.extended import ExtendedGenerator
        generator = ExtendedGenerator(ROOT, 'cpu')
        pixels = np.zeros((1, 64, 64, 3), dtype=np.uint8)
        pixels[:, 20:40, 20:40, 1] = 255
        codes = generator.vision().encode(pixels)
        reconstructed = generator.vision().decode(codes)
        assert codes.shape == (1, 8, 8) and reconstructed.shape == pixels.shape
        assert np.isfinite(reconstructed).all()
        print('Pinned visual codec CPU encode/decode passed', flush=True)
        del generator.codec
        generator.codec = None
        gc.collect()
        codec = generator.audio()
        waveform = (.05 * np.sin(np.arange(2880) * 2 * np.pi * 440 / 24000)).astype(np.float32)
        packet = codec.encode(waveform, 24000)
        reconstructed = codec.decode(packet)
        assert reconstructed.shape == waveform.shape and np.isfinite(reconstructed).all()
        print('Pinned aligned audio codec CPU encode/decode passed', flush=True)
    print('Release verified on CPU; no optimizer or GPU job started.', flush=True)


if __name__ == '__main__':
    main()