mail2couch/rust/src/imap.rs

//! IMAP client functionality for mail2couch
//! 
//! This module provides IMAP client operations for connecting to mail servers,
//! listing mailboxes, and retrieving messages.

use crate::config::{MailSource, MessageFilter};
use crate::schemas::MailDocument;
use anyhow::{anyhow, Result};
use async_imap::types::Fetch;
use async_imap::{Client, Session};
use async_native_tls::{TlsConnector, TlsStream};
use async_std::io::{Read, Write};
use async_std::net::TcpStream;
use async_std::stream::StreamExt;
use async_std::task::{Context, Poll};
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::pin::Pin;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum ImapError {
    #[error("Connection failed: {0}")]
    Connection(String),
    #[error("Authentication failed: {0}")]
    Authentication(String),
    #[error("IMAP operation failed: {0}")]
    Operation(String),
    #[error("Message parsing failed: {0}")]
    Parsing(String),
}

/// Wrapper for both TLS and plain TCP streams
pub enum ImapStream {
    Plain(TcpStream),
    Tls(TlsStream<TcpStream>),
}

impl Read for ImapStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<std::io::Result<usize>> {
        match self.get_mut() {
            ImapStream::Plain(stream) => Pin::new(stream).poll_read(cx, buf),
            ImapStream::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
        }
    }
}

impl Write for ImapStream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        match self.get_mut() {
            ImapStream::Plain(stream) => Pin::new(stream).poll_write(cx, buf),
            ImapStream::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            ImapStream::Plain(stream) => Pin::new(stream).poll_flush(cx),
            ImapStream::Tls(stream) => Pin::new(stream).poll_flush(cx),
        }
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            ImapStream::Plain(stream) => Pin::new(stream).poll_close(cx),
            ImapStream::Tls(stream) => Pin::new(stream).poll_close(cx),
        }
    }
}

impl std::fmt::Debug for ImapStream {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ImapStream::Plain(_) => write!(f, "ImapStream::Plain(_)"),
            ImapStream::Tls(_) => write!(f, "ImapStream::Tls(_)"),
        }
    }
}

pub type ImapSession = Session<ImapStream>;

/// IMAP client for mail operations
pub struct ImapClient {
    session: Option<ImapSession>,
    source: MailSource,
}

/// Represents a mailbox on the IMAP server
#[derive(Debug, Clone)]
pub struct MailboxInfo {
    pub name: String,
    pub exists: u32,
    pub recent: u32,
    pub uid_validity: Option<u32>,
    pub uid_next: Option<u32>,
}

impl ImapClient {
    /// Create a new IMAP client and connect to the server
    pub async fn connect(source: MailSource) -> Result<Self> {
        let mut client = ImapClient {
            session: None,
            source,
        };
        
        client.establish_connection().await?;
        Ok(client)
    }

    /// Establish connection to IMAP server
    async fn establish_connection(&mut self) -> Result<()> {
        // Connect to the IMAP server
        let addr = format!("{}:{}", self.source.host, self.source.port);
        let tcp_stream = TcpStream::connect(&addr).await
            .map_err(|e| ImapError::Connection(format!("Failed to connect to {}: {}", addr, e)))?;

        // Determine if we should use TLS based on port
        let imap_stream = if self.should_use_tls() {
            // Use TLS for secure connection (typically port 993)
            let tls_connector = TlsConnector::new();
            let tls_stream = tls_connector.connect(&self.source.host, tcp_stream).await
                .map_err(|e| ImapError::Connection(format!("TLS connection failed: {}", e)))?;
            ImapStream::Tls(tls_stream)
        } else {
            // Use plain connection (typically port 143 or test environments)
            ImapStream::Plain(tcp_stream)
        };

        // Log connection type for debugging
        let connection_type = if self.should_use_tls() { "TLS" } else { "Plain" };
        log::debug!(
            "Connecting to {}:{} using {} connection",
            self.source.host,
            self.source.port,
            connection_type
        );

        // Create IMAP client
        let client = Client::new(imap_stream);

        // Perform IMAP login
        let session = client
            .login(&self.source.user, &self.source.password)
            .await
            .map_err(|e| ImapError::Authentication(format!("Login failed: {:?}", e)))?;

        self.session = Some(session);
        Ok(())
    }

