certifai/src/models/chat.rs

use serde::{Deserialize, Serialize};

/// The role of a participant in a chat conversation.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ChatRole {
    /// Message sent by the human user
    User,
    /// Message generated by the AI assistant
    Assistant,
    /// System-level instruction (not displayed in UI)
    System,
}

/// Namespace for grouping chat sessions in the sidebar.
///
/// Sessions are visually separated in the chat sidebar by namespace,
/// with `News` sessions appearing under a dedicated "News Chats" header.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub enum ChatNamespace {
    /// General user-initiated chat conversations.
    #[default]
    General,
    /// Chats originating from news article follow-ups on the dashboard.
    News,
}

/// The type of file attached to a chat message.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum AttachmentKind {
    /// Image file (png, jpg, webp, etc.)
    Image,
    /// Document file (pdf, docx, txt, etc.)
    Document,
    /// Source code file
    Code,
}

/// A file attachment on a chat message.
///
/// # Fields
///
/// * `name` - Original filename
/// * `kind` - Type of attachment for rendering
/// * `size_bytes` - File size in bytes
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Attachment {
    pub name: String,
    pub kind: AttachmentKind,
    pub size_bytes: u64,
}

/// A persisted chat session stored in MongoDB.
///
/// Messages are stored separately in the `chat_messages` collection
/// and loaded on demand when the user opens a session.
///
/// # Fields
///
/// * `id` - MongoDB document ID (hex string)
/// * `user_sub` - Keycloak subject ID (session owner)
/// * `title` - Display title (auto-generated or user-renamed)
/// * `namespace` - Grouping for sidebar sections
/// * `provider` - LLM provider used (e.g. "litellm", "openai")
/// * `model` - Model ID used (e.g. "qwen3-32b")
/// * `created_at` - ISO 8601 creation timestamp
/// * `updated_at` - ISO 8601 last-activity timestamp
/// * `article_url` - Source article URL (for News namespace sessions)
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatSession {
    #[serde(default, alias = "_id", skip_serializing_if = "String::is_empty")]
    pub id: String,
    pub user_sub: String,
    pub title: String,
    #[serde(default)]
    pub namespace: ChatNamespace,
    pub provider: String,
    pub model: String,
    pub created_at: String,
    pub updated_at: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub article_url: Option<String>,
}

