laser-sdk 0.8.2

use crate::error::LaserError;
use crate::fork::{ForkError, ForkInfo, ForkKind};
use crate::laser::Laser;
use crate::query::{
    AGDX_FORK_CREATE_CODE, AGDX_FORK_DELETE_CODE, AGDX_FORK_LIST_CODE, AGDX_FORK_PROMOTE_CODE,
    AGDX_FORK_PUT_CODE, FORK_OP_VERSION,
};
use laser_wire::fork::{
    ForkCreate, ForkDelete, ForkList, ForkOutcome, ForkPromote, ForkPut, ForkReply,
};
use laser_wire::framing::encode_named;
#[cfg(test)]
use laser_wire::limits::MAX_FORK_ID_BYTES;
use serde::Serialize;
use std::collections::BTreeMap;

impl Laser {
    /// A handle to one fork by id. Cheap to create, since it borrows the connection.
    /// Open the fork with [`ForkHandle::create`], then write rows, query it
    /// (`query(...).fork(id)`), and finally [`promote`](ForkHandle::promote) or
    /// [`squash`](ForkHandle::squash) it. The id is sent as
    /// [`resource_name`](Self::resource_name) names it.
    pub fn fork<'a>(&'a self, fork_id: impl Into<String>) -> ForkHandle<'a> {
        ForkHandle {
            laser: self,
            fork_id: self.resource_name(&fork_id.into()),
        }
    }

    /// Every open fork for the authenticated user. A handle that scopes its
    /// resources to a stream lists only that stream's forks, under the ids the
    /// caller gave them.
    pub async fn forks(&self) -> Result<Vec<ForkInfo>, LaserError> {
        match self
            .execute_fork(AGDX_FORK_LIST_CODE, &ForkList { v: FORK_OP_VERSION })
            .await?
        {
            ForkOutcome::List(forks) => Ok(self.local_forks(forks)),
            other => Err(unexpected("list", &other)),
        }
    }

    fn local_forks(&self, forks: Vec<ForkInfo>) -> Vec<ForkInfo> {
        forks
            .into_iter()
            .filter_map(|mut fork| {
                fork.fork_id = self.local_resource_name(&fork.fork_id)?.to_owned();
                if let Some(parent) = fork.parent.take() {
                    let local = self.local_resource_name(&parent).map(str::to_owned);
                    fork.parent = Some(local.unwrap_or(parent));
                }
                Some(fork)
            })
            .collect()
    }

    // Send one fork command over the binary connection and decode the reply.
    // Gated on `forks` (set by the connect-time probe). Without it, `Unsupported`.
    pub(crate) async fn execute_fork(
        &self,
        code: u32,
        request: &impl Serialize,
    ) -> Result<ForkOutcome, LaserError> {
        let capabilities = self.capabilities().await;
        if !capabilities.forks {
            return Err(LaserError::unsupported(
                "fork",
                "forks are not served by this deployment",
            ));
        }
        // Fail fast on advertised version skew (see `execute_query`): a server
        // that advertised its accepted fork op version spares the round-trip.
        if let Some(versions) = capabilities.versions
            && versions.fork != FORK_OP_VERSION
        {
            return Err(ForkError::Version {
                expected: versions.fork,
                got: FORK_OP_VERSION,
            }
            .into());
        }
        let payload = encode_named(request)
            .map_err(|error| LaserError::Codec(format!("encode fork request: {error}")))?;
        let payload = self.send_raw_with_response(code, payload).await?;
        match crate::error::decode_managed_reply::<ForkReply>(&payload)? {
            ForkReply::Ok(outcome) => Ok(outcome),
            ForkReply::Err(error) => Err(error.into()),
            _ => Err(LaserError::Protocol(
                "fork: unknown reply variant".to_owned(),
            )),
        }
    }
}

/// A handle to one fork. Build it with [`Laser::fork`].
pub struct ForkHandle<'a> {
    laser: &'a Laser,
    fork_id: String,
}