    /// Determine if TLS should be used based on port and configuration
    fn should_use_tls(&self) -> bool {
        // Standard IMAPS port (993) uses TLS
        // Port 143 is typically plain IMAP
        // Port 3143 is used in test environments (plain)
        match self.source.port {
            993 => true,   // Standard IMAPS port
            143 => false,  // Standard IMAP port
            3143 => false, // Test environment port
            _ => {
                // For other ports, default to TLS for security
                // but log a warning
                log::warn!(
                    "Unknown IMAP port {}, defaulting to TLS. Use port 143 for plain IMAP or 993 for TLS",
                    self.source.port
                );
                true
            }
        }
    }

    /// List all mailboxes
    pub async fn list_mailboxes(&mut self) -> Result<Vec<String>> {
        let session = self.session.as_mut()
            .ok_or_else(|| anyhow!("Not connected to IMAP server"))?;

        let mut mailboxes = session.list(Some(""), Some("*")).await
            .map_err(|e| ImapError::Operation(format!("Failed to list mailboxes: {:?}", e)))?;

        let mut mailbox_names = Vec::new();
        while let Some(mailbox_result) = mailboxes.next().await {
            match mailbox_result {
                Ok(mailbox) => mailbox_names.push(mailbox.name().to_string()),
                Err(e) => return Err(ImapError::Operation(format!("Error processing mailbox: {:?}", e)).into()),
            }
        }

        Ok(mailbox_names)
    }

    /// Select a mailbox
    pub async fn select_mailbox(&mut self, mailbox: &str) -> Result<MailboxInfo> {
        let session = self.session.as_mut()
            .ok_or_else(|| anyhow!("Not connected to IMAP server"))?;

        let mailbox_data = session.select(mailbox).await
            .map_err(|e| ImapError::Operation(format!("Failed to select mailbox {}: {:?}", mailbox, e)))?;

        Ok(MailboxInfo {
            name: mailbox.to_string(),
            exists: mailbox_data.exists,
            recent: mailbox_data.recent,
            uid_validity: mailbox_data.uid_validity,
            uid_next: mailbox_data.uid_next,
        })
    }

    /// Search for messages using IMAP SEARCH command
    /// Returns UIDs of matching messages
    pub async fn search_messages(&mut self, since_date: Option<&DateTime<Utc>>) -> Result<Vec<u32>> {
        let session = self.session.as_mut()
            .ok_or_else(|| anyhow!("Not connected to IMAP server"))?;

        let search_query = if let Some(since) = since_date {
            // Format date as required by IMAP (DD-MMM-YYYY)
            let formatted_date = since.format("%d-%b-%Y").to_string();
            format!("SINCE {}", formatted_date)
        } else {
            "ALL".to_string()
        };

        let uids = session.uid_search(&search_query).await
            .map_err(|e| ImapError::Operation(format!("Search failed: {:?}", e)))?;

        Ok(uids.into_iter().collect())
    }

    /// Fetch message by UID
    pub async fn fetch_message(&mut self, uid: u32) -> Result<MailDocument> {
        let session = self.session.as_mut()
            .ok_or_else(|| anyhow!("Not connected to IMAP server"))?;

        // Fetch message headers and body
        let mut messages = session.uid_fetch(format!("{}", uid), "RFC822").await
            .map_err(|e| ImapError::Operation(format!("Failed to fetch message {}: {:?}", uid, e)))?;

        // Collect the first message
        if let Some(message_result) = messages.next().await {
            match message_result {
                Ok(message) => {
                    // Drop the messages stream to release the session borrow
                    drop(messages);
                    self.parse_message(&message, uid).await
                }
                Err(e) => Err(ImapError::Operation(format!("Failed to process message {}: {:?}", uid, e)).into()),
            }
        } else {
            Err(anyhow!("Message {} not found", uid))
        }
    }

