# tests/test_wave0_assets.py
"""Wave 0 asset verification tests. Run after wave0_download_assets.py."""
from pathlib import Path
import ast
import json
import pytest

ROOT = Path(__file__).resolve().parent.parent

MATH_PATH = ROOT / "results" / "benchmark_datasets" / "math_competition" / "math_level3_5.jsonl"


def test_math_exists():
    assert MATH_PATH.exists(), f"Missing: {MATH_PATH}"


def test_math_nonempty():
    lines = MATH_PATH.read_text().splitlines()
    assert len(lines) >= 100, f"Expected >= 100 records, got {len(lines)}"


def test_math_schema():
    rec = json.loads(MATH_PATH.read_text().splitlines()[0])
    assert "problem" in rec, "missing 'problem'"
    assert "solution" in rec, "missing 'solution'"
    assert "level" in rec, "missing 'level'"
    assert rec["level"] in {"Level 3", "Level 4", "Level 5"}, f"Bad level: {rec['level']}"


def test_math_type_field():
    rec = json.loads(MATH_PATH.read_text().splitlines()[0])
    assert "type" in rec  # algebra, geometry, etc.


MMMU_PATH = ROOT / "results" / "benchmark_datasets" / "mmmu" / "mmmu_math_physics.jsonl"

def test_mmmu_exists():
    assert MMMU_PATH.exists(), f"Missing: {MMMU_PATH}"

def test_mmmu_nonempty():
    lines = MMMU_PATH.read_text().splitlines()
    assert len(lines) >= 50, f"Expected >= 50, got {len(lines)}"

def test_mmmu_schema():
    rec = json.loads(MMMU_PATH.read_text().splitlines()[0])
    assert "question" in rec
    assert "answer" in rec
    assert "subject" in rec
    assert rec["subject"] in {"Math", "Physics"}

def test_mmmu_options_parsed():
    """Options must be a real list (not a string) after download."""
    rec = json.loads(MMMU_PATH.read_text().splitlines()[0])
    assert isinstance(rec["options"], list), \
        f"options should be list, got {type(rec['options'])}: {rec['options']}"

def test_mmmu_has_image_field():
    rec = json.loads(MMMU_PATH.read_text().splitlines()[0])
    assert "has_image" in rec
    assert isinstance(rec["has_image"], bool)

XQUAD_DIR = ROOT / "results" / "benchmark_datasets" / "xquad"

@pytest.mark.parametrize("lang", ["ar", "zh", "es"])
def test_xquad_exists(lang):
    assert (XQUAD_DIR / f"xquad_{lang}.jsonl").exists()

@pytest.mark.parametrize("lang", ["ar", "zh", "es"])
def test_xquad_nonempty(lang):
    lines = (XQUAD_DIR / f"xquad_{lang}.jsonl").read_text().splitlines()
    assert len(lines) >= 100, f"xquad_{lang}: only {len(lines)} records"

@pytest.mark.parametrize("lang", ["ar", "zh", "es"])
def test_xquad_schema(lang):
    rec = json.loads((XQUAD_DIR / f"xquad_{lang}.jsonl").read_text().splitlines()[0])
    assert "question" in rec
    assert "context" in rec
    assert "answers" in rec
    assert isinstance(rec["answers"], list)  # list of answer strings
    assert rec["lang"] == lang

GAIA_PATH = ROOT / "results" / "benchmark_datasets" / "gaia" / "gaia_level2_3.jsonl"

def test_gaia_exists():
    assert GAIA_PATH.exists(), f"Missing: {GAIA_PATH}"

def test_gaia_nonempty():
    lines = GAIA_PATH.read_text().splitlines()
    assert len(lines) >= 20, f"Expected >= 20, got {len(lines)}"

def test_gaia_schema():
    rec = json.loads(GAIA_PATH.read_text().splitlines()[0])
    assert "task" in rec
    assert "level" in rec
    assert rec["level"] in {2, 3}  # stored as int in our JSONL

def test_gaia_levels_only_2_and_3():
    recs = [json.loads(l) for l in GAIA_PATH.read_text().splitlines()]
    levels = {r["level"] for r in recs}
    assert levels.issubset({2, 3}), f"Unexpected levels: {levels}"

MISTRAL_DIR = ROOT / "models" / "mistral-7b"

def test_mistral_dir_exists():
    assert MISTRAL_DIR.exists()

def test_mistral_config_valid():
    cfg = json.loads((MISTRAL_DIR / "config.json").read_text())
    assert cfg.get("model_type") == "mistral"
    assert cfg.get("num_hidden_layers", 0) > 0

def test_mistral_has_weights():
    shards = list(MISTRAL_DIR.glob("*.safetensors"))
    assert len(shards) > 0, f"No safetensors shards in {MISTRAL_DIR}"

def test_mistral_tokenizer_present():
    # tokenizer.model is required for SentencePiece tokenizer
    assert (MISTRAL_DIR / "tokenizer.model").exists(), \
        "tokenizer.model missing — do not exclude it from snapshot_download"

MINILM_DIR = ROOT / "models" / "multilingual-minilm"

def test_minilm_dir_exists():
    assert MINILM_DIR.exists()

def test_minilm_config_valid():
    cfg = json.loads((MINILM_DIR / "config.json").read_text())
    assert cfg.get("hidden_size") == 384

@pytest.mark.slow
def test_minilm_encodes_multilingual():
    """Functional test: model loads and produces 384-dim vectors for ar/zh/es text."""
    try:
        from sentence_transformers import SentenceTransformer
    except ImportError:
        pytest.skip("sentence-transformers not installed")
    model = SentenceTransformer(str(MINILM_DIR))
    texts = ["مرحبا", "你好", "Hola"]  # Arabic, Chinese, Spanish
    vecs = model.encode(texts)
    assert vecs.shape == (3, 384), f"Expected (3, 384), got {vecs.shape}"
