"""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
- the pump drains at `flow_max * capacity * speed` while DO0.0 is ON;
  the speed fraction is AO0.0 normalized over its display range
  (FR-10.2.14)
- the level integrates the balance, clamped at the brim; AI0.0
  reports the level scaled over the channel's display range
  (default 0-4095 counts = 100 %)
- a high-level switch DI0.1 closes when the level rises above 2.7 m
- optional is_done: stop after a given sim time (default: run forever)
"""

from __future__ import annotations

from wago_sim.scenario.plantenv import PlantEnv, PlantState

TANK_AREA = 5.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_SENSOR_MAX = 5.0   # m — tank level meter full-scale
LEVEL_BRIM = 2.9         # m — tank brim (physical clamp)
LEVEL_HIGH = 2.7         # m — high-level switch threshold


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

    TICK_MS = 9.0        # overrides PlantEnv TICK_MS

    def __init__(self, *args, stop_after_ms: float | None = None, **kw):
        super().__init__(*args, **kw)
        self._level_m = 1.2                      # 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 (master side: DO/AO)
        pump_running = bool(state.by_code("DO0.0"))

        # speed setpoint from AO0.0, normalized over its display range
        ao_lo, ao_hi = state.channel_range("AO0.0")
        ao_value = state.by_code("AO0.0")
        speed_raw = float(ao_value) if isinstance(ao_value, (int, float)) else 0.0
        try:
            speed = min(1.0, max(0.0, (speed_raw - ao_lo) / (ao_hi - ao_lo)))
        except ZeroDivisionError:
            speed = 0.0

        # physics: water balance with pump wear (parameters attachable)
        capacity = state.parameters.get("pump_capacity", 1.0)
        outflow = FLOW_MAX * capacity * speed if pump_running else 0.0
        d_level_m = (INFLOW - outflow) / TANK_AREA * dt_h
        self._level_m = min(LEVEL_BRIM, max(0.0, self._level_m + d_level_m))

        # the field the PLC reads back: level over the AI display range
        # (dense store, FR-10.2.14 — defaults to 0-4095 counts = 100 %)
        ai_lo, ai_hi = state.channel_range("AI0.0")
        level_raw = int((self._level_m / LEVEL_SENSOR_MAX) * (ai_hi - ai_lo)
                        + ai_lo)
        level_raw = min(ai_hi, max(ai_lo, level_raw))
        high_switch = self._level_m > LEVEL_HIGH       # DI0.1

        return {
            "AI0.0": level_raw,
            "DI0.1": high_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