core/

mailbox.rs

//! Mailbox is an actor to handle message routing and persistence.

// TODO: Split into many files, cleanup dead code, and maybe finish implementing?

use chrono::Utc;
use sqlx::{PgPool, postgres::PgPoolOptions};
use std::cmp::Ordering;
use std::collections::BinaryHeap;
use std::env;
use std::fs::File;
use std::io::{BufRead, BufReader, Seek, SeekFrom};
use std::path::PathBuf;
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tracing::{debug, error, info};
use types::{
    cv::Gps,
    mailbox::{GNCTelemetry, MailboxMessage, MailboxMessageType},
};

// Wrapper struct for items in the priority queue
#[derive(Debug, Clone, Eq, PartialEq)]
struct QueueItem {
    priority: u8,
    message: MailboxMessage,
}

impl Ord for QueueItem {
    fn cmp(&self, other: &Self) -> Ordering {
        other.priority.cmp(&self.priority)
    }
}

impl PartialOrd for QueueItem {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

// Messages that can be sent to the Mailbox actor
#[allow(dead_code)]
enum MailboxActorMessage {
    QueueMessage {
        message: MailboxMessage,
    },
    ProcessQueue,
    SaveGpsToDb {
        gps: Gps,
    },
    GetNextMessage {
        respond_to: tokio::sync::oneshot::Sender<Option<MailboxMessage>>,
    },
}

// The Mailbox actor
struct Mailbox {
    receiver: mpsc::Receiver<MailboxActorMessage>,
    queue: BinaryHeap<QueueItem>,
    pool: PgPool,
    telemetry_sender: mpsc::Sender<String>,
    log_sender: mpsc::Sender<String>,
}

impl Mailbox {
    async fn new(receiver: mpsc::Receiver<MailboxActorMessage>) -> Self {
        info!("Mailbox Actor Started");

        let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
        debug!("Mailbox connecting to database: {}", database_url);

        let pool = PgPoolOptions::new()
            .max_connections(5)
            .connect(&database_url)
            .await
            .expect("Failed to create pool for Mailbox");

        if let Err(e) = Self::create_mailbox_schema(&pool).await {
            error!("Failed to create mailbox schema: {}", e);
        }

        // Create a channel for general-purpose telemetry broadcasting
        let (telemetry_sender, telemetry_receiver) = mpsc::channel::<String>(100);

        // Create a channel for log broadcasting
        let (log_sender, log_receiver) = mpsc::channel::<String>(100);

        // Spawn telemetry broadcaster in a separate task
        tokio::task::spawn_blocking(move || {
            Self::run_telemetry_broadcaster(telemetry_receiver);
        });

        // Spawn log broadcaster in a separate task
        tokio::task::spawn_blocking(move || {
            Self::run_log_broadcaster(log_receiver);
        });

        Mailbox {
            receiver,
            queue: BinaryHeap::new(),
            pool,
            telemetry_sender,
            log_sender,
        }
    }

    fn run_telemetry_broadcaster(mut receiver: mpsc::Receiver<String>) {
        let context = zmq::Context::new();
        let publisher = context
            .socket(zmq::PUB)
            .expect("Failed to create ZMQ PUB socket for telemetry");

        let telemetry_port = env::var("TELEMETRY_PORT").unwrap_or("5557".to_string());
        let telemetry_addr = format!("tcp://0.0.0.0:{}", telemetry_port);

        publisher
            .bind(&telemetry_addr)
            .expect("Failed to bind telemetry publisher");
        info!(
            "General-Purpose Telemetry port opened on {}",
            telemetry_addr
        );

        // Block and wait for messages to broadcast
        while let Some(msg) = receiver.blocking_recv() {
            if let Err(e) = publisher.send(&msg, 0) {
                error!("Failed to broadcast telemetry: {}", e);
            }
        }
    }

    fn run_log_broadcaster(mut receiver: mpsc::Receiver<String>) {
        let context = zmq::Context::new();
        let publisher = context
            .socket(zmq::PUB)
            .expect("Failed to create ZMQ PUB socket for logs");

        let logs_port = env::var("LOGS_PORT").unwrap_or("5558".to_string());
        let logs_addr = format!("tcp://0.0.0.0:{}", logs_port);

        publisher
            .bind(&logs_addr)
            .expect("Failed to bind log publisher");
        info!("Log broadcasting port opened on {}", logs_addr);

        // Block and wait for messages to broadcast
        while let Some(msg) = receiver.blocking_recv() {
            if let Err(e) = publisher.send(&msg, 0) {
                error!("Failed to broadcast log: {}", e);
            }
        }
    }

