rlmesh 0.1.0

use std::sync::Arc;

use rlmesh_grpc::wire::Bytes;

use crate::spaces;
use crate::{Error, Result};

/// One row's episode identity: the runtime-minted id (UUIDv7, never repeats)
/// and the explicit seed the episode was reset with, if any. Mirrors the wire's
/// `EpisodeInfo` row, so identity and seed cannot fall out of alignment.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct EpisodeInfo {
    /// Minted by the runtime (UUIDv7); never repeats.
    pub episode_id: String,
    /// The explicit reset seed; `None` when the episode was not explicitly
    /// seeded (a genuine seed of 0 is `Some(0)`).
    pub seed: Option<i64>,
}

/// The per-predict sampling seed: a reproducible mix of the episode's reset seed
/// and the re-plan ordinal within that episode.
///
/// A stochastic policy seeded once per episode cannot be reproduced under
/// interleaving — another episode's forward advances the same global generators
/// between two of ours. Deriving a fresh seed per predict from `(episode_seed,
/// predict_index)` makes each forward reproducible on its own, independent of
/// what ran between. The SDK stamps the result on the predict context as
/// `predict_seed`; the same function is exported to Python as
/// `rlmesh.predict_seed`, so a model can derive it itself.
///
/// FNV-1a over the seed and the ordinal, masked to 32 bits — the narrow range
/// every framework's seeding API accepts.
pub fn predict_seed(episode_seed: i64, predict_index: u64) -> i64 {
    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;

    fn update(mut hash: u64, bytes: &[u8]) -> u64 {
        for byte in bytes {
            hash ^= u64::from(*byte);
            hash = hash.wrapping_mul(FNV_PRIME);
        }
        hash
    }

    let mut hash = FNV_OFFSET;
    hash = update(hash, &episode_seed.to_le_bytes());
    hash = update(hash, &[0xfe]);
    hash = update(hash, &predict_index.to_le_bytes());
    (hash & u64::from(u32::MAX)) as i64
}

/// Routing metadata attached to a [`ModelObservation`].
///
/// Identifies the env (adapter) the request belongs to and the ordered per-row
/// episode identity — the self-describing batch. Row `i` of the batched
/// observation belongs to `episodes[i]`. There is no positional lane/slot
/// concept; the model keys all per-episode state by `episode_id`, never by
/// position.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ModelRouteContext {
    /// Identifier of the session this request belongs to (correlation only).
    pub session_id: String,
    /// The connected env container (UUIDv7) this adapter serves. The routing key.
    pub env_id: String,
    /// Identifier of this individual request.
    pub request_id: String,
    /// Ordered per-row episode identity (length `== num_envs`).
    pub episodes: Vec<EpisodeInfo>,
}

impl ModelRouteContext {
    /// The first row's episode id, or `""` if the batch is empty. The natural
    /// choice for single-env handlers.
    pub fn primary_episode_id(&self) -> &str {
        self.episodes
            .first()
            .map(|episode| episode.episode_id.as_str())
            .unwrap_or("")
    }

    /// The episode id of every row, in batch order.
    pub fn episode_ids(&self) -> Vec<String> {
        self.episodes
            .iter()
            .map(|episode| episode.episode_id.clone())
            .collect()
    }
}

/// An observation delivered to [`ModelHandler::predict`](crate::ModelHandler::predict).
///
/// Carries the encoded observation payload plus the routing context and the
/// env contract / batch size needed to decode it.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ModelObservation {
    /// Raw per-leaf observation wire bytes, if present. `None` is the only
    /// "absent value" signal (§5) — decode via
    /// [`decoded`](ModelObservation::decoded)/[`decoded_lanes`](ModelObservation::decoded_lanes),
    /// never read these bytes directly.
    pub observation: Option<Vec<Bytes>>,
    /// Routing/episode metadata for this request.
    pub route: ModelRouteContext,
    /// Number of rows (sub-environments) in this batch (`== route.episodes.len()`).
    pub num_envs: usize,
    /// The env contract (spaces/metadata) for decoding the observation. Shared
    /// (`Arc`) so the per-predict hot path clones a refcount, not the contract.
    pub env_contract: Option<Arc<spaces::EnvContract>>,
    /// The env steps executed since the previous predict without one (replayed
    /// chunk frames), oldest first. Non-empty only on a route that negotiated
    /// history; the engine advances its frame windows through them before
    /// this request's own observation.
    pub history: Vec<HistoryFrame>,
    /// The producer's per-group env-step counter at `observation`; required
    /// once history is negotiated (a gap or repeat is a rejected request).
    pub step: Option<i64>,
}

