haem-frames 0.1.1

use std::collections::BTreeSet;
use std::path::Path;

use super::{
    Embed, Error, FailedCode, ItemCode, Kind, Limits, MINIMUM_FRAME_BYTES, Model, Outcome,
    PROTOCOL, PROTOCOL_WINDOWS, Ready, Start, Vectors, Window, Windows, decode_vector,
};

fn limits(message: impl Into<String>) -> Error {
    Error::new(FailedCode::EmbedLimits, message)
}

/// Validate the protocol, model path and positive exchange limits.
///
/// # Errors
/// Returns `embed_protocol` for an unsupported protocol or windows shape, and
/// `embed_limits` for invalid limits or a relative model path.
pub fn check_start(start: &Start) -> Result<(), Error> {
    // EMB-E2: windows are asked for by protocol 2 and only by it.
    let windows = match (start.protocol, start.windows) {
        (PROTOCOL, None) => None,
        (PROTOCOL_WINDOWS, Some(windows)) => Some(windows),
        _ => {
            return Err(Error::new(
                FailedCode::EmbedProtocol,
                "unsupported embedding protocol",
            ));
        }
    };
    let value = &start.limits;
    if [
        value.memory_bytes,
        value.threads,
        value.batch_items,
        value.input_bytes,
        value.tokens,
    ]
    .contains(&0)
        || value.frame_bytes < MINIMUM_FRAME_BYTES
    {
        return Err(limits(
            "all limits must be positive and frame_bytes must be at least 8192",
        ));
    }
    if !Path::new(&start.model_dir).is_absolute() {
        return Err(limits("model_dir must be absolute"));
    }
    if windows.is_some_and(|windows| {
        windows.max_windows < 2 || windows.overlap_tokens >= start.limits.tokens
    }) {
        return Err(limits(
            "max_windows must be at least 2 and overlap_tokens below tokens",
        ));
    }
    Ok(())
}

/// Ready names protocol 1 or, for a Start that asked for windows, 2; the
/// service checks it names the protocol its Start did.
///
/// # Errors
/// Returns `embed_protocol` for an invalid protocol or model identity.
pub fn check_ready(ready: &Ready) -> Result<(), Error> {
    let model = &ready.model;
    if ![PROTOCOL, PROTOCOL_WINDOWS].contains(&ready.protocol)
        || model.name.is_empty()
        || model.manifest_sha256.len() != 64
        || !model
            .manifest_sha256
            .bytes()
            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
        || model.dimensions == 0
        || model.max_tokens == 0
        || model.special_tokens >= model.max_tokens
    {
        return Err(Error::new(
            FailedCode::EmbedProtocol,
            "Ready has invalid protocol or model identity",
        ));
    }
    Ok(())
}

/// Upper bound for tag plus compact JSON, not the usual message size.
/// Includes maximally escaped ids and the widest whole-number token counts.
///
/// # Errors
/// Returns `embed_limits` for zero counts or an overflowing message size.
pub fn vectors_message_size(batch_items: u64, dimensions: u64) -> Result<u64, Error> {
    if batch_items == 0 || dimensions == 0 {
        return Err(limits("batch_items and dimensions must be positive"));
    }
    // Each legal single-byte control character can take six JSON bytes.
    let id = 512_u64 * 6;
    let fixed = u64::try_from(
        b"{\"outcome\":\"vector\",\"id\":\"\",\"vector\":\"\",\"tokens_read\":,\"tokens_total\":}"
            .len(),
    )
    .map_err(|_| limits("message size overflows"))?;
    let item = dimensions
        .checked_mul(8)
        .and_then(|vector| vector.checked_add(id + fixed + 40))
        .ok_or_else(|| limits("message size overflows"))?;
    let items = batch_items
        .checked_mul(item)
        .and_then(|value| value.checked_add(batch_items - 1))
        .ok_or_else(|| limits("message size overflows"))?;
    items
        .checked_add(13)
        .ok_or_else(|| limits("message size overflows"))
}

