frxd/src/message.rs

use crate::crypto::{Keypair, now_ts, random_nonce, signing_bytes, verify_envelope};
use anyhow::{Context, Result, anyhow};
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::Value;

pub const TYPE_QUERY: &str = "query";
pub const TYPE_RESPONSE: &str = "response";
pub const TYPE_AGGREGATE: &str = "aggregate";

pub const EXPOSURE_METADATA: &str = "metadata";
pub const EXPOSURE_FULL: &str = "full";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Envelope {
    #[serde(rename = "type")]
    pub msg_type: String,
    pub from: String,
    pub ts: u64,
    pub nonce: String,
    pub body: Value,
    pub sig: String,
}

impl Envelope {
    pub fn new(key: &Keypair, msg_type: &str, body: Value) -> Self {
        let mut envelope = Self {
            msg_type: msg_type.to_string(),
            from: key.public_hex(),
            ts: now_ts(),
            nonce: random_nonce(),
            body,
            sig: String::new(),
        };
        envelope.sig = key.sign(&signing_bytes(&envelope));
        envelope
    }

    pub fn verify(&self) -> Result<()> {
        verify_envelope(self)
    }

    pub fn parse_body<T: DeserializeOwned>(&self) -> Result<T> {
        serde_json::from_value(self.body.clone()).context("malformed body")
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Budget {
    pub max_results: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryBody {
    pub qid: String,
    pub text: String,
    #[serde(default)]
    pub entities: Vec<String>,
    pub budget: Budget,
}

impl QueryBody {
    pub fn new(text: &str, max_results: usize) -> Self {
        Self {
            qid: crate::crypto::random_id(),
            text: text.to_string(),
            entities: Vec::new(),
            budget: Budget {
                max_results: max_results.clamp(1, 1000),
            },
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResponseItem {
    pub url: String,
    pub title: String,
    pub summary: String,
    pub published: String,
    pub exposure: String,
    pub content: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResponseBody {
    pub qid: String,
    pub results: Vec<ResponseItem>,
    pub truncated: bool,
    pub more_available: u64,
    pub cursor: Option<String>,
}

pub fn build_response(
    qid: &str,
    hits: Vec<ResponseItem>,
    total: u64,
    max_results: usize,
) -> ResponseBody {
    let max_results = max_results.max(1);
    let mut results = hits;
    results.truncate(max_results);
    let more_available = total.saturating_sub(results.len() as u64);
    ResponseBody {
        qid: qid.to_string(),
        results,
        truncated: more_available > 0,
        more_available,
        cursor: None,
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AggregateBody {
    pub period: String,
    pub sent: u64,
    pub passed: u64,
}

pub fn require_type(envelope: &Envelope, expected: &str) -> Result<()> {
    if envelope.msg_type != expected {
        return Err(anyhow!(
            "expected message type {expected}, got {}",
            envelope.msg_type
        ));
    }
    Ok(())
}

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

    fn item(n: usize) -> ResponseItem {
        ResponseItem {
            url: format!("file:///doc{n}.txt"),
            title: format!("doc {n}"),
            summary: String::new(),
            published: String::new(),
            exposure: EXPOSURE_METADATA.to_string(),
            content: None,
        }
    }

    #[test]
    fn max_results_is_respected() {
        let response = build_response("q1", vec![item(1), item(2), item(3)], 10, 2);
        assert_eq!(response.results.len(), 2);
    }

    #[test]
    fn truncation_is_honest() {
        let response = build_response("q1", vec![item(1), item(2), item(3)], 10, 2);
        assert!(response.truncated);
        assert_eq!(response.more_available, 8);

        let response = build_response("q2", vec![item(1), item(2)], 2, 5);
        assert!(!response.truncated);
        assert_eq!(response.more_available, 0);
    }

    #[test]
    fn response_carries_no_scores() {
        let response = build_response("q1", vec![item(1)], 1, 5);
        let json = serde_json::to_string(&response).unwrap();
        assert!(!json.contains("score"));
        assert!(!json.contains("relevance"));
    }

    #[test]
    fn zero_budget_is_clamped() {
        let query = QueryBody::new("anything", 0);
        assert_eq!(query.budget.max_results, 1);
    }

    #[test]
    fn envelope_rejects_wrong_type() {
        let key = Keypair::generate();
        let envelope = Envelope::new(&key, TYPE_RESPONSE, json!({}));
        assert!(require_type(&envelope, TYPE_QUERY).is_err());
    }
}