"""Test-only environment: is_done() after 30 ms of sim time.

Used by the web-UI env-control tests (FR-10.2.15): the paced ticker
must notice is_done(), set state "done" (no auto-restart), and keep
the instance for inspection.
"""

from __future__ import annotations

from wago_sim.scenario.plantenv import PlantEnv, PlantState


class PumpStationEnv(PlantEnv):
    """Minimal env: fills AI0.0, done after 30 ms."""

    def __init__(self, *args, stop_after_ms: float | None = 30.0, **kw):
        super().__init__(*args, **kw)
        self.stop_after_ms = stop_after_ms

    def model(self, state: PlantState, dt_s: float) -> dict[str, object]:
        level = float(state.by_code("AI0.0"))
        return {"AI0.0": level + 1.0}

    def is_done(self, state: PlantState) -> bool:
        return state.sim_time_ms >= (self.stop_after_ms or 0.0)


Env = PumpStationEnv