/// EMB-E2: [`vectors_message_size`] when each item may answer up to
/// `max_windows` windows, each its bounds, its vector and its token count.
///
/// # Errors
/// `embed_limits` when a count is zero or the size overflows.
pub fn windows_message_size(
    batch_items: u64,
    dimensions: u64,
    max_windows: u64,
) -> Result<u64, Error> {
    let overflow = || limits("message size overflows");
    let fixed = |bytes: &[u8]| u64::try_from(bytes.len()).map_err(|_| overflow());
    let item = fixed(b"{\"outcome\":\"windows\",\"id\":\"\",\"windows\":[],\"tokens_total\":}")?;
    let window = fixed(b"{\"start\":,\"end\":,\"vector\":\"\",\"tokens_read\":},")?;
    // Each legal single-byte control character can take six JSON bytes.
    let id = 512_u64 * 6;
    let windows = dimensions
        .checked_mul(8)
        .and_then(|vector| vector.checked_add(window + 60))
        .and_then(|one| one.checked_mul(max_windows))
        .ok_or_else(overflow)?;
    let one = windows.checked_add(item + id + 20).ok_or_else(overflow)?;
    let ordinary = vectors_message_size(batch_items, dimensions)?;
    batch_items
        .checked_mul(one)
        .and_then(|items| items.checked_add(batch_items + 13))
        .map(|size| size.max(ordinary))
        .ok_or_else(overflow)
}

/// Check that token and frame limits fit the model and its largest answer.
///
/// # Errors
/// Returns `embed_limits`, naming the incompatible limit or size overflow.
pub fn check_model_limits(limits_value: &Limits, model: &Model) -> Result<(), Error> {
    check_model_limits_windowed(limits_value, model, None)
}

/// [`check_model_limits`], and with `windows` (EMB-E2) a frame that carries
/// the largest windowed answer and an overlap below the model's content
/// tokens, so every window moves forward through its text.
///
/// # Errors
/// `embed_limits`, naming the limit.
pub fn check_model_limits_windowed(
    limits_value: &Limits,
    model: &Model,
    windows: Option<Windows>,
) -> Result<(), Error> {
    let minimum_tokens = model
        .special_tokens
        .checked_add(model.prefix_tokens)
        .and_then(|reserved| reserved.checked_add(1));
    if limits_value.tokens > model.max_tokens
        || minimum_tokens.is_none_or(|minimum| limits_value.tokens < minimum)
    {
        return Err(limits(format!(
            "tokens {} must be at least special_tokens {} + prefix_tokens {} + 1 and not exceed max_tokens {}",
            limits_value.tokens, model.special_tokens, model.prefix_tokens, model.max_tokens
        )));
    }
    let largest = match windows {
        None => vectors_message_size(limits_value.batch_items, model.dimensions)?,
        Some(windows) => windows_message_size(
            limits_value.batch_items,
            model.dimensions,
            windows.max_windows,
        )?,
    };
    if u64::from(limits_value.frame_bytes) < largest {
        return Err(limits(
            "frame_bytes cannot carry the largest legal Vectors batch",
        ));
    }
    let content = model
        .max_tokens
        .saturating_sub(model.special_tokens)
        .saturating_sub(model.prefix_tokens);
    if windows.is_some_and(|windows| windows.overlap_tokens >= content) {
        return Err(limits(format!(
            "overlap_tokens must be below the model's {content} content tokens"
        )));
    }
    Ok(())
}

fn id_valid(id: &str) -> bool {
    !id.is_empty() && id.len() <= 512 && !id.contains('\0')
}