/// One replayed env step delivered as history: the batched observation and
/// the row-aligned episode identity AT THAT STEP, plus its step counter.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct HistoryFrame {
    /// Raw per-leaf observation wire bytes (see [`ModelObservation::observation`]).
    pub observation: Option<Vec<Bytes>>,
    /// Row `i` of `observation` belongs to `episodes[i]`.
    pub episodes: Vec<EpisodeInfo>,
    /// The step counter this frame was observed at.
    pub step: i64,
}

impl ModelObservation {
    /// The first row's episode id (see [`ModelRouteContext::primary_episode_id`]).
    pub fn episode_id(&self) -> &str {
        self.route.primary_episode_id()
    }

    /// The episode id of every row (see [`ModelRouteContext::episode_ids`]).
    pub fn episode_ids(&self) -> Vec<String> {
        self.route.episode_ids()
    }

    /// Validate the request carries everything a decode needs (observation, env
    /// contract, observation space) WITHOUT paying the batch decode. Lets a
    /// malformed request fail fast even on a pure chunk-replay step that skips the
    /// actual decode below, instead of silently replaying buffered actions.
    pub fn ensure_decodable(&self) -> Result<()> {
        if self.observation.is_none() {
            return Err(Error::model(
                "observation absent; a predict request must carry an observation",
            ));
        }
        let contract = self
            .env_contract
            .as_ref()
            .ok_or_else(|| Error::model("observation missing env contract; cannot decode"))?;
        if contract.observation_space.is_none() {
            return Err(Error::model("env contract missing observation space"));
        }
        Ok(())
    }

    /// Decode the observation into one typed [`SpaceValue`](spaces::SpaceValue) per lane (length
    /// `== num_envs`), using the route's pinned observation space.
    ///
    /// Errors if the observation is absent — an absent value yields no lanes, so
    /// returning it as an empty `Vec` would silently violate the `== num_envs`
    /// contract. Guard with `observation.is_some()` for optional observations.
    pub fn decoded_lanes(&self) -> Result<Vec<spaces::SpaceValue>> {
        let leaves = self.observation.as_ref().ok_or_else(|| {
            Error::model("observation absent; cannot decode lanes (check is_some() first)")
        })?;
        self.decode_rows(leaves, self.num_envs)
    }

    /// Decode one history frame into one typed value per row (length `==
    /// frame.episodes.len()`), with the same observation space as the request.
    pub fn decoded_history_lanes(&self, frame: &HistoryFrame) -> Result<Vec<spaces::SpaceValue>> {
        let leaves = frame
            .observation
            .as_ref()
            .ok_or_else(|| Error::model("history frame carries no observation"))?;
        self.decode_rows(leaves, frame.episodes.len())
    }

    fn decode_rows(&self, leaves: &[Bytes], rows: usize) -> Result<Vec<spaces::SpaceValue>> {
        let contract = self
            .env_contract
            .as_ref()
            .ok_or_else(|| Error::model("observation missing env contract; cannot decode"))?;
        let space = contract
            .observation_space
            .as_ref()
            .ok_or_else(|| Error::model("env contract missing observation space"))?;
        let value = rlmesh_grpc::wire::leaves_value(leaves.to_vec());
        rlmesh_grpc::wire::decode_batched_partial_values(Some(&value), space, rows)
            .map_err(|err| Error::model(err.to_string()))
    }

