keep/src/services/async_item_service.rs

use crate::common::PIPESIZE;
use crate::config::Settings;
use crate::services::error::CoreError;
use crate::services::item_service::ItemService;
use crate::services::types::{ItemWithContent, ItemWithMeta};
use clap::Command;
use rusqlite::Connection;
use std::collections::HashMap;
use std::io::Read;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;

/// An asynchronous wrapper around the `ItemService` for use in async contexts like the web server.
/// It uses `tokio::task::spawn_blocking` to run synchronous database and filesystem operations
/// on a dedicated thread pool, preventing them from blocking the async runtime.
#[allow(dead_code)]
pub struct AsyncItemService {
    pub data_dir: PathBuf,
    db: Arc<Mutex<Connection>>,
    item_service: Arc<ItemService>,
    cmd: Arc<Mutex<Command>>,
    settings: Arc<Settings>,
}

#[allow(dead_code)]
impl AsyncItemService {
    pub fn new(
        data_dir: PathBuf, 
        db: Arc<Mutex<Connection>>, 
        item_service: Arc<ItemService>,
        cmd: Arc<Mutex<Command>>,
        settings: Arc<Settings>,
    ) -> Self {
        Self { 
            data_dir, 
            db, 
            item_service,
            cmd,
            settings,
        }
    }

    pub async fn get_item(&self, id: i64) -> Result<ItemWithMeta, CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();