impl<'a> ForkHandle<'a> {
    /// The fork id this handle is bound to, as the caller named it.
    pub fn id(&self) -> &str {
        self.laser
            .local_resource_name(&self.fork_id)
            .unwrap_or(&self.fork_id)
    }

    /// The fork id this handle sends, scoped by the connection's
    /// [`ResourceNaming`](crate::laser::ResourceNaming). Pass it to
    /// `query(..).fork(..)` or any other request naming this fork.
    pub fn resource_id(&self) -> &str {
        &self.fork_id
    }

    /// Start opening this fork. Choose `.severed()` or `.continuous()` (default
    /// continuous), optionally narrow the snapshot with `.tables([...])`, then
    /// `.send().await` for the fork's metadata.
    pub fn create(&self) -> ForkCreateRequest<'a> {
        ForkCreateRequest {
            laser: self.laser,
            fork_id: self.fork_id.clone(),
            parent: None,
            kind: ForkKind::Continuous,
            tables: Vec::new(),
        }
    }

    /// Promote this fork: splice its speculative rows onto the trunk (and apply
    /// its tombstones), then squash it. Returns the number of rows applied.
    pub async fn promote(&self) -> Result<usize, LaserError> {
        match self
            .laser
            .execute_fork(
                AGDX_FORK_PROMOTE_CODE,
                &ForkPromote {
                    v: FORK_OP_VERSION,
                    fork_id: self.fork_id.clone(),
                },
            )
            .await?
        {
            ForkOutcome::Promoted { rows } => Ok(rows),
            other => Err(unexpected("promote", &other)),
        }
    }

    /// Squash this fork: discard its speculative rows. Returns `true` when an open
    /// fork was removed, `false` when none existed.
    pub async fn squash(&self) -> Result<bool, LaserError> {
        match self
            .laser
            .execute_fork(
                AGDX_FORK_DELETE_CODE,
                &ForkDelete {
                    v: FORK_OP_VERSION,
                    fork_id: self.fork_id.clone(),
                },
            )
            .await?
        {
            ForkOutcome::Deleted(existed) => Ok(existed),
            other => Err(unexpected("squash", &other)),
        }
    }

    /// Write one speculative row into this fork at `(table, partition_id, offset)`.
    /// Add indexed fields with `.field`, an opaque body with `.payload`, an
    /// embedding with `.embedding`, or mark `.tombstone()` to hide the trunk row
    /// at that coordinate from the fork's view. Finish with `.send().await`.
    pub fn put_row(
        &self,
        table: impl Into<String>,
        partition_id: u32,
        offset: u64,
    ) -> ForkPutRequest<'a> {
        ForkPutRequest {
            laser: self.laser,
            fork_id: self.fork_id.clone(),
            table: self.laser.resource_name(&table.into()),
            partition_id,
            offset,
            projection_id: String::new(),
            projection_version: 0,
            fields: BTreeMap::new(),
            metadata: BTreeMap::new(),
            payload: None,
            embedding: None,
            tombstone: false,
        }
    }
}

/// Fluent builder for [`ForkHandle::create`].
#[must_use = "call .send().await to create the fork"]
pub struct ForkCreateRequest<'a> {
    laser: &'a Laser,
    fork_id: String,
    parent: Option<String>,
    kind: ForkKind,
    tables: Vec<String>,
}

impl<'a> ForkCreateRequest<'a> {
    /// Make a frozen snapshot at the trunk's current offsets (later appends hidden).
    pub fn severed(mut self) -> Self {
        self.kind = ForkKind::Severed;
        self
    }

    /// Make a live branch that keeps seeing new trunk appends (the default).
    pub fn continuous(mut self) -> Self {
        self.kind = ForkKind::Continuous;
        self
    }

    /// Record this fork's parent (audit only, the fork still branches off the trunk).
    pub fn parent(mut self, parent: impl Into<String>) -> Self {
        self.parent = Some(self.laser.resource_name(&parent.into()));
        self
    }

