import json

from conftest import GOOD_SCENE

from worker.sweep import SweepConfig, run_sweep_job, sample_worlds, summarize


def test_worlds_are_reproducible():
    a, b = sample_worlds(SweepConfig(n_worlds=5, seed=7)), sample_worlds(SweepConfig(n_worlds=5, seed=7))
    assert a == b and a != sample_worlds(SweepConfig(n_worlds=5, seed=8))
    assert all(0.5 <= w.friction <= 1.5 and 0 <= w.push_n <= 80 for w in a)


def test_sweep_job_on_cpu(menagerie):
    out = run_sweep_job(
        {"mjcf": GOOD_SCENE, "robot": "unitree_g1", "n_worlds": 3, "duration_s": 0.3, "push_n": [0, 20]},
        menagerie_dir=menagerie,
        engine="mujoco",
    )
    s = out["summary"]
    assert s["worlds"] == 3 and s["engine"] == "MuJoCo (CPU)" and "fall_rate" in s and "/3 fell" in s["headline"]
    table = json.loads(out["artifacts"]["sweep.json"])
    assert len(table["worlds"]) == 3 and {"push_n", "friction", "fell", "min_height"} <= set(table["worlds"][0])


def test_summary_without_fall_detection():
    from worker.sweep import World

    result = summarize([World(1.0, 1.0, 10.0, 0.0)], SweepConfig(n_worlds=1), "franka_emika_panda", "x", 1.0)
    assert "no fall detection" in result["summary"]["headline"] and "fall_rate" not in result["summary"]


def test_sweep_endpoint(run_env):
    client, keys_, backend, services, gen_id = run_env
    h = {"X-API-Key": keys_["alice"]}
    run_id = client.post("/simulate/run", json={"generation_id": gen_id}, headers=h).json()["run_id"]
    r = client.post("/simulate/sweep", json={"run_id": run_id, "n_worlds": 16, "push_max_n": 120}, headers=h)
    assert r.status_code == 200, r.text
    job = r.json()
    assert job["kind"] == "sweep" and backend.functions[-1] == "sweep"
    sent = backend.submitted[-1]
    assert sent["mjcf"] == GOOD_SCENE and sent["n_worlds"] == 16 and sent["push_n"] == [0.0, 120.0]
    row = services.logs.get_run(job["id"])
    assert row["parent_id"] == run_id and row["generation_id"] == gen_id

    worlds = [
        {"friction": 0.6, "object_mass": 1.0, "push_n": 110.0, "push_dir_deg": 0, "fell": True},
        {"friction": 1.4, "object_mass": 1.0, "push_n": 10.0, "push_dir_deg": 90, "fell": False},
    ]
    summary = {"headline": "1/2 fell (50%)", "worlds": 2, "fall_rate": 0.5, "ranges": {}, "engine": "MuJoCo Warp (GPU)"}
    backend.jobs[row["job_id"]] = {
        "status": "COMPLETED",
        "output": {
            "summary": summary,
            "worker_seconds": 20,
            "artifacts": {"sweep.json": json.dumps({"worlds": worlds}).encode()},
        },
    }
    page = client.get(job["report_url"].split("testserver")[-1]).text
    assert "1/2 fell (50%)" in page and "<svg" in page and "push 110.0 N" in page
    parent = client.get(f"/jobs/{run_id}", headers=h).json()["report_url"]
    assert "Robustness sweep" in client.get(parent.split("testserver")[-1]).text


def test_sweep_needs_exactly_one_source(run_env):
    client, keys_, backend, services, gen_id = run_env
    h = {"X-API-Key": keys_["alice"]}
    assert client.post("/simulate/sweep", json={}, headers=h).status_code == 422
    both = {"generation_id": gen_id, "mjcf": GOOD_SCENE}
    assert client.post("/simulate/sweep", json=both, headers=h).status_code == 422
    assert client.post("/simulate/sweep", json={"generation_id": gen_id, "n_worlds": 999}, headers=h).status_code == 422