botlib/src/models.rs

use crate::message_types::MessageType;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ApiResponse<T> {
    pub success: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<T>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub code: Option<String>,
}

impl<T> ApiResponse<T> {
    pub fn success(data: T) -> Self {
        Self {
            success: true,
            data: Some(data),
            error: None,
            message: None,
            code: None,
        }
    }

    pub fn success_with_message(data: T, message: impl Into<String>) -> Self {
        Self {
            success: true,
            data: Some(data),
            error: None,
            message: Some(message.into()),
            code: None,
        }
    }

    pub fn error(message: impl Into<String>) -> Self {
        Self {
            success: false,
            data: None,
            error: Some(message.into()),
            message: None,
            code: None,
        }
    }

    pub fn error_with_code(message: impl Into<String>, code: impl Into<String>) -> Self {
        Self {
            success: false,
            data: None,
            error: Some(message.into()),
            message: None,
            code: Some(code.into()),
        }
    }

    pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> ApiResponse<U> {
        ApiResponse {
            success: self.success,
            data: self.data.map(f),
            error: self.error,
            message: self.message,
            code: self.code,
        }
    }

    pub fn is_success(&self) -> bool {
        self.success
    }

    pub fn is_error(&self) -> bool {
        !self.success
    }
}

impl<T: Default> Default for ApiResponse<T> {
    fn default() -> Self {
        Self::success(T::default())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    pub id: Uuid,
    pub user_id: Uuid,
    pub bot_id: Uuid,
    pub title: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<DateTime<Utc>>,
}

impl Session {
    pub fn new(user_id: Uuid, bot_id: Uuid, title: impl Into<String>) -> Self {
        let now = Utc::now();
        Self {
            id: Uuid::new_v4(),
            user_id,
            bot_id,
            title: title.into(),
            created_at: now,
            updated_at: now,
            expires_at: None,
        }
    }

    pub fn with_expiry(mut self, expires_at: DateTime<Utc>) -> Self {
        self.expires_at = Some(expires_at);
        self
    }

    pub fn is_expired(&self) -> bool {
        self.expires_at.map(|exp| Utc::now() > exp).unwrap_or(false)
    }

    pub fn is_active(&self) -> bool {
        !self.is_expired()
    }