/// Check a batch has a bounded item count and distinct valid identifiers.
///
/// # Errors
/// Returns `embed_protocol` for an empty batch or invalid ids, and
/// `embed_batch_too_large` when the item count exceeds the limit.
pub fn check_embed(embed: &Embed, limits_value: &Limits) -> Result<(), Error> {
    if embed.items.is_empty() {
        return Err(Error::new(
            FailedCode::EmbedProtocol,
            "Embed must contain at least one item",
        ));
    }
    let count = u64::try_from(embed.items.len()).map_err(|_| {
        Error::new(
            FailedCode::EmbedBatchTooLarge,
            "batch count exceeds its limit",
        )
    })?;
    if count > limits_value.batch_items {
        return Err(Error::new(
            FailedCode::EmbedBatchTooLarge,
            "Embed exceeds batch_items",
        ));
    }
    let mut seen = BTreeSet::new();
    for item in &embed.items {
        if !id_valid(&item.id) || !seen.insert(item.id.as_str()) {
            return Err(Error::new(
                FailedCode::EmbedProtocol,
                "Embed has a bad or repeated id",
            ));
        }
    }
    Ok(())
}

pub fn input_refusal(text: &str, input_bytes: u64) -> Option<ItemCode> {
    if text.chars().all(char::is_whitespace) {
        return Some(ItemCode::EmbedInputEmpty);
    }
    match u64::try_from(text.len()) {
        Ok(bytes) if bytes <= input_bytes => None,
        Ok(_) | Err(_) => Some(ItemCode::EmbedInputTooLarge),
    }
}

fn invalid_output() -> Error {
    Error::new(
        FailedCode::EmbedOutputInvalid,
        "Vectors content does not satisfy the request and model",
    )
}

/// The content tokens the model reads after its special and prefix tokens.
const fn content_tokens(ready: &Ready) -> u64 {
    ready
        .model
        .max_tokens
        .saturating_sub(ready.model.special_tokens)
        .saturating_sub(ready.model.prefix_tokens)
}

/// A refusal is valid when it names the byte classifier's own code and carries no token counts,
/// or, for a text the classifier passed, says `text_longer_than_model` with the text's content
/// tokens beyond what the model reads after its special and prefix tokens, and the model's limit
/// as Ready declared it. EMB-E2: `windows_over_cap` is the same, and only when windows were
/// asked for.
fn refusal_fits(
    code: ItemCode,
    expected: Option<ItemCode>,
    (tokens_total, tokens_limit): (Option<u64>, Option<u64>),
    ready: &Ready,
    windowed: bool,
) -> bool {
    let longer = |total: u64, limit: u64| {
        expected.is_none() && limit == ready.model.max_tokens && total > content_tokens(ready)
    };
    match (code, tokens_total, tokens_limit) {
        (ItemCode::TextLongerThanModel, Some(total), Some(limit)) => longer(total, limit),
        (ItemCode::WindowsOverCap, Some(total), Some(limit)) => windowed && longer(total, limit),
        (ItemCode::TextLongerThanModel | ItemCode::WindowsOverCap, _, _)
        | (_, Some(_), _)
        | (_, _, Some(_)) => false,
        (_, None, None) => expected == Some(code),
    }
}

/// One encoded vector of the model's dimensions, of length one when the
/// model says its vectors are normalised.
fn check_vector(vector: &str, ready: &Ready) -> Result<(), Error> {
    let decoded = decode_vector(vector)?;
    let dimensions = u64::try_from(decoded.len()).map_err(|_| invalid_output())?;
    if dimensions != ready.model.dimensions {
        return Err(invalid_output());
    }
    if ready.model.normalised {
        let squared = decoded
            .iter()
            .map(|value| {
                let member = f64::from(*value);
                member * member
            })
            .sum::<f64>();
        if (squared.sqrt() - 1.0).abs() > 0.0001 {
            return Err(invalid_output());
        }
    }
    Ok(())
}