        tokio::task::spawn_blocking(move || {
            let conn = db.blocking_lock();
            item_service.get_item(&conn, id)
        })
        .await
        .unwrap() // Propagate panics from spawn_blocking
    }

    pub async fn get_item_content(&self, id: i64) -> Result<ItemWithContent, CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();

        tokio::task::spawn_blocking(move || {
            let conn = db.blocking_lock();
            item_service.get_item_content(&conn, id)
        })
        .await
        .unwrap()
    }

    pub async fn get_item_content_info(
        &self,
        id: i64,
        head_bytes: Option<usize>,
        head_words: Option<usize>,
        head_lines: Option<usize>,
        tail_bytes: Option<usize>,
        tail_words: Option<usize>,
        tail_lines: Option<usize>,
        line_start: Option<usize>,
        line_end: Option<usize>,
        grep: Option<String>,
    ) -> Result<(Vec<u8>, String, bool), CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();

        tokio::task::spawn_blocking(move || {
            let conn = db.blocking_lock();
            item_service.get_item_content_info(
                &conn, 
                id, 
                head_bytes, 
                head_words, 
                head_lines, 
                tail_bytes, 
                tail_words, 
                tail_lines, 
                line_start, 
                line_end,
                grep
            )
        })
        .await
        .unwrap()
    }

    pub async fn stream_item_content_by_id(
        &self,
        item_id: i64,
        allow_binary: bool,
        offset: u64,
        length: u64,
    ) -> Result<(std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<tokio_util::bytes::Bytes, std::io::Error>> + Send>>, String), CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();

        // Get item content
        let content = tokio::task::spawn_blocking(move || {
            let conn = db.blocking_lock();
            let item_with_content = item_service.get_item_content(&conn, item_id)?;
            Ok::<_, CoreError>(item_with_content.content)
        })
        .await
        .unwrap()?;

        // Clone content for use in the binary check closure
        let content_clone = content.clone();

        // Get metadata to determine MIME type and binary status
        let (mime_type, is_binary) = {
            let db = self.db.clone();
            let item_service = self.item_service.clone();
            tokio::task::spawn_blocking(move || {
                let conn = db.blocking_lock();
                let item_with_meta = item_service.get_item(&conn, item_id)?;
                let metadata = item_with_meta.meta_as_map();

                let mime_type = metadata
                    .get("mime_type")
                    .map(|s| s.to_string())
                    .unwrap_or_else(|| "application/octet-stream".to_string());

                let is_binary = if let Some(text_val) = metadata.get("text") {
                    text_val == "false"
                } else {
                    crate::common::is_binary::is_binary(&content_clone)
                };

                Ok::<_, CoreError>((mime_type, is_binary))
            })
            .await
            .unwrap()?
        };

        // Check if content is binary when allow_binary is false
        if !allow_binary && is_binary {
            return Err(CoreError::InvalidInput("Binary content not allowed".to_string()));
        }

        // Create a stream that reads only the requested portion
        let content_len = content.len() as u64;

        // Apply offset and length constraints
        let start = std::cmp::min(offset, content_len);
        let end = if length > 0 {
            std::cmp::min(start + length, content_len)
        } else {
            content_len
        };

        let stream = if start < content_len {
            let chunk = tokio_util::bytes::Bytes::from(content[start as usize..end as usize].to_vec());
            Box::pin(tokio_stream::iter(vec![Ok(chunk)]))
        } else {
            Box::pin(tokio_stream::iter(vec![]))
        };

        Ok((stream, mime_type))
    }

    pub async fn stream_item_content_by_id_with_metadata(
        &self,
        item_id: i64,
        metadata: &HashMap<String, String>,
        allow_binary: bool,
        offset: u64,
        length: u64,
        head_bytes: Option<usize>,
        head_words: Option<usize>,
        head_lines: Option<usize>,
        tail_bytes: Option<usize>,
        tail_words: Option<usize>,
        tail_lines: Option<usize>,
        line_start: Option<usize>,
        line_end: Option<usize>,
        grep: Option<String>,
    ) -> Result<(std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<tokio_util::bytes::Bytes, std::io::Error>> + Send>>, String), CoreError> {
        // Use provided metadata to determine MIME type and binary status
        let mime_type = metadata
            .get("mime_type")
            .map(|s| s.to_string())
            .unwrap_or_else(|| "application/octet-stream".to_string());

        // Check if content is binary when allow_binary is false
        if !allow_binary {
            let is_binary = if let Some(text_val) = metadata.get("text") {
                text_val == "false"
            } else {
                // Get binary status using streaming approach
                let (_, _, is_binary) = self.get_item_content_info_streaming(
                    item_id,
                    head_bytes,
                    head_words,
                    head_lines,
                    tail_bytes,
                    tail_words,
                    tail_lines,
                    line_start,
                    line_end
                ).await?;
                is_binary
            };

            if is_binary {
                return Err(CoreError::InvalidInput("Binary content not allowed".to_string()));
            }
        }

        // Get a streaming reader for the content with filtering applied
        let reader = {
            let db = self.db.clone();
            let item_service = self.item_service.clone();
            tokio::task::spawn_blocking(move || {
                let conn = db.blocking_lock();
                item_service.get_item_content_info_streaming(
                    &conn, 
                    item_id, 
                    head_bytes, 
                    head_words, 
                    head_lines, 
                    tail_bytes, 
                    tail_words, 
                    tail_lines, 
                    line_start, 
                    line_end,
                    grep
                ).map(|(reader, _, _)| reader)
            })
            .await
            .unwrap()?
        };

        // Convert the reader into an async stream manually
        use tokio_util::bytes::Bytes;

        // Create a channel to stream data between the blocking thread and async runtime
        let (tx, rx) = tokio::sync::mpsc::channel(1);

        // Spawn a blocking task to read from the reader and send chunks
        tokio::task::spawn_blocking(move || {
            let mut reader = reader;
            // Apply offset by reading and discarding bytes
            if offset > 0 {
                let mut remaining = offset;
                let mut buf = [0; PIPESIZE];
                while remaining > 0 {
                    let to_read = std::cmp::min(remaining, buf.len() as u64);
                    match reader.read(&mut buf[..to_read as usize]) {
                        Ok(0) => break, // EOF reached before offset
                        Ok(n) => remaining -= n as u64,
                        Err(e) => {
                            let _ = tx.blocking_send(Err(e));
                            return;
                        }
                    }
                }
            }

            // Read and send data up to the specified length
            let mut remaining_length = length;
            let mut buffer = [0; PIPESIZE];

            loop {
                // Determine how much to read in this iteration
                let to_read = if length > 0 {
                    // If length is specified, don't read more than remaining_length
                    std::cmp::min(remaining_length, buffer.len() as u64) as usize
                } else {
                    buffer.len()
                };

                if to_read == 0 {
                    break; // We've read the requested length
                }

                match reader.read(&mut buffer[..to_read]) {
                    Ok(0) => break, // EOF
                    Ok(n) => {
                        let chunk = Bytes::copy_from_slice(&buffer[..n]);
                        // Block on sending to the channel
                        if tx.blocking_send(Ok(chunk)).is_err() {
                            break; // Receiver dropped
                        }
                        if length > 0 {
                            remaining_length -= n as u64;
                            if remaining_length == 0 {
                                break; // Reached the requested length
                            }
                        }
                    }
                    Err(e) => {
                        let _ = tx.blocking_send(Err(e));
                        break;
                    }
                }
            }
        });

        // Convert the receiver into a stream
        let stream = tokio_stream::wrappers::ReceiverStream::new(rx);

        Ok((Box::pin(stream), mime_type))
    }

    pub async fn get_item_content_info_streaming(
        &self,
        item_id: i64,
        head_bytes: Option<usize>,
        head_words: Option<usize>,
        head_lines: Option<usize>,
        tail_bytes: Option<usize>,
        tail_words: Option<usize>,
        tail_lines: Option<usize>,
        line_start: Option<usize>,
        line_end: Option<usize>,
        grep: Option<String>,
    ) -> Result<(Box<dyn Read + Send>, String, bool), CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();

        tokio::task::spawn_blocking(move || {
            let conn = db.blocking_lock();
            item_service.get_item_content_info_streaming(
                &conn, 
                item_id, 
                head_bytes, 
                head_words, 
                head_lines, 
                tail_bytes, 
                tail_words, 
                tail_lines, 
                line_start, 
                line_end,
                grep
            )
        })
        .await
        .unwrap()
    }

    pub async fn find_item(
        &self,
        ids: Vec<i64>,
        tags: Vec<String>,
        meta: HashMap<String, String>,
    ) -> Result<ItemWithMeta, CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();
        
        tokio::task::spawn_blocking(move || {
            let conn = db.blocking_lock();
            item_service.find_item(&conn, &ids, &tags, &meta)
        })
        .await
        .unwrap()
    }

    pub async fn list_items(
        &self,
        tags: Vec<String>,
        meta: HashMap<String, String>,
    ) -> Result<Vec<ItemWithMeta>, CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();
        
        tokio::task::spawn_blocking(move || {
            let conn = db.blocking_lock();
            item_service.list_items(&conn, &tags, &meta)
        })
        .await
        .unwrap()
    }

    pub async fn delete_item(&self, id: i64) -> Result<(), CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();
        
        tokio::task::spawn_blocking(move || {
            let mut conn = db.blocking_lock();
            item_service.delete_item(&mut conn, id)
        })
        .await
        .unwrap()
    }

    pub async fn save_item_from_mcp(
        &self,
        content: Vec<u8>,
        tags: Vec<String>,
        metadata: HashMap<String, String>,
    ) -> Result<ItemWithMeta, CoreError> {
        let db = self.db.clone();
        let item_service = self.item_service.clone();
        let cmd = self.cmd.clone();
        let settings = self.settings.clone();

        tokio::task::spawn_blocking(move || {
            let mut conn = db.blocking_lock();
            let mut cmd = cmd.blocking_lock();
            let settings = settings.as_ref();
            item_service.save_item_from_mcp(&content, &tags, &metadata, &mut cmd, settings, &mut conn)
        })
        .await
        .unwrap()
    }
}