    pub fn remaining_time(&self) -> Option<chrono::Duration> {
        self.expires_at.map(|exp| exp - Utc::now())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserMessage {
    pub bot_id: String,
    pub user_id: String,
    pub session_id: String,
    pub channel: String,
    pub content: String,
    pub message_type: MessageType,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub media_url: Option<String>,
    pub timestamp: DateTime<Utc>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub context_name: Option<String>,
}

impl UserMessage {
    pub fn text(
        bot_id: impl Into<String>,
        user_id: impl Into<String>,
        session_id: impl Into<String>,
        channel: impl Into<String>,
        content: impl Into<String>,
    ) -> Self {
        Self {
            bot_id: bot_id.into(),
            user_id: user_id.into(),
            session_id: session_id.into(),
            channel: channel.into(),
            content: content.into(),
            message_type: MessageType::USER,
            media_url: None,
            timestamp: Utc::now(),
            context_name: None,
        }
    }

    pub fn with_media(mut self, url: impl Into<String>) -> Self {
        self.media_url = Some(url.into());
        self
    }

    pub fn with_context(mut self, context: impl Into<String>) -> Self {
        self.context_name = Some(context.into());
        self
    }

    pub fn has_media(&self) -> bool {
        self.media_url.is_some()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Suggestion {
    pub text: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub context: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub action: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub icon: Option<String>,
}

impl Suggestion {
    pub fn new(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            context: None,
            action: None,
            icon: None,
        }
    }

    pub fn with_context(mut self, context: impl Into<String>) -> Self {
        self.context = Some(context.into());
        self
    }

    pub fn with_action(mut self, action: impl Into<String>) -> Self {
        self.action = Some(action.into());
        self
    }

    pub fn with_icon(mut self, icon: impl Into<String>) -> Self {
        self.icon = Some(icon.into());
        self
    }
}

impl<S: Into<String>> From<S> for Suggestion {
    fn from(text: S) -> Self {
        Self::new(text)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BotResponse {
    pub bot_id: String,
    pub user_id: String,
    pub session_id: String,
    pub channel: String,
    pub content: String,
    pub message_type: MessageType,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream_token: Option<String>,
    pub is_complete: bool,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub suggestions: Vec<Suggestion>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub context_name: Option<String>,
    #[serde(default)]
    pub context_length: usize,
    #[serde(default)]
    pub context_max_length: usize,
}

impl BotResponse {
    pub fn new(
        bot_id: impl Into<String>,
        session_id: impl Into<String>,
        user_id: impl Into<String>,
        content: impl Into<String>,
        channel: impl Into<String>,
    ) -> Self {
        Self {
            bot_id: bot_id.into(),
            user_id: user_id.into(),
            session_id: session_id.into(),
            channel: channel.into(),
            content: content.into(),
            message_type: MessageType::BOT_RESPONSE,
            stream_token: None,
            is_complete: true,
            suggestions: Vec::new(),
            context_name: None,
            context_length: 0,
            context_max_length: 0,
        }
    }

    pub fn streaming(
        bot_id: impl Into<String>,
        session_id: impl Into<String>,
        user_id: impl Into<String>,
        channel: impl Into<String>,
        stream_token: impl Into<String>,
    ) -> Self {
        Self {
            bot_id: bot_id.into(),
            user_id: user_id.into(),
            session_id: session_id.into(),
            channel: channel.into(),
            content: String::new(),
            message_type: MessageType::BOT_RESPONSE,
            stream_token: Some(stream_token.into()),
            is_complete: false,
            suggestions: Vec::new(),
            context_name: None,
            context_length: 0,
            context_max_length: 0,
        }
    }

    pub fn with_suggestions<I, S>(mut self, suggestions: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<Suggestion>,
    {
        self.suggestions = suggestions.into_iter().map(Into::into).collect();
        self
    }

    pub fn add_suggestion(mut self, suggestion: impl Into<Suggestion>) -> Self {
        self.suggestions.push(suggestion.into());
        self
    }

    pub fn with_context(
        mut self,
        name: impl Into<String>,
        length: usize,
        max_length: usize,
    ) -> Self {
        self.context_name = Some(name.into());
        self.context_length = length;
        self.context_max_length = max_length;
        self
    }

    pub fn append_content(&mut self, chunk: &str) {
        self.content.push_str(chunk);
    }

    pub fn complete(mut self) -> Self {
        self.is_complete = true;
        self
    }

    pub fn is_streaming(&self) -> bool {
        self.stream_token.is_some() && !self.is_complete
    }

    pub fn has_suggestions(&self) -> bool {
        !self.suggestions.is_empty()
    }
}

impl Default for BotResponse {
    fn default() -> Self {
        Self {
            bot_id: String::new(),
            user_id: String::new(),
            session_id: String::new(),
            channel: String::new(),
            content: String::new(),
            message_type: MessageType::BOT_RESPONSE,
            stream_token: None,
            is_complete: true,
            suggestions: Vec::new(),
            context_name: None,
            context_length: 0,
            context_max_length: 0,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Attachment {
    pub attachment_type: AttachmentType,
    pub url: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mime_type: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub filename: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub size: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thumbnail_url: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AttachmentType {
    Image,
    Audio,
    Video,
    Document,
    File,
}

impl Attachment {
    pub fn new(attachment_type: AttachmentType, url: impl Into<String>) -> Self {
        Self {
            attachment_type,
            url: url.into(),
            mime_type: None,
            filename: None,
            size: None,
            thumbnail_url: None,
        }
    }

    pub fn image(url: impl Into<String>) -> Self {
        Self::new(AttachmentType::Image, url)
    }

    pub fn audio(url: impl Into<String>) -> Self {
        Self::new(AttachmentType::Audio, url)
    }

    pub fn video(url: impl Into<String>) -> Self {
        Self::new(AttachmentType::Video, url)
    }

    pub fn document(url: impl Into<String>) -> Self {
        Self::new(AttachmentType::Document, url)
    }

    pub fn file(url: impl Into<String>) -> Self {
        Self::new(AttachmentType::File, url)
    }

    pub fn with_mime_type(mut self, mime_type: impl Into<String>) -> Self {
        self.mime_type = Some(mime_type.into());
        self
    }

    pub fn with_filename(mut self, filename: impl Into<String>) -> Self {
        self.filename = Some(filename.into());
        self
    }

    pub fn with_size(mut self, size: u64) -> Self {
        self.size = Some(size);
        self
    }

    pub fn with_thumbnail(mut self, thumbnail_url: impl Into<String>) -> Self {
        self.thumbnail_url = Some(thumbnail_url.into());
        self
    }

    pub fn is_image(&self) -> bool {
        self.attachment_type == AttachmentType::Image
    }

    pub fn is_media(&self) -> bool {
        matches!(
            self.attachment_type,
            AttachmentType::Image | AttachmentType::Audio | AttachmentType::Video
        )
    }
}

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

    #[test]
    fn test_api_response_success() {
        let response: ApiResponse<String> = ApiResponse::success("test".to_string());
        assert!(response.is_success());
        assert!(!response.is_error());
        assert_eq!(response.data, Some("test".to_string()));
        assert!(response.error.is_none());
    }

    #[test]
    fn test_api_response_error() {
        let response: ApiResponse<String> = ApiResponse::error("something went wrong");
        assert!(!response.is_success());
        assert!(response.is_error());
        assert!(response.data.is_none());
        assert_eq!(response.error, Some("something went wrong".to_string()));
    }

    #[test]
    fn test_api_response_map() {
        let response: ApiResponse<i32> = ApiResponse::success(42);
        let mapped = response.map(|n| n.to_string());
        assert_eq!(mapped.data, Some("42".to_string()));
    }

    #[test]
    fn test_session_creation() {
        let user_id = Uuid::new_v4();
        let bot_id = Uuid::new_v4();
        let session = Session::new(user_id, bot_id, "Test Session");

        assert_eq!(session.user_id, user_id);
        assert_eq!(session.bot_id, bot_id);
        assert_eq!(session.title, "Test Session");
        assert!(session.is_active());
        assert!(!session.is_expired());
    }

    #[test]
    fn test_user_message_creation() {
        let msg =
            UserMessage::text("bot1", "user1", "sess1", "web", "Hello!").with_context("greeting");

        assert_eq!(msg.content, "Hello!");
        assert_eq!(msg.message_type, MessageType::USER);
        assert_eq!(msg.context_name, Some("greeting".to_string()));
    }

    #[test]
    fn test_bot_response_creation() {
        let response = BotResponse::new("bot1", "sess1", "user1", "Hi there!", "web")
            .add_suggestion("Option 1")
            .add_suggestion("Option 2");

        assert!(response.is_complete);
        assert!(!response.is_streaming());
        assert!(response.has_suggestions());
        assert_eq!(response.suggestions.len(), 2);
    }

    #[test]
    fn test_bot_response_streaming() {
        let mut response = BotResponse::streaming("bot1", "sess1", "user1", "web", "token123");
        assert!(response.is_streaming());
        assert!(!response.is_complete);

        response.append_content("Hello ");
        response.append_content("World!");
        assert_eq!(response.content, "Hello World!");

        let response = response.complete();
        assert!(!response.is_streaming());
        assert!(response.is_complete);
    }

    #[test]
    fn test_attachment_creation() {
        let attachment = Attachment::image("https://example.com/photo.jpg")
            .with_filename("photo.jpg")
            .with_size(1024)
            .with_mime_type("image/jpeg");

        assert!(attachment.is_image());
        assert!(attachment.is_media());
        assert_eq!(attachment.filename, Some("photo.jpg".to_string()));
        assert_eq!(attachment.size, Some(1024));
    }

    #[test]
    fn test_suggestion_from_string() {
        let suggestion: Suggestion = "Click here".into();
        assert_eq!(suggestion.text, "Click here");
        assert!(suggestion.context.is_none());
    }
}