"""Synthetic protocol checks; these do not qualify rendering or AI model quality."""

import base64
import hashlib
import io
import json

import numpy as np
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from PIL import Image
from pydantic import ValidationError
from test_observatory_recorded_jobs import _definitions

from k1link.observatory.recorded_jobs import ObservatoryRecordedJobQueue
from k1link.simulation.ai_polygon.contracts import (
    Decision,
    RunApplied,
    RunCreate,
    RunSample,
    WorkerHello,
    WorkerWorldCreate,
    WorldCreate,
    WorldSettings,
)
from k1link.simulation.ai_polygon.policy import RoadPolicy
from k1link.simulation.ai_polygon.runs import RunStore
from k1link.simulation.ai_polygon.worlds import WorldStore, inspect_gaussian_ply
from k1link.web.ai_polygon_api import build_ai_polygon_router


def ply(value=0.0):
    names = [
        "x",
        "y",
        "z",
        "opacity",
        "f_dc_0",
        "f_dc_1",
        "f_dc_2",
        "scale_0",
        "scale_1",
        "scale_2",
        "rot_0",
        "rot_1",
        "rot_2",
        "rot_3",
    ]
    header = "ply\nformat binary_little_endian 1.0\nelement vertex 1\n"
    header += "".join(f"property float {name}\n" for name in names) + "end_header\n"
    return header.encode() + np.full(14, value, dtype="<f4").tobytes()


def make_world(worlds):
    data = ply()
    world = worlds.create(
        WorldCreate(
            name="Synthetic protocol fixture",
            filename="fixture.PLY",
            byte_length=len(data),
            author="test",
            license="CC0",
        )
    )
    worlds.append(world["world_id"], 0, data)
    worlds.complete(world["world_id"])
    return worlds.configure(world["world_id"], WorldSettings(prepared=True))


@pytest.fixture
def runs(tmp_path):
    queue = ObservatoryRecordedJobQueue(tmp_path, definitions=_definitions())
    worlds = WorldStore(tmp_path)
    store = RunStore(worlds, queue)
    hello = WorkerHello(
        worker_id="worker-006",
        instance_id="a" * 32,
        runtime="isaac-sim-6.1",
        model_ids=["ddrnet-goose-pytorch-reference", "rf_detr_large"],
        profile_sha256="b" * 64,
        runtime_sources={key: "c" * 64 for key in ("worker", "scene", "models", "robot")},
    )
    store.register(hello)
    return store, hello, make_world(worlds)


def sample(sequence=0, **changes):
    stream = io.BytesIO()
    Image.new("RGB", (800, 600)).save(stream, "JPEG")
    row = dict(
        sequence=sequence,
        simulation_time_ns=sequence * 100_000_000,
        inference_ms=5,
        pose_xy=(0, 0),
        decision=Decision(
            speed_mps=0, yaw_rate_rps=0, reason="no-road", road_fraction=0, obstacle_count=0
        ),
        image_jpeg_base64=base64.b64encode(stream.getvalue()).decode(),
    )
    return RunSample(**{**row, **changes})


def test_resume_and_atomic_publication_recovery(tmp_path):
    worlds = WorldStore(tmp_path)
    data = ply()
    row = worlds.create(
        WorldCreate(
            name=" Test ", filename="scan.ply", byte_length=len(data), author="Test", license="CC0"
        )
    )
    key = row["world_id"]
    worlds.append(key, 0, data[:20])
    assert worlds.prefix_hashes(key)["chunks"] == [
        {"byte_length": 20, "sha256": hashlib.sha256(data[:20]).hexdigest()}
    ]
    with pytest.raises(RuntimeError):
        worlds.append(key, 0, data[:20])
    with pytest.raises(RuntimeError):
        worlds.complete(key)
    worlds.append(key, 20, data[20:])
    directory = worlds.directory(key)
    (directory / "source.part").replace(directory / "source.ply")
    recovered = WorldStore(tmp_path).complete(key)
    assert recovered["status"] == "available"
    assert recovered["sha256"] == hashlib.sha256(data).hexdigest()
    assert recovered["name"] == "Test"
    assert worlds.complete(key) == recovered


@pytest.mark.parametrize(
    "data", [b"not ply", ply()[:-1], ply() + b"xxxx", ply(float("nan")), ply(float("inf"))]
)
def test_invalid_ply_rejected(tmp_path, data):
    path = tmp_path / "bad.ply"
    path.write_bytes(data)
    with pytest.raises(ValueError):
        inspect_gaussian_ply(path)


@pytest.mark.parametrize(
    "settings",
    [{"spawn_xy": [float("nan"), 0]}, {"rotation_degrees": [0, 0, 900]}, {"max_speed_mps": 2}],
)
def test_nonfinite_or_out_of_bounds_settings_rejected(settings):
    with pytest.raises(ValidationError):
        WorldSettings(**settings)