/// EMB-E2: a text's windows cover it whole. The first starts at 0 and the
/// last ends at its length; each starts after the one before it starts, no
/// later than it ends, and ends after it; every bound is on a char boundary;
/// there are 2 to `max_windows` of them, for a text longer than the model
/// reads; and each window was read whole.
fn check_windows(
    windows: &[Window],
    tokens_total: u64,
    text: &str,
    ready: &Ready,
    limits_value: &Limits,
    asked: Windows,
) -> Result<(), Error> {
    let count = u64::try_from(windows.len()).map_err(|_| invalid_output())?;
    let length = u64::try_from(text.len()).map_err(|_| invalid_output())?;
    if count < 2 || count > asked.max_windows || tokens_total <= content_tokens(ready) {
        return Err(invalid_output());
    }
    let boundary = |at: u64| usize::try_from(at).is_ok_and(|at| text.is_char_boundary(at));
    let mut previous: Option<&Window> = None;
    for window in windows {
        let follows = previous.map_or(window.start == 0, |before| {
            (before.start..=before.end).contains(&window.start)
                && window.start != before.start
                && window.end > before.end
        });
        if !follows
            || window.end <= window.start
            || window.end > length
            || !boundary(window.start)
            || !boundary(window.end)
            || window.tokens_read == 0
            || window.tokens_read > limits_value.tokens
        {
            return Err(invalid_output());
        }
        check_vector(&window.vector, ready)?;
        previous = Some(window);
    }
    if previous.map(|last| last.end) != Some(length) {
        return Err(invalid_output());
    }
    Ok(())
}

/// Check answer count, order, vectors and refusal outcomes against the request.
///
/// # Errors
/// Returns `embed_protocol` for a count or order mismatch, and
/// `embed_output_invalid` for outcomes the request and model do not allow.
pub fn check_vectors(
    vectors: &Vectors,
    embed: &Embed,
    ready: &Ready,
    limits_value: &Limits,
) -> Result<(), Error> {
    check_vectors_windowed(vectors, embed, ready, limits_value, None)
}

/// [`check_vectors`] for a worker started with `windows` (EMB-E2): an item of
/// a document Embed may also answer its windows, or `windows_over_cap`; a
/// question is never windowed.
///
/// # Errors
/// `embed_protocol` for a count or order that does not match the Embed;
/// `embed_output_invalid` for any outcome the request and model do not allow.
pub fn check_vectors_windowed(
    vectors: &Vectors,
    embed: &Embed,
    ready: &Ready,
    limits_value: &Limits,
    windows: Option<Windows>,
) -> Result<(), Error> {
    if vectors.items.len() != embed.items.len() {
        return Err(Error::new(
            FailedCode::EmbedProtocol,
            "Vectors count does not match Embed",
        ));
    }
    let windows = windows.filter(|_| embed.kind == Kind::Document);
    for (outcome, input) in vectors.items.iter().zip(&embed.items) {
        if outcome.id() != input.id {
            return Err(Error::new(
                FailedCode::EmbedProtocol,
                "Vectors ids do not match Embed order",
            ));
        }
        let expected = input_refusal(&input.text, limits_value.input_bytes);
        match outcome {
            Outcome::Refused {
                code,
                tokens_total,
                tokens_limit,
                ..
            } => {
                let counts = (*tokens_total, *tokens_limit);
                if !refusal_fits(*code, expected, counts, ready, windows.is_some()) {
                    return Err(invalid_output());
                }
            }
            Outcome::Vector {
                vector,
                tokens_read,
                tokens_total,
                ..
            } => {
                if expected.is_some()
                    || *tokens_read == 0
                    || tokens_read > tokens_total
                    || *tokens_read > limits_value.tokens
                {
                    return Err(invalid_output());
                }
                check_vector(vector, ready)?;
            }
            Outcome::Windows {
                windows: made,
                tokens_total,
                ..
            } => {
                let Some(asked) = windows.filter(|_| expected.is_none()) else {
                    return Err(invalid_output());
                };
                check_windows(made, *tokens_total, &input.text, ready, limits_value, asked)?;
            }
        }
    }
    Ok(())
}