    async fn create_mailbox_schema(pool: &PgPool) -> Result<(), sqlx::Error> {
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS mailbox_messages (
                id SERIAL PRIMARY KEY,
                timestamp BIGINT NOT NULL,
                msg_type TEXT NOT NULL,
                data JSONB NOT NULL,
                processed BOOLEAN DEFAULT FALSE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
            "#,
        )
        .execute(pool)
        .await?;

        sqlx::query(
            r#"
            CREATE INDEX IF NOT EXISTS idx_mailbox_timestamp ON mailbox_messages(timestamp);
            "#,
        )
        .execute(pool)
        .await?;

        sqlx::query(
            r#"
            CREATE INDEX IF NOT EXISTS idx_mailbox_processed ON mailbox_messages(processed);
            "#,
        )
        .execute(pool)
        .await?;

        info!("Mailbox schema created successfully");
        Ok(())
    }

    async fn handle_message(&mut self, msg: MailboxActorMessage) {
        match msg {
            MailboxActorMessage::QueueMessage { message } => {
                self.add_message(message).await;
            }
            MailboxActorMessage::ProcessQueue => {
                self.process_queue().await;
            }
            MailboxActorMessage::SaveGpsToDb { gps } => {
                if let Err(e) = self.save_gps_to_db(gps).await {
                    error!("Failed to save GPS to database: {}", e);
                }
            }
            MailboxActorMessage::GetNextMessage { respond_to } => {
                let next_msg = self.get_next_message();
                let _ = respond_to.send(next_msg);
            }
        }
    }

    async fn add_message(&mut self, message: MailboxMessage) {
        let priority = Self::get_priority(&message.msg_type);
        debug!(
            "Queuing message: {:?} with priority {}",
            message.msg_type, priority
        );

        // Broadcast GPTelemetry messages via ZMQ (bypass database)
        if matches!(message.msg_type, MailboxMessageType::GPTelemetry)
            && let Ok(msg_json) = serde_json::to_string(&message)
        {
            let _ = self.telemetry_sender.send(msg_json).await;
        }

        // Broadcast Log messages via ZMQ (bypass database)
        if matches!(message.msg_type, MailboxMessageType::Log)
            && let Ok(msg_json) = serde_json::to_string(&message)
        {
            let _ = self.log_sender.send(msg_json).await;
        }

        // Save to database (all messages)
        if let Err(e) = self.save_message_to_db(&message).await {
            error!("Failed to save message to database: {}", e);
        }

        self.queue.push(QueueItem { priority, message });
    }

    fn get_next_message(&mut self) -> Option<MailboxMessage> {
        self.queue.pop().map(|item| item.message)
    }

    async fn save_message_to_db(&self, message: &MailboxMessage) -> Result<(), sqlx::Error> {
        let timestamp = Utc::now().timestamp_millis();
        let msg_type = format!("{:?}", message.msg_type);

        sqlx::query(
            r#"
            INSERT INTO mailbox_messages (timestamp, msg_type, data)
            VALUES ($1, $2, $3)
            "#,
        )
        .bind(timestamp)
        .bind(msg_type)
        .bind(&message.data)
        .execute(&self.pool)
        .await?;

        Ok(())
    }

    async fn save_gps_to_db(&self, gps: Gps) -> Result<(), sqlx::Error> {
        let timestamp = gps.time.timestamp_millis();

        sqlx::query(
            r#"
            INSERT INTO gps (timestamp, lat, long, alt, heading)
            VALUES ($1, $2, $3, $4, $5)
            ON CONFLICT (timestamp) DO UPDATE
            SET lat = EXCLUDED.lat,
                long = EXCLUDED.long,
                alt = EXCLUDED.alt,
                heading = EXCLUDED.heading
            "#,
        )
        .bind(timestamp)
        .bind(gps.lat)
        .bind(gps.long)
        .bind(gps.alt)
        .bind(gps.heading)
        .execute(&self.pool)
        .await?;

        debug!("GPS data saved to database: timestamp={}", timestamp);
        Ok(())
    }