def test_start_snapshots_and_idempotency(runs):
    store, hello, world = runs
    request = RunCreate(world_id=world["world_id"], max_steps=2)
    row = store.start(request, "request-001")
    assert store.start(request, "request-001") == row
    store.worlds.configure(world["world_id"], WorldSettings(max_speed_mps=0.5, prepared=True))
    assert store.get(row["run_id"])["world"]["settings"]["max_speed_mps"] == 0.3
    with pytest.raises(RuntimeError):
        store.start(RunCreate(world_id=world["world_id"], max_steps=3), "request-001")
    with pytest.raises(RuntimeError):
        store.start(request, "request-002")
    assert store.poll(hello.instance_id, None)["action"] == "load"


def test_sequence_clock_pause_step_stop_and_resource_release(runs):
    store, hello, world = runs
    key = store.start(RunCreate(world_id=world["world_id"], max_steps=2), "request-001")["run_id"]
    with pytest.raises(RuntimeError):
        store.control(key, "step")
    store.poll(hello.instance_id, key)
    store.sample(key, hello.instance_id, sample())
    with pytest.raises(RuntimeError):
        store.sample(key, hello.instance_id, sample())
    store.applied(
        key,
        hello.instance_id,
        RunApplied(sequence=0, simulation_time_ns=100_000_000, physics_steps=6, pose_xy=(0, 0)),
    )
    with pytest.raises(ValueError):
        store.sample(key, hello.instance_id, sample(1, simulation_time_ns=999))
    store.control(key, "pause")
    assert store.get(key)["state"] == "running"
    assert store.poll(hello.instance_id, key)["run"]["state"] == "paused"
    store.control(key, "step")
    assert store.poll(hello.instance_id, key)["action"] == "step"
    assert store.poll(hello.instance_id, key)["action"] == "pause"
    store.control(key, "stop")
    with pytest.raises(RuntimeError):
        store.sample(key, hello.instance_id, sample(1))
    store.finish(key, hello.instance_id, "stopped", "")
    store.queue.reserve_simulation("airun-" + "b" * 32)


def test_expired_worker_stops_run_but_does_not_claim_gpu_released(runs):
    store, hello, world = runs
    key = store.start(RunCreate(world_id=world["world_id"]), "request-001")["run_id"]
    store.seen -= 21
    assert not store.status()["available"]
    assert store.get(key)["state"] == "failed"
    with pytest.raises(RuntimeError):
        store.sample(key, hello.instance_id, sample())
    with pytest.raises(RuntimeError):
        store.queue.reserve_simulation("airun-" + "b" * 32)
    store.register(hello)
    assert store.finish(key, hello.instance_id, "failed", "")["state"] == "failed"
    store.queue.reserve_simulation("airun-" + "b" * 32)


def test_restart_fences_worker_and_preserves_journal(runs):
    store, hello, world = runs
    key = store.start(RunCreate(world_id=world["world_id"], max_steps=2), "request-001")["run_id"]
    store.poll(hello.instance_id, key)
    store.sample(key, hello.instance_id, sample())
    restarted = RunStore(store.worlds, store.queue)
    assert restarted.get(key)["state"] == "failed"
    with pytest.raises(RuntimeError):
        restarted.sample(key, hello.instance_id, sample(1))
    saved = json.loads((store.directory(key) / "decisions.jsonl").read_text())
    assert (
        saved["image_sha256"]
        == hashlib.sha256((store.directory(key) / "frames/000000.jpg").read_bytes()).hexdigest()
    )
    assert "image_jpeg_base64" not in saved


def test_completion_requires_expected_steps(runs):
    store, hello, world = runs
    key = store.start(RunCreate(world_id=world["world_id"], max_steps=1), "request-001")["run_id"]
    with pytest.raises(RuntimeError):
        store.finish(key, hello.instance_id, "completed", "")
    store.poll(hello.instance_id, key)
    store.sample(key, hello.instance_id, sample())
    with pytest.raises(RuntimeError):
        store.finish(key, hello.instance_id, "completed", "")
    store.applied(
        key,
        hello.instance_id,
        RunApplied(sequence=0, simulation_time_ns=100_000_000, physics_steps=6, pose_xy=(0, 0)),
    )
    assert store.finish(key, hello.instance_id, "completed", "")["state"] == "completed"


def test_api_auth_and_bounded_upload(tmp_path):
    app = FastAPI()
    app.include_router(build_ai_polygon_router(tmp_path))
    with TestClient(app) as client:
        root = "/api/v1/ai-polygon"
        assert client.get(root + "/catalog").json()["runtime"]["available"] is False
        assert client.post(root + "/worker/poll", json={"instance_id": "a" * 32}).status_code == 401
        assert client.post(root + "/worlds", json={"filename": "../bad.ply"}).status_code == 422
        world = make_world(WorldStore(tmp_path))
        response = client.get(root + f"/worlds/{world['world_id']}/source.ply")
        assert response.status_code == 200 and response.content == ply()
        assert (
            client.post(
                root + "/runs",
                headers={"Idempotency-Key": "test-run-001"},
                json={"world_id": world["world_id"]},
            ).status_code
            == 409
        )
        assert not any((tmp_path / "ai-polygon/runs").iterdir())