    /// Narrow a severed snapshot to these tables. Empty captures every table.
    pub fn tables(mut self, tables: impl IntoIterator<Item = impl Into<String>>) -> Self {
        self.tables = tables
            .into_iter()
            .map(|table| self.laser.resource_name(&table.into()))
            .collect();
        self
    }

    /// Open the fork. Returns its metadata.
    pub async fn send(self) -> Result<ForkInfo, LaserError> {
        validated_fork_id(&self.fork_id)?;
        let request = ForkCreate {
            v: FORK_OP_VERSION,
            fork_id: self.fork_id,
            parent: self.parent,
            kind: self.kind,
            tables: self.tables,
        };
        match self
            .laser
            .execute_fork(AGDX_FORK_CREATE_CODE, &request)
            .await?
        {
            ForkOutcome::Created(info) => Ok(info),
            other => Err(unexpected("create", &other)),
        }
    }
}

/// Fluent builder for [`ForkHandle::put_row`].
#[must_use = "call .send().await to write the speculative row"]
pub struct ForkPutRequest<'a> {
    laser: &'a Laser,
    fork_id: String,
    table: String,
    partition_id: u32,
    offset: u64,
    projection_id: String,
    projection_version: u32,
    fields: BTreeMap<String, String>,
    metadata: BTreeMap<String, String>,
    payload: Option<Vec<u8>>,
    embedding: Option<Vec<f32>>,
    tombstone: bool,
}

impl<'a> ForkPutRequest<'a> {
    /// Set the projection id/version this speculative row belongs to.
    pub fn projection(mut self, id: impl Into<String>, version: u32) -> Self {
        self.projection_id = self.laser.resource_name(&id.into());
        self.projection_version = version;
        self
    }

    /// Add one indexed field (the columns queries filter and order on).
    pub fn field(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.fields.insert(name.into(), value.into());
        self
    }

    /// Add one metadata header (non-indexed).
    pub fn metadata(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.metadata.insert(name.into(), value.into());
        self
    }

    /// Attach an opaque payload body.
    pub fn payload(mut self, payload: impl AsRef<[u8]>) -> Self {
        self.payload = Some(payload.as_ref().to_vec());
        self
    }

    /// Attach an embedding vector. It travels as a JSON array literal, and
    /// `send` rejects a non-finite component with `LaserError::Invalid`.
    pub fn embedding(mut self, embedding: impl IntoIterator<Item = f32>) -> Self {
        self.embedding = Some(embedding.into_iter().collect());
        self
    }

    /// Mark this as a tombstone: hide the trunk row at this coordinate from the
    /// fork's view instead of writing a value.
    pub fn tombstone(mut self) -> Self {
        self.tombstone = true;
        self
    }

    /// Write the speculative row.
    pub async fn send(self) -> Result<(), LaserError> {
        let embedding = self
            .embedding
            .as_deref()
            .map(embedding_literal)
            .transpose()?;
        let request = ForkPut {
            v: FORK_OP_VERSION,
            fork_id: self.fork_id,
            table: self.table,
            partition_id: self.partition_id,
            offset: self.offset,
            projection_id: self.projection_id,
            projection_version: self.projection_version,
            fields: self.fields,
            metadata: self.metadata,
            payload: self.payload,
            embedding,
            tombstone: self.tombstone,
        };
        match self
            .laser
            .execute_fork(AGDX_FORK_PUT_CODE, &request)
            .await?
        {
            ForkOutcome::Written => Ok(()),
            other => Err(unexpected("put_row", &other)),
        }
    }
}

fn unexpected(op: &str, outcome: &ForkOutcome) -> LaserError {
    LaserError::Protocol(format!("fork {op}: unexpected outcome {outcome:?}"))
}