    /// Fetch multiple messages by UIDs
    pub async fn fetch_messages(&mut self, uids: &[u32], max_count: Option<u32>) -> Result<Vec<MailDocument>> {
        if uids.is_empty() {
            return Ok(Vec::new());
        }

        let session = self.session.as_mut()
            .ok_or_else(|| anyhow!("Not connected to IMAP server"))?;

        // Limit the number of messages if specified
        let uids_to_fetch = if let Some(max) = max_count {
            if uids.len() > max as usize {
                &uids[..max as usize]
            } else {
                uids
            }
        } else {
            uids
        };

        // Create UID sequence
        let uid_sequence = uids_to_fetch.iter()
            .map(|uid| uid.to_string())
            .collect::<Vec<_>>()
            .join(",");

        // Fetch messages
        let mut messages = session.uid_fetch(&uid_sequence, "RFC822").await
            .map_err(|e| ImapError::Operation(format!("Failed to fetch messages: {:?}", e)))?;

        // Collect all messages first to avoid borrowing issues
        let mut fetched_messages = Vec::new();
        while let Some(message_result) = messages.next().await {
            match message_result {
                Ok(message) => fetched_messages.push(message),
                Err(e) => log::warn!("Failed to fetch message: {:?}", e),
            }
        }
        
        // Drop the messages stream to release the session borrow
        drop(messages);

        let mut mail_documents = Vec::new();
        for (i, message) in fetched_messages.iter().enumerate() {
            if let Some(&uid) = uids_to_fetch.get(i) {
                match self.parse_message(message, uid).await {
                    Ok(doc) => mail_documents.push(doc),
                    Err(e) => {
                        log::warn!("Failed to parse message {}: {}", uid, e);
                    }
                }
            }
        }

        Ok(mail_documents)
    }

    /// Parse a raw IMAP message into a MailDocument
    async fn parse_message(&self, message: &Fetch, uid: u32) -> Result<MailDocument> {
        let body = message.body()
            .ok_or_else(|| ImapError::Parsing("No message body found".to_string()))?;

        // Parse the email using a simple RFC822 parser
        // This is a basic implementation - a production version would use a proper email parser
        let email_str = String::from_utf8_lossy(body);
        let (headers, body_content) = self.parse_rfc822(&email_str)?;

        // Extract key fields
        let from = self.parse_addresses(&headers, "from")?;
        let to = self.parse_addresses(&headers, "to")?;
        let subject = headers.get("subject")
            .and_then(|v| v.first())
            .unwrap_or(&"No Subject".to_string())
            .clone();

        // Parse date
        let date = self.parse_date(&headers)?;

        // Get current mailbox name (this would need to be passed in properly)
        let mailbox = "INBOX".to_string(); // Placeholder - should be passed from caller

        let mail_doc = MailDocument::new(
            uid.to_string(),
            mailbox,
            from,
            to,
            subject,
            date,
            body_content,
            headers,
            false, // TODO: Check for attachments
        );

        Ok(mail_doc)
    }

    /// Basic RFC822 header and body parser
    fn parse_rfc822(&self, email: &str) -> Result<(HashMap<String, Vec<String>>, String)> {
        let mut headers = HashMap::new();
        let lines = email.lines();
        let mut body_lines = Vec::new();
        let mut in_body = false;

        for line in lines {
            if in_body {
                body_lines.push(line);
            } else if line.trim().is_empty() {
                in_body = true;
            } else if line.starts_with(' ') || line.starts_with('\t') {
                // Continuation of previous header
                // Skip for simplicity in this basic implementation
                continue;
            } else if let Some(colon_pos) = line.find(':') {
                let header_name = line[..colon_pos].trim().to_lowercase();
                let header_value = line[colon_pos + 1..].trim().to_string();
                
                headers.entry(header_name)
                    .or_insert_with(Vec::new)
                    .push(header_value);
            }
        }

        let body = body_lines.join("\n");
        Ok((headers, body))
    }

    /// Parse email addresses from headers
    fn parse_addresses(&self, headers: &HashMap<String, Vec<String>>, header_name: &str) -> Result<Vec<String>> {
        let addresses = headers.get(header_name)
            .map(|values| values.clone())
            .unwrap_or_default();
        
        // Basic email extraction - just return the raw values for now
        // A production implementation would properly parse RFC822 addresses
        Ok(addresses)
    }