def test_camera_policy_brakes_and_waits_before_resume():
    policy = RoadPolicy(0.3)
    road = np.ones((512, 512), dtype=bool)
    assert policy.decide(road, []).speed_mps == 0
    policy.decide(road, [])
    assert policy.decide(road, []).speed_mps == 0.3
    obstacle = policy.decide(road, [(0.4, 0.2, 0.6, 0.8)])
    assert obstacle.reason == "obstacle" and obstacle.speed_mps == 0
    assert policy.decide(road, []).speed_mps == 0
    assert policy.decide(np.zeros_like(road), []).reason == "no-road"
    with pytest.raises(ValueError):
        policy.decide(road, [(float("nan"), 0, 1, 1)])


def worker_asset():
    return WorkerWorldCreate(
        name="Synthetic paired scene",
        filename="fixture.ply",
        byte_length=len(ply()),
        author="test",
        license="CC0",
        sha256="a" * 64,
        collider_sha256="b" * 64,
        splat_count=1,
        settings=WorldSettings(prepared=True),
    )


def test_worker_asset_is_manifest_only_and_retry_preserves_settings(tmp_path):
    worlds = WorldStore(tmp_path)
    request = worker_asset()
    row = worlds.register_worker_asset(request, "worker-006")
    assert {p.name for p in worlds.directory(row["world_id"]).iterdir()} == {"world.json"}
    changed = worlds.configure(row["world_id"], WorldSettings(spawn_xy=(3, 4), prepared=True))
    assert worlds.register_worker_asset(request, "worker-006") == changed
    with pytest.raises(RuntimeError):
        worlds.register_worker_asset(
            request.model_copy(update={"collider_sha256": "c" * 64}), "worker-006"
        )
    with pytest.raises(RuntimeError):
        worlds.prefix_hashes(row["world_id"])


def test_worker_asset_api_requires_current_authenticated_instance(tmp_path):
    app = FastAPI()
    app.include_router(build_ai_polygon_router(tmp_path))
    base = "/api/v1/ai-polygon"
    store = RunStore(WorldStore(tmp_path))
    hello = WorkerHello(
        worker_id="worker-006",
        instance_id="a" * 32,
        runtime="isaac-sim-6.1",
        model_ids=["test"],
        runtime_sources={k: "c" * 64 for k in ("worker", "scene", "models", "robot")},
        profile_sha256="b" * 64,
    )
    auth = {"Authorization": "Bearer " + store.token, "Worker-Instance": hello.instance_id}
    body = worker_asset().model_dump(mode="json")
    with TestClient(app) as client:
        assert (
            client.post(
                base + "/worker/worlds", json=body, headers={"Worker-Instance": hello.instance_id}
            ).status_code
            == 401
        )
        assert client.post(base + "/worker/worlds", json=body, headers=auth).status_code == 409
        assert (
            client.post(
                base + "/worker/register", json=hello.model_dump(), headers=auth
            ).status_code
            == 200
        )
        response = client.post(base + "/worker/worlds", json=body, headers=auth)
        assert response.status_code == 201
        row = response.json()
        assert row["storage"] == {"kind": "worker", "worker_id": "worker-006"}
        assert client.get(base + f"/worlds/{row['world_id']}/source.ply").status_code == 409
        assert (
            client.post(
                base + "/worker/worlds", json=body, headers={**auth, "Worker-Instance": "d" * 32}
            ).status_code
            == 409
        )


def test_camera_policy_turns_towards_connected_road():
    policy = RoadPolicy(0.3)
    road = np.zeros((512, 512), dtype=bool)
    road[:, 100:260] = True
    policy.decide(road, [])
    policy.decide(road, [])
    decision = policy.decide(road, [])
    assert decision.yaw_rate_rps > 0 and 0 < decision.speed_mps < 0.3


def test_progress_does_not_claim_camera_ready(runs):
    store, hello, world = runs
    run = store.start(RunCreate(world_id=world["world_id"]), "progress-case-001")
    store.progress(run["run_id"], hello.instance_id, "scene")
    assert store.get(run["run_id"])["phase"] == "scene"
    assert store.get(run["run_id"])["state"] == "starting"
    store.control(run["run_id"], "stop")
    assert store.progress(run["run_id"], hello.instance_id, "models")["control"] == "stop"
    assert store.get(run["run_id"])["state"] == "stopping"


def test_worker_stop_race_accepts_already_exited_process(monkeypatch):
    import importlib.util
    from pathlib import Path

    runtime = Path(__file__).resolve().parents[1] / "simulation/ai-polygon"
    monkeypatch.syspath_prepend(str(runtime))
    spec = importlib.util.spec_from_file_location("polygon_worker_stop_test", runtime / "worker.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    calls = []

    class Child:
        pid = 123

        def poll(self):
            return None

        def wait(self, timeout):
            calls.append(("wait", timeout))
            return 0

    def taskkill(*args, **kwargs):
        # Windows reports a non-zero taskkill when the process exits in the race.
        assert kwargs["check"] is False
        calls.append(("taskkill", 255))

    monkeypatch.setattr(module.subprocess, "run", taskkill)
    module.terminate_episode(Child())
    assert calls == [("taskkill", 255), ("wait", 30)]