aimo-core 0.1.9

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use solana_sdk::pubkey::Pubkey;

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RequestReceipt {
    pub sender: String,
    pub request_id: String,
    pub service_id: String,
    pub max_est_cost: u64,
    pub timestamp: i64,
    pub finish_reason: String,
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub amount: u64,
    pub token_mint: String,
    pub model_name: String,
    pub cache_discount: u64,

    // Model Statistics
    pub streamed: Option<bool>,
    pub cancelled: Option<bool>,
    pub latency: Option<i64>,
    pub generation_time: Option<i64>,
    pub moderation_latency: Option<i64>,
    pub num_media_prompt: Option<u64>,
    pub num_media_completion: Option<u64>,
    pub num_search_results: Option<u64>,

    // Optional Fields
    pub app_id: Option<String>,
    pub origin: Option<String>,
}

impl RequestReceipt {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn sender(self, sender: Pubkey) -> Self {
        Self {
            sender: sender.to_string(),
            ..self
        }
    }

    pub fn request_id(self, request_id: Pubkey) -> Self {
        Self {
            request_id: request_id.to_string(),
            ..self
        }
    }

    pub fn service_id(self, service_id: impl Into<String>) -> Self {
        Self {
            service_id: service_id.into(),
            ..self
        }
    }

    pub fn max_est_cost(self, max_est_cost: u64) -> Self {
        Self {
            max_est_cost,
            ..self
        }
    }

    pub fn timestamp(self, timestamp: i64) -> Self {
        Self {
            timestamp,
            ..self
        }
    }

    pub fn finish_reason(self, finish_reason: impl Into<String>) -> Self {
        Self {
            finish_reason: finish_reason.into(),
            ..self
        }
    }

    pub fn prompt_tokens(self, prompt_tokens: u64) -> Self {
        Self {
            prompt_tokens,
            ..self
        }
    }

    pub fn completion_tokens(self, completion_tokens: u64) -> Self {
        Self {
            completion_tokens,
            ..self
        }
    }

    pub fn amount(self, amount: u64) -> Self {
        Self {
            amount,
            ..self
        }
    }

    pub fn token_mint(self, token_mint: impl Into<String>) -> Self {
        Self {
            token_mint: token_mint.into(),
            ..self
        }
    }

    pub fn model_name(self, model_name: impl Into<String>) -> Self {
        Self {
            model_name: model_name.into(),
            ..self
        }
    }

    pub fn cache_discount(self, cache_discount: u64) -> Self {
        Self {
            cache_discount,
            ..self
        }
    }

    pub fn streamed(self, streamed: bool) -> Self {
        Self {
            streamed: Some(streamed),
            ..self
        }
    }

    pub fn cancelled(self, cancelled: bool) -> Self {
        Self {
            cancelled: Some(cancelled),
            ..self
        }
    }

    pub fn latency(self, latency: i64) -> Self {
        Self {
            latency: Some(latency),
            ..self
        }
    }

    pub fn generation_time(self, generation_time: i64) -> Self {
        Self {
            generation_time: Some(generation_time),
            ..self
        }
    }

    pub fn moderation_latency(self, moderation_latency: i64) -> Self {
        Self {
            moderation_latency: Some(moderation_latency),
            ..self
        }
    }

    pub fn num_media_prompt(self, num_media_prompt: u64) -> Self {
        Self {
            num_media_prompt: Some(num_media_prompt),
            ..self
        }
    }

    pub fn num_media_completion(self, num_media_completion: u64) -> Self {
        Self {
            num_media_completion: Some(num_media_completion),
            ..self
        }
    }

    pub fn num_search_results(self, num_search_results: u64) -> Self {
        Self {
            num_search_results: Some(num_search_results),
            ..self
        }
    }

    pub fn app_id(self, app_id: impl Into<String>) -> Self {
        Self {
            app_id: Some(app_id.into()),
            ..self
        }
    }

    pub fn origin(self, origin: impl Into<String>) -> Self {
        Self {
            origin: Some(origin.into()),
            ..self
        }
    }
}

impl RequestReceipt {
    pub fn to_bytes(self) -> anyhow::Result<Vec<u8>> {
        let json = serde_json::to_value(self)?;
        let string = canonical_json::to_string(&json)?;
        Ok(string.into_bytes())
    }

    pub fn to_hash(self) -> anyhow::Result<[u8; 32]> {
        let mut hasher = Sha256::new();
        hasher.update(self.to_bytes()?);
        Ok(hasher.finalize().into())
    }
}