"""Jet amplitude update failure contract (host doubles only, no CUDA numerics).

Both owners validate before any device write, commit host state only after the
upload returns, and refuse every later call except close() after an uncertain
device write.
"""
from types import SimpleNamespace

import numpy as np
import pytest

pytest.importorskip("pycuda.driver")
from CelerisLab import simulation as module
from CelerisLab.body.jets import Jet, Radial, Sector
from tests.unit.test_pinned_body_remap import ready
from tests.unit.test_simulation_lifecycle import make_simulation

BAD = (float("nan"), float("inf"), 1.)


class Curved:
    """Records uploads; ``fail`` raises once the write was attempted."""
    def __init__(self):
        self.writes, self.fail = [], False

    def set_jet_amplitude(self, index, amplitude, stream):
        self.writes.append((index, amplitude))
        if self.fail:
            raise RuntimeError("upload failed")


def simulation(monkeypatch):
    sim = make_simulation(monkeypatch)
    jet = Jet(Sector((6., 5.), 0., .5, 4.), Radial((6., 5.)), "uniform", 1e-3)
    sim._failed, sim._initialized = None, True
    sim.bodies, sim.field = SimpleNamespace(jets=[jet]), SimpleNamespace(curved=Curved())
    return sim, jet


def test_simulation_invalid_input_writes_nothing_and_stays_usable(monkeypatch):
    sim, jet = simulation(monkeypatch)
    for amplitude in BAD:
        with pytest.raises(ValueError):
            sim.set_jet_amplitude(0, amplitude)
    with pytest.raises(IndexError):
        sim.set_jet_amplitude(1, 2e-3)
    assert sim.field.curved.writes == [] and sim.bodies.jets[0] is jet and sim._failed is None
    sim.set_jet_amplitude(0, 2e-3)  # retry after a rejected input
    assert sim.field.curved.writes == [(0, 2e-3)] and sim.bodies.jets[0].amplitude == 2e-3
    sim.close()


def test_simulation_upload_failure_keeps_host_and_blocks_reuse(monkeypatch):
    sim, jet = simulation(monkeypatch)
    sim.field.curved.fail = True
    with pytest.raises(RuntimeError, match="upload failed"):
        sim.set_jet_amplitude(0, 2e-3)
    assert sim.bodies.jets[0] is jet and sim._failed is not None
    sim.field.curved.fail = False
    for call in (lambda: sim.set_jet_amplitude(0, 2e-3), lambda: sim.run(1), sim.step,
                 sim.sync_bodies, sim.initialize, sim.get_ddf, sim.get_macroscopic,
                 lambda: sim.read_surface_exchange(0), lambda: sim.read_body(0),
                 sim.read_bodies, sim.snapshot, sim.restore, sim.save_checkpoint,
                 lambda: sim.load_checkpoint("x.h5"), lambda: sim.update_runtime_params(omega=1.)):
        with pytest.raises(RuntimeError, match="Simulation failed"):
            call()
    assert sim.field.curved.writes == [(0, 2e-3)]
    sim.close()  # cleanup still works
    assert sim._closed and sim.ctx.is_closed


def test_simulation_stream_creation_failure_fails_closed(monkeypatch):
    sim, jet = simulation(monkeypatch)
    created = []
    def stream():
        created.append(1)
        raise RuntimeError("stream creation failed")
    monkeypatch.setattr(module.cuda, "Stream", stream, raising=False)
    sim._stream = None
    with pytest.raises(RuntimeError, match="stream creation failed"):
        sim.set_jet_amplitude(0, 2e-3)
    assert sim.bodies.jets[0] is jet and sim._failed is not None
    assert sim.field.curved.writes == [] and created == [1]
    for call in (lambda: sim.stream, lambda: sim.set_jet_amplitude(0, 2e-3)):
        with pytest.raises(RuntimeError, match="Simulation failed"):
            call()
    assert created == [1]  # the public stream property does not retry creation
    sim.ctx.detach_failure = True
    with pytest.raises(RuntimeError, match="detach failed"):
        sim.close()
    assert not sim._closed
    sim.close()  # close stays retryable after the failure
    assert sim._closed and sim.ctx.is_closed


def test_simulation_before_initialize_commits_host_only(monkeypatch):
    sim, _ = simulation(monkeypatch)
    sim._initialized = False
    sim.set_jet_amplitude(0, -2e-3)
    assert sim.field.curved.writes == [] and sim.bodies.jets[0].amplitude == -2e-3
    sim.close()


def stepper_state(run):
    boundary = run._level_boundaries[2]
    return boundary.tables["surface_velocity"].copy(), boundary.jets, run._level_jets[2]


def test_stepper_invalid_input_writes_nothing_and_stays_ready():
    run, module, _, _, _ = ready()
    try:
        velocity, jets, level_jets = stepper_state(run)
        module.cuda.memcpy_htod.reset_mock()
        for args, kw, error in [((0, a), {"level": 2}, ValueError) for a in BAD] + [
                ((1, .01), {"level": 2}, IndexError), ((0, .01), {"level": 3}, ValueError)]:
            with pytest.raises(error):
                run.set_jet_amplitude(*args, **kw)
        module.cuda.memcpy_htod.assert_not_called()
        np.testing.assert_array_equal(run._level_boundaries[2].tables["surface_velocity"], velocity)
        assert (run._level_boundaries[2].jets, run._level_jets[2]) == (jets, level_jets)
        assert run._status == "ready"
        run.set_jet_amplitude(0, .017, level=2)
        assert run._level_jets[2][0].amplitude == .017
    finally:
        run.close()


@pytest.mark.parametrize("where", ["upload", "synchronize"])
def test_stepper_device_failure_keeps_host_and_blocks_reuse(where):
    run, module, _, _, _ = ready()
    try:
        velocity, jets, level_jets = stepper_state(run)
        if where == "upload":
            module.cuda.memcpy_htod.side_effect = RuntimeError("upload failed")
        else:
            run._stream.synchronize.side_effect = RuntimeError("upload failed")
        with pytest.raises(RuntimeError, match="upload failed"):
            run.set_jet_amplitude(0, .017, level=2)
        np.testing.assert_array_equal(run._level_boundaries[2].tables["surface_velocity"], velocity)
        assert (run._level_boundaries[2].jets, run._level_jets[2]) == (jets, level_jets)
        assert run._status == "failed"
        module.cuda.memcpy_htod.side_effect = None
        run._stream.synchronize.side_effect = None
        for call in (lambda: run.set_jet_amplitude(0, .017, level=2), run.step, run.download):
            with pytest.raises(RuntimeError, match="failed"):
                call()
    finally:
        run.close()
    assert run._status == "closed"