"""Example pump-station plant environment (spec FR-8.2.17).

Demonstrates environment authoring: subclass PlantEnv, override
``model()`` — read what the PLC commanded (DO/AO), compute what the
field now shows (DI/AI). Run it live against a Modbus master:

    wago-sim serve --node config/pump_node.yaml --port 5020 \
                   --env envs.pump_station_env:PumpStationEnv

or headless for model development:

    wago-sim sim --node config/pump_node.yaml \
                 --env envs.pump_station_env:PumpStationEnv

Physics (deliberately simple, for demonstration): one tank.
- inflow is constant (or a small sinusoidal demand profile)
- the pump drains at `flow_max * capacity * speed` while DO0.0 is ON
- the level integrates the balance; AI0.0 reports it as 0.1 mm/count
- a low-level switch DI0.1 opens when the tank drains below 500 mm
- optional is_done: stop after a sim hour (default: run forever)
"""

from __future__ import annotations

from wago_sim.scenario.plantenv import PlantEnv, PlantState

TANK_AREA = 2.0          # m^2
FLOW_MAX = 30.0          # m^3/h the pump moves at full speed
INFLOW = 12.0            # m^3/h constant inflow
LEVEL_MAX = 9000.0       # mm — tank brim
SWITCH_AT = 500.0        # mm — low-level switch threshold


class PumpStationEnv(PlantEnv):
    """Closed-loop tank + pump physics over the pump_node layout."""

    def __init__(self, *args, stop_after_ms: float | None = None, **kw):
        super().__init__(*args, **kw)
        self._level_mm = 3000.0          # initial tank level
        self.stop_after_ms = stop_after_ms

    def model(self, state: PlantState, dt_s: float) -> dict[str, object]:
        dt_h = dt_s / 3600.0
        # what the PLC currently commands
        pump_on = bool(state.by_code("DO0.0"))
        # optional speed setpoint from AO0.0 (counts -> 0..1 fraction)
        speed_raw = float(state.by_code("AO0.0") or 0)
        speed = min(1.0, max(0.0, speed_raw / 4095.0))
        # physics: water balance with pump wear (parameters attachable)
        capacity = state.parameters.get("pump_capacity", 1.0)
        outflow = FLOW_MAX * capacity * speed if pump_on else 0.0
        d_level_mm = (INFLOW - outflow) / TANK_AREA * 1000.0 * dt_h
        self._level_mm = min(LEVEL_MAX, max(0.0, self._level_mm + d_level_mm))
        # the field the PLC reads back
        level_counts = int(self._level_mm / 0.1)      # AI0.0: 0.1 mm/count
        low_switch = self._level_mm < SWITCH_AT       # DI0.1
        return {
            "AI0.0": level_counts,
            "DI0.1": low_switch,
        }

    def is_done(self, state: PlantState) -> bool:
        if self.stop_after_ms is None:
            return False                 # default: run until stopped
        return state.sim_time_ms >= self.stop_after_ms


def make_env(node=None, engine=None) -> PumpStationEnv:
    """Factory form (FR-8.2.17): envs.pump_station_env:make_env."""
    return PumpStationEnv(node=node, engine=engine)


# bare-module form: `--env envs.pump_station_env` picks this up
Env = PumpStationEnv