    async fn process_queue(&mut self) {
        while let Some(message) = self.get_next_message() {
            debug!("Processing message: {:?}", message.msg_type);
        }
    }

    fn get_priority(msg_type: &MailboxMessageType) -> u8 {
        match msg_type {
            MailboxMessageType::Error => 1,
            MailboxMessageType::TargetFound => 2,
            MailboxMessageType::Image => 3,
            MailboxMessageType::GNCTelemetry => 2,
            MailboxMessageType::GPTelemetry => 3,
            MailboxMessageType::Log => 5,
        }
    }
}

#[derive(Clone)]
pub struct MailboxHandle {
    sender: mpsc::Sender<MailboxActorMessage>,
}

#[allow(dead_code)]
impl MailboxHandle {
    pub async fn new() -> Self {
        let (sender, receiver) = mpsc::channel(100);
        let actor = Mailbox::new(receiver).await;

        tokio::spawn(run_mailbox(actor));

        Self { sender }
    }

    pub async fn queue_message(&self, message: MailboxMessage) {
        let msg = MailboxActorMessage::QueueMessage { message };
        let _ = self.sender.send(msg).await;
    }

    pub async fn save_gps(&self, gps: Gps) {
        let msg = MailboxActorMessage::SaveGpsToDb { gps };
        let _ = self.sender.send(msg).await;
    }

    pub async fn process_queue(&self) {
        let msg = MailboxActorMessage::ProcessQueue;
        let _ = self.sender.send(msg).await;
    }

