"""Run with no source checkpoint and disable PyTorch neural weight loading."""

import asyncio
import json
from pathlib import Path

import numpy as np
from PIL import Image
from transformers import PreTrainedModel

from coreai_embeddinggemma2 import EmbeddingGemma2CoreAI


def forbidden(*args, **kwargs):
    raise AssertionError(
        "PyTorch neural weight loading is forbidden in this smoke test"
    )


PreTrainedModel.from_pretrained = classmethod(forbidden)


async def main():
    root = Path("bundles/embeddinggemma-2-CoreAI")
    runtime = await EmbeddingGemma2CoreAI.load(root)
    audio = (0.2 * np.sin(2 * np.pi * 440 * np.arange(8000) / 16000)).astype(np.float32)
    image = Image.new("RGB", (64, 96), (30, 70, 200))
    video = np.stack([np.asarray(image), np.asarray(image)])
    inputs = [
        {"text": "What causes auroras?", "prompt_name": "SearchQuery"},
        {"images": image},
        {"audio": audio},
        {"videos": video},
        {"text": "<|image|> Blue. <|audio|> Tone.", "images": image, "audio": audio},
    ]
    results = []
    for data in inputs:
        v = await runtime.encode(**data)
        assert (
            v.shape == (768,)
            and np.isfinite(v).all()
            and abs(np.linalg.norm(v) - 1) < 1e-5
        )
        results.append(
            {
                "modalities": list(data),
                "shape": list(v.shape),
                "unit_norm": float(np.linalg.norm(v)),
            }
        )
    for data in [
        {"text": "word " * 3000},
        {"audio": np.zeros(179201, np.float32)},
        {"audio": audio, "sampling_rate": 8000},
        {"text": "hello", "truncate_dim": 769},
    ]:
        try:
            await runtime.encode(**data)
        except ValueError:
            pass
        else:
            raise AssertionError("Invalid input accepted")
    (root / "smoke.json").write_text(
        json.dumps(
            {
                "pytorch_neural_loading_disabled": True,
                "cases": results,
                "invalid_inputs_rejected": 4,
            },
            indent=2,
        )
        + "\n"
    )
    print("Core AI runtime-only smoke passed", flush=True)


if __name__ == "__main__":
    asyncio.run(main())