"""Tests for offline synthetic generation (T1.5)."""
import json

import numpy as np
import pytest

from tsmm.data.gen_pipeline import gen_one_sample, write_jsonl


def test_gen_one_sample_has_required_schema():
    sample = gen_one_sample(T=128, C=2, seed=0)
    required = {
        "series",
        "attributes",
        "timestamps",
        "events",
        "instruction",
        "answer",
        "labels",
        "category",
    }
    assert required.issubset(sample.keys()), f"missing keys: {required - sample.keys()}"


def test_gen_one_sample_shapes_and_types():
    T, C = 128, 3
    s = gen_one_sample(T=T, C=C, seed=1)
    assert np.asarray(s["series"]).shape == (T, C)
    assert np.asarray(s["labels"]).shape == (T,)
    assert len(s["attributes"]) == C
    assert len(s["timestamps"]) == T
    # instruction + answer are non-empty strings
    assert isinstance(s["instruction"], str) and s["instruction"].strip()
    assert isinstance(s["answer"], str) and s["answer"].strip()
    # category is one of the 6
    from tsmm.data.instruct import INSTRUCTION_CATEGORIES
    assert s["category"] in INSTRUCTION_CATEGORIES


def test_gen_one_sample_reproducible():
    a = gen_one_sample(T=64, C=2, seed=42)
    b = gen_one_sample(T=64, C=2, seed=42)
    # JSON-serializable equality
    assert json.dumps(a, sort_keys=True, default=str) == json.dumps(b, sort_keys=True, default=str)


def test_gen_one_sample_anomaly_answer_parses_and_matches_labels():
    s = gen_one_sample(T=128, C=2, seed=3, force_category="anomaly")
    assert s["category"] == "anomaly"
    parsed = json.loads(s["answer"])
    recon = np.zeros(128, dtype=int)
    for r in parsed["anomaly_regions"]:
        recon[r["start"]: r["end"]] = 1
    assert np.array_equal(recon, np.asarray(s["labels"]).astype(int))


def test_gen_one_sample_json_serializable():
    s = gen_one_sample(T=64, C=2, seed=0)
    # must round-trip through json with no custom types
    blob = json.dumps(s)  # raises if non-serializable
    json.loads(blob)


def test_write_jsonl_creates_valid_file(tmp_path):
    samples = [gen_one_sample(T=32, C=2, seed=i) for i in range(5)]
    out = tmp_path / "out.jsonl"
    write_jsonl(samples, str(out))
    assert out.exists()
    lines = out.read_text().strip().splitlines()
    assert len(lines) == 5
    for line in lines:
        obj = json.loads(line)
        assert "series" in obj and "instruction" in obj