    /// Decode a single-env observation. Errors unless exactly one lane is present.
    pub fn decoded(&self) -> Result<spaces::SpaceValue> {
        let mut lanes = self.decoded_lanes()?;
        if lanes.len() != 1 {
            return Err(Error::model(format!(
                "decoded() requires a single-env observation, got {} lanes",
                lanes.len()
            )));
        }
        Ok(lanes.remove(0))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::spaces::{DType, SpaceValue, Tensor};

    fn obs_space() -> spaces::SpaceSpec {
        spaces::spaces::BoxSpaceBuilder::scalar(0.0, 255.0, vec![1])
            .dtype(DType::Uint8)
            .build()
            .unwrap()
    }

    fn box_u8(v: u8) -> SpaceValue {
        SpaceValue::Box(Tensor::from_vec(vec![v], vec![1], DType::Uint8).unwrap())
    }

    /// A `ModelObservation` carrying `values` encoded against `obs_space()`.
    fn observation(values: &[SpaceValue]) -> ModelObservation {
        let space = obs_space();
        let wire = rlmesh_grpc::wire::encode_batched_partial_values(values, &space).unwrap();
        let contract = spaces::EnvContract {
            id: "T".into(),
            action_space: None,
            observation_space: Some(space),
            metadata: None,
            render_mode: String::new(),
            num_envs: values.len() as u32,
            autoreset_mode: Default::default(),
        };
        ModelObservation {
            history: Vec::new(),
            step: None,
            observation: Some(wire.leaves),
            route: ModelRouteContext::default(),
            num_envs: values.len(),
            env_contract: Some(Arc::new(contract)),
        }
    }

    #[test]
    fn predict_seed_law() {
        // Pure function of (episode_seed, predict_index): same inputs, same seed.
        assert_eq!(predict_seed(7, 3), predict_seed(7, 3));
        // Both ingredients move the result -- a per-episode seed alone would give
        // every re-plan of an episode the same noise.
        assert_ne!(predict_seed(7, 0), predict_seed(7, 1));
        assert_ne!(predict_seed(7, 0), predict_seed(8, 0));
        // Always a non-negative u32: the range every framework's seeding API takes.
        for (seed, index) in [(0, 0), (-1, 0), (i64::MIN, u64::MAX), (i64::MAX, 1)] {
            let derived = predict_seed(seed, index);
            assert!((0..=i64::from(u32::MAX)).contains(&derived), "{derived}");
        }
    }

    #[test]
    fn decoded_lanes_roundtrips_each_lane() {
        let lanes = vec![box_u8(5), box_u8(9), box_u8(0)];
        assert_eq!(observation(&lanes).decoded_lanes().unwrap(), lanes);
    }

    #[test]
    fn ensure_decodable_catches_malformed_requests_without_decoding() {
        // Happy path: present observation + contract + space validates.
        assert!(observation(&[box_u8(3)]).ensure_decodable().is_ok());
        // Absent observation errors up front.
        let mut no_obs = observation(&[box_u8(3)]);
        no_obs.observation = None;
        assert!(no_obs.ensure_decodable().is_err());
        // Missing env contract errors up front -- the gap that used to be deferred
        // on a step where every lane is mid-chunk-replay and skips the decode.
        let mut no_contract = observation(&[box_u8(3)]);
        no_contract.env_contract = None;
        assert!(no_contract.ensure_decodable().is_err());
    }

    #[test]
    fn decoded_requires_exactly_one_lane() {
        // Single lane decodes to the value; multi-lane errors rather than
        // silently taking lane 0.
        assert_eq!(observation(&[box_u8(7)]).decoded().unwrap(), box_u8(7));
        assert!(observation(&[box_u8(1), box_u8(2)]).decoded().is_err());
    }

    #[test]
    fn absent_observation_errors_not_empty_vec() {
        // An absent observation must error, never return an empty Vec that would
        // silently satisfy a `== num_envs` check for num_envs == 0 callers.
        let mut obs = observation(&[box_u8(3)]);
        obs.observation = None;
        assert!(obs.decoded_lanes().is_err());
    }
}