// Reject a malformed fork id before the round-trip, using the wire contract's
// canonical `validate_fork_id` (non-empty, within `MAX_FORK_ID_BYTES`, and the
// shared charset safelist) so a client fails fast on exactly what the managed
// side would reject, and the anti-injection charset rule lives in one place.
fn validated_fork_id(fork_id: &str) -> Result<(), LaserError> {
    laser_wire::fork::validate_fork_id(fork_id).map_err(|error| LaserError::Invalid(error.0))
}

fn embedding_literal(embedding: &[f32]) -> Result<String, LaserError> {
    if embedding.iter().any(|component| !component.is_finite()) {
        return Err(LaserError::Invalid(
            "an embedding component must be a finite number".to_owned(),
        ));
    }
    serde_json::to_string(embedding).map_err(|error| LaserError::Codec(error.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn info(fork_id: &str, parent: Option<&str>) -> ForkInfo {
        ForkInfo {
            fork_id: fork_id.to_owned(),
            parent: parent.map(str::to_owned),
            kind: ForkKind::Continuous,
            user_id: 1,
            status: laser_wire::fork::ForkStatus::Open,
            created_at_micros: 0,
            row_count: 0,
        }
    }

    #[test]
    fn given_a_default_stream_when_opening_a_fork_then_should_scope_its_names() {
        let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default())
            .with_default_stream("acme");
        let fork = laser.fork("experiment");
        assert_eq!(fork.resource_id(), "stream:acme/experiment");
        assert_eq!(fork.id(), "experiment");
        validated_fork_id(fork.resource_id()).expect("a scoped fork id is valid");
        let create = fork.create().parent("baseline").tables(["readings"]);
        assert_eq!(create.parent.as_deref(), Some("stream:acme/baseline"));
        assert_eq!(create.tables, ["stream:acme/readings"]);
        let put = fork.put_row("readings", 0, 1).projection("reading.v1", 1);
        assert_eq!(put.table, "stream:acme/readings");
        assert_eq!(put.projection_id, "stream:acme/reading.v1");
    }

    #[test]
    fn given_listed_forks_when_localized_then_should_keep_only_the_own_stream() {
        let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default())
            .with_default_stream("acme");
        let local = laser.local_forks(vec![
            info("stream:acme/experiment", Some("stream:acme/baseline")),
            info("stream:other/experiment", None),
            info("legacy", None),
        ]);
        assert_eq!(local.len(), 1);
        assert_eq!(local[0].fork_id, "experiment");
        assert_eq!(local[0].parent.as_deref(), Some("baseline"));
    }

    #[test]
    fn given_an_over_long_or_empty_fork_id_when_validated_then_should_reject() {
        assert!(validated_fork_id("").is_err());
        assert!(validated_fork_id("experiment-2026-q2").is_ok());
        let too_long = "f".repeat(MAX_FORK_ID_BYTES + 1);
        assert!(validated_fork_id(&too_long).is_err());
        let at_cap = "f".repeat(MAX_FORK_ID_BYTES);
        assert!(validated_fork_id(&at_cap).is_ok());
    }

    #[test]
    fn given_fork_error_when_mapped_then_should_nest_the_typed_error() {
        let error: LaserError = ForkError::NotFound("f9".to_owned()).into();
        assert!(matches!(error, LaserError::Fork(ForkError::NotFound(_))));
        assert!(error.is_not_found());
        let unsupported: LaserError = ForkError::Unsupported("x".to_owned()).into();
        assert!(matches!(
            unsupported,
            LaserError::Fork(ForkError::Unsupported(_))
        ));
        assert!(unsupported.is_unsupported());
    }

    #[test]
    fn given_an_embedding_when_encoded_then_should_be_a_json_array_and_reject_non_finite() {
        assert_eq!(
            embedding_literal(&[0.5, 1.0]).expect("finite components encode"),
            "[0.5,1.0]"
        );
        assert!(matches!(
            embedding_literal(&[f32::NAN]),
            Err(LaserError::Invalid(_))
        ));
    }
}