    pub async fn get_next_message(&self) -> Option<MailboxMessage> {
        let (send, recv) = tokio::sync::oneshot::channel();
        let msg = MailboxActorMessage::GetNextMessage { respond_to: send };

        if self.sender.send(msg).await.is_ok() {
            recv.await.ok().flatten()
        } else {
            None
        }
    }
}

async fn run_mailbox(mut actor: Mailbox) {
    while let Some(msg) = actor.receiver.recv().await {
        actor.handle_message(msg).await;
    }
}

pub async fn start_gnc_listener(
    mailbox: MailboxHandle,
    shutdown_token: tokio_util::sync::CancellationToken,
) {
    info!("Starting GNC listener");

    tokio::task::spawn_blocking(move || {
        if let Err(e) = run_gnc_listener_blocking(mailbox, shutdown_token) {
            error!("GNC listener error: {}", e);
        }
    });
}

fn run_gnc_listener_blocking(
    mailbox: MailboxHandle,
    shutdown_token: tokio_util::sync::CancellationToken,
) -> Result<(), Box<dyn std::error::Error>> {
    let context = zmq::Context::new();
    let subscriber = context.socket(zmq::SUB)?;

    let gnc_port = env::var("GNC_PORT").unwrap_or("5556".to_string());
    let gnc_addr = format!("tcp://0.0.0.0:{}", gnc_port);

    subscriber.connect(&gnc_addr)?;
    subscriber.set_subscribe(b"")?; // Subscribe to all messages

    info!("GNC listener connected to {}", gnc_addr);

    subscriber.set_rcvtimeo(1000)?;

    loop {
        if shutdown_token.is_cancelled() {
            info!("GNC listener shutting down");
            break;
        }

        match subscriber.recv_bytes(0) {
            Ok(msg_bytes) => {
                match serde_json::from_slice::<MailboxMessage>(&msg_bytes) {
                    Ok(message) => {
                        debug!("[GNC] Received message: {:?}", message.msg_type);

                        // Special handling for GNC Telemetry data
                        if let MailboxMessageType::GNCTelemetry = message.msg_type
                            && let Ok(gnc_telem) =
                                serde_json::from_value::<GNCTelemetry>(message.data.clone())
                        {
                            // Convert to GPS for database storage
                            let gps = Gps {
                                lat: gnc_telem.lat,
                                long: gnc_telem.long,
                                alt: gnc_telem.alt,
                                heading: gnc_telem.heading,
                                time: chrono::DateTime::from_timestamp_millis(
                                    (gnc_telem.time * 1000.0) as i64,
                                )
                                .unwrap_or_else(chrono::Utc::now),
                            };

                            // Save GPS data to database
                            let mailbox_clone = mailbox.clone();
                            tokio::runtime::Handle::current().block_on(async {
                                mailbox_clone.save_gps(gps).await;
                            });
                        }

                        // Queue the message
                        let mailbox_clone = mailbox.clone();
                        tokio::runtime::Handle::current().block_on(async {
                            mailbox_clone.queue_message(message).await;
                        });
                    }
                    Err(e) => {
                        error!("Failed to parse GNC message: {}", e);
                    }
                }
            }
            Err(e) => {
                if e == zmq::Error::EAGAIN {
                    continue;
                } else {
                    error!("Failed to receive from GNC: {}", e);
                    std::thread::sleep(std::time::Duration::from_millis(100));
                }
            }
        }
    }

    Ok(())
}

// Logfile streaming server
pub async fn start_logfile_broadcaster(
    _mailbox: MailboxHandle,
    shutdown_token: tokio_util::sync::CancellationToken,
) {
    let log_addr = format!(
        "0.0.0.0:{}",
        env::var("LOGFILE_STREAM_PORT").unwrap_or("5559".to_string())
    );

    info!("Starting logfile broadcaster on {}", log_addr);

    tokio::spawn(async move {
        if let Err(e) = run_logfile_broadcaster(&log_addr, shutdown_token).await {
            error!("Logfile broadcaster error: {}", e);
        }
    });
}

async fn run_logfile_broadcaster(
    addr: &str,
    shutdown_token: tokio_util::sync::CancellationToken,
) -> Result<(), Box<dyn std::error::Error>> {
    let listener = TcpListener::bind(addr).await?;
    info!("Logfile broadcaster listening on {}", addr);

    let logfile_path =
        PathBuf::from(env::var("LOG_FILE_PATH").unwrap_or("./logs/feonix.log".to_string()));

    loop {
        tokio::select! {
            result = listener.accept() => {
                match result {
                    Ok((socket, addr)) => {
                        info!("New logfile subscriber connected: {}", addr);
                        let path = logfile_path.clone();
                        let shutdown_token = shutdown_token.clone();
                        tokio::spawn(async move {
                            if let Err(e) = stream_logfile(socket, path, shutdown_token).await {
                                error!("Error streaming logfile to {}: {}", addr, e);
                            }
                        });
                    }
                    Err(e) => {
                        error!("Failed to accept connection: {}", e);
                        break;
                    }
                }
            }
            _ = shutdown_token.cancelled() => {
                info!("Logfile broadcaster shutting down");
                break;
            }
        }
    }

    Ok(())
}

async fn stream_logfile(
    mut socket: TcpStream,
    logfile_path: PathBuf,
    shutdown_token: tokio_util::sync::CancellationToken,
) -> Result<(), Box<dyn std::error::Error>> {
    // Try to open file, create if doesn't exist
    let file = match File::open(&logfile_path) {
        Ok(f) => f,
        Err(_) => {
            // If file doesn't exist yet, wait and retry
            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
            File::open(&logfile_path)?
        }
    };

    let mut reader = BufReader::new(file);
    let mut position = 0u64;

    // Send existing log contents
    loop {
        let mut line = String::new();
        let bytes_read = reader.read_line(&mut line)?;

        if bytes_read == 0 {
            break;
        }

        socket.write_all(line.as_bytes()).await?;
        position += bytes_read as u64;
    }

    loop {
        tokio::select! {
            _ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => {
                let file = File::open(&logfile_path)?;
                let mut reader = BufReader::new(file);
                reader.seek(SeekFrom::Start(position))?;

                loop {
                    let mut line = String::new();
                    let bytes_read = reader.read_line(&mut line)?;

                    if bytes_read == 0 {
                        break;
                    }

                    socket.write_all(line.as_bytes()).await?;
                    position += bytes_read as u64;
                }
            }
            _ = shutdown_token.cancelled() => {
                info!("Logfile streaming cancelled, closing connection");
                break;
            }
        }
    }

    Ok(())
}

pub async fn start_telemetry_server(
    mailbox: MailboxHandle,
    shutdown_token: tokio_util::sync::CancellationToken,
) {
    info!("Telemetry server integrated with mailbox actor");
    start_gnc_listener(mailbox, shutdown_token).await;
}