    /// Parse date from headers
    fn parse_date(&self, headers: &HashMap<String, Vec<String>>) -> Result<DateTime<Utc>> {
        let default_date = Utc::now().to_rfc2822();
        let date_str = headers.get("date")
            .and_then(|v| v.first())
            .unwrap_or(&default_date);

        // Try to parse RFC2822 date format
        // For simplicity, fall back to current time if parsing fails
        DateTime::parse_from_rfc2822(date_str)
            .map(|dt| dt.with_timezone(&Utc))
            .or_else(|_| {
                log::warn!("Failed to parse date '{}', using current time", date_str);
                Ok(Utc::now())
            })
    }

    /// Close the IMAP connection
    pub async fn close(self) -> Result<()> {
        if let Some(mut session) = self.session {
            session.logout().await
                .map_err(|e| ImapError::Operation(format!("Logout failed: {:?}", e)))?;
        }
        Ok(())
    }
}

/// Apply message filters to determine if a message should be processed
pub fn should_process_message(
    mail_doc: &MailDocument,
    filter: &MessageFilter,
) -> bool {
    // Check subject keywords
    if !filter.subject_keywords.is_empty() {
        let subject_lower = mail_doc.subject.to_lowercase();
        let has_subject_keyword = filter.subject_keywords.iter()
            .any(|keyword| subject_lower.contains(&keyword.to_lowercase()));
        if !has_subject_keyword {
            return false;
        }
    }

    // Check sender keywords
    if !filter.sender_keywords.is_empty() {
        let has_sender_keyword = mail_doc.from.iter()
            .any(|from_addr| {
                let from_lower = from_addr.to_lowercase();
                filter.sender_keywords.iter()
                    .any(|keyword| from_lower.contains(&keyword.to_lowercase()))
            });
        if !has_sender_keyword {
            return false;
        }
    }

    // Check recipient keywords
    if !filter.recipient_keywords.is_empty() {
        let has_recipient_keyword = mail_doc.to.iter()
            .any(|to_addr| {
                let to_lower = to_addr.to_lowercase();
                filter.recipient_keywords.iter()
                    .any(|keyword| to_lower.contains(&keyword.to_lowercase()))
            });
        if !has_recipient_keyword {
            return false;
        }
    }

    true
}

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

    #[test]
    fn test_message_filtering() {
        let mail_doc = MailDocument::new(
            "123".to_string(),
            "INBOX".to_string(),
            vec!["sender@example.com".to_string()],
            vec!["recipient@test.com".to_string()],
            "Urgent: Meeting tomorrow".to_string(),
            Utc::now(),
            "Test body".to_string(),
            HashMap::new(),
            false,
        );

        // Test subject keyword filtering
        let mut filter = MessageFilter {
            subject_keywords: vec!["urgent".to_string()],
            ..Default::default()
        };
        assert!(should_process_message(&mail_doc, &filter));

        filter.subject_keywords = vec!["spam".to_string()];
        assert!(!should_process_message(&mail_doc, &filter));

        // Test sender keyword filtering
        filter = MessageFilter {
            sender_keywords: vec!["@example.com".to_string()],
            ..Default::default()
        };
        assert!(should_process_message(&mail_doc, &filter));

        filter.sender_keywords = vec!["@spam.com".to_string()];
        assert!(!should_process_message(&mail_doc, &filter));
    }

    #[test]
    fn test_rfc822_parsing() {
        let client = ImapClient {
            session: None,
            source: MailSource {
                name: "test".to_string(),
                enabled: true,
                protocol: "imap".to_string(),
                host: "localhost".to_string(),
                port: 143,
                user: "test".to_string(),
                password: "test".to_string(),
                mode: "archive".to_string(),
                folder_filter: Default::default(),
                message_filter: Default::default(),
            },
        };

        let email = "From: sender@example.com\r\nTo: recipient@example.com\r\nSubject: Test\r\n\r\nTest body\r\n";
        let (headers, body) = client.parse_rfc822(email).unwrap();

        assert_eq!(headers.get("from").unwrap()[0], "sender@example.com");
        assert_eq!(headers.get("to").unwrap()[0], "recipient@example.com");
        assert_eq!(headers.get("subject").unwrap()[0], "Test");
        assert_eq!(body.trim(), "Test body");
    }
}