/// A single persisted message within a chat session.
///
/// Stored in the `chat_messages` MongoDB collection, linked to a
/// `ChatSession` via `session_id`.
///
/// # Fields
///
/// * `id` - MongoDB document ID (hex string)
/// * `session_id` - Foreign key to `ChatSession.id`
/// * `role` - Who sent this message
/// * `content` - Message text content (may contain markdown)
/// * `attachments` - File attachments (Phase 2, currently empty)
/// * `timestamp` - ISO 8601 timestamp
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatMessage {
    #[serde(default, alias = "_id", skip_serializing_if = "String::is_empty")]
    pub id: String,
    pub session_id: String,
    pub role: ChatRole,
    pub content: String,
    #[serde(default)]
    pub attachments: Vec<Attachment>,
    pub timestamp: String,
}

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

    #[test]
    fn chat_namespace_default_is_general() {
        assert_eq!(ChatNamespace::default(), ChatNamespace::General);
    }

    #[test]
    fn chat_role_serde_round_trip() {
        for role in [ChatRole::User, ChatRole::Assistant, ChatRole::System] {
            let json =
                serde_json::to_string(&role).unwrap_or_else(|_| panic!("serialize {:?}", role));
            let back: ChatRole =
                serde_json::from_str(&json).unwrap_or_else(|_| panic!("deserialize {:?}", role));
            assert_eq!(role, back);
        }
    }

    #[test]
    fn chat_namespace_serde_round_trip() {
        for ns in [ChatNamespace::General, ChatNamespace::News] {
            let json = serde_json::to_string(&ns).unwrap_or_else(|_| panic!("serialize {:?}", ns));
            let back: ChatNamespace =
                serde_json::from_str(&json).unwrap_or_else(|_| panic!("deserialize {:?}", ns));
            assert_eq!(ns, back);
        }
    }

    #[test]
    fn attachment_kind_serde_round_trip() {
        for kind in [
            AttachmentKind::Image,
            AttachmentKind::Document,
            AttachmentKind::Code,
        ] {
            let json =
                serde_json::to_string(&kind).unwrap_or_else(|_| panic!("serialize {:?}", kind));
            let back: AttachmentKind =
                serde_json::from_str(&json).unwrap_or_else(|_| panic!("deserialize {:?}", kind));
            assert_eq!(kind, back);
        }
    }

    #[test]
    fn attachment_serde_round_trip() {
        let att = Attachment {
            name: "photo.png".into(),
            kind: AttachmentKind::Image,
            size_bytes: 2048,
        };
        let json = serde_json::to_string(&att).expect("serialize Attachment");
        let back: Attachment = serde_json::from_str(&json).expect("deserialize Attachment");
        assert_eq!(att, back);
    }

    #[test]
    fn chat_session_serde_round_trip() {
        let session = ChatSession {
            id: "abc123".into(),
            user_sub: "user-1".into(),
            title: "Test Chat".into(),
            namespace: ChatNamespace::General,
            provider: "litellm".into(),
            model: "qwen3-32b".into(),
            created_at: "2025-01-01T00:00:00Z".into(),
            updated_at: "2025-01-01T01:00:00Z".into(),
            article_url: None,
        };
        let json = serde_json::to_string(&session).expect("serialize ChatSession");
        let back: ChatSession = serde_json::from_str(&json).expect("deserialize ChatSession");
        assert_eq!(session, back);
    }

    #[test]
    fn chat_session_id_alias_deserialization() {
        // MongoDB returns `_id` instead of `id`
        let json = r#"{
            "_id": "mongo-id",
            "user_sub": "u1",
            "title": "t",
            "provider": "litellm",
            "model": "m",
            "created_at": "2025-01-01",
            "updated_at": "2025-01-01"
        }"#;
        let session: ChatSession = serde_json::from_str(json).expect("deserialize with _id");
        assert_eq!(session.id, "mongo-id");
    }

    #[test]
    fn chat_session_empty_id_skips_serialization() {
        let session = ChatSession {
            id: String::new(),
            user_sub: "u1".into(),
            title: "t".into(),
            namespace: ChatNamespace::default(),
            provider: "litellm".into(),
            model: "m".into(),
            created_at: "2025-01-01".into(),
            updated_at: "2025-01-01".into(),
            article_url: None,
        };
        let json = serde_json::to_string(&session).expect("serialize");
        // `id` field should be absent when empty due to skip_serializing_if
        assert!(!json.contains("\"id\""));
    }

    #[test]
    fn chat_session_none_article_url_skips_serialization() {
        let session = ChatSession {
            id: "s1".into(),
            user_sub: "u1".into(),
            title: "t".into(),
            namespace: ChatNamespace::default(),
            provider: "litellm".into(),
            model: "m".into(),
            created_at: "2025-01-01".into(),
            updated_at: "2025-01-01".into(),
            article_url: None,
        };
        let json = serde_json::to_string(&session).expect("serialize");
        assert!(!json.contains("article_url"));
    }

    #[test]
    fn chat_message_serde_round_trip() {
        let msg = ChatMessage {
            id: "msg-1".into(),
            session_id: "s1".into(),
            role: ChatRole::User,
            content: "Hello AI".into(),
            attachments: vec![Attachment {
                name: "doc.pdf".into(),
                kind: AttachmentKind::Document,
                size_bytes: 4096,
            }],
            timestamp: "2025-01-01T00:00:00Z".into(),
        };
        let json = serde_json::to_string(&msg).expect("serialize ChatMessage");
        let back: ChatMessage = serde_json::from_str(&json).expect("deserialize ChatMessage");
        assert_eq!(msg, back);
    }

    #[test]
    fn chat_message_id_alias_deserialization() {
        let json = r#"{
            "_id": "mongo-msg-id",
            "session_id": "s1",
            "role": "User",
            "content": "hi",
            "timestamp": "2025-01-01"
        }"#;
        let msg: ChatMessage = serde_json::from_str(json).expect("deserialize with _id");
        assert_eq!(msg.id, "mongo-msg-id");
    }
}