as-rust/strategy/src/incremental_json_writer.rs

use serde::Serialize;
use serde_json;
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufWriter, Write};
use std::path::PathBuf;
use rust_decimal::Decimal;

// --------------------- 错误处理 ---------------------
#[derive(Debug)]
pub enum WriterError {
    Serde(serde_json::Error),
    Io(io::Error),
    InvalidPath(String), // 用于处理路径错误
}

impl std::fmt::Display for WriterError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            WriterError::Serde(e) => write!(f, "JSON serialization error: {}", e),
            WriterError::Io(e) => write!(f, "File I/O error: {}", e),
            WriterError::InvalidPath(s) => write!(f, "Invalid path: {}", s),
        }
    }
}

impl std::error::Error for WriterError {}

impl From<serde_json::Error> for WriterError {
    fn from(err: serde_json::Error) -> Self {
        WriterError::Serde(err)
    }
}

impl From<io::Error> for WriterError {
    fn from(err: io::Error) -> Self {
        WriterError::Io(err)
    }
}

// --------------------- 核心写入器 ---------------------
pub struct IncrementalJsonWriter<T> where T: Serialize {
    base_file_path: PathBuf, // 文件的基础路径和名称前缀 (e.g., "data/my_log")
    current_file_idx: usize, // 当前文件的索引 (e.g., 0, 1, 2...)
    max_file_size_bytes: u64, // 最大文件大小 (字节)
    current_file_size_bytes: u64, // 当前文件已写入的字节数
    batch_size: usize,
    buffer: Vec<T>,
    // 使用 BufWriter 来提高写入效率，每次写入不需要立即刷新到磁盘
    current_writer: Option<BufWriter<File>>,
}

impl<T> IncrementalJsonWriter<T> where T: Serialize {
    /// 创建一个新的 IncrementalJsonWriter 实例。
    /// `base_path_prefix` 是文件路径和前缀，例如 "logs/events" 会生成 "logs/events_0.json", "logs/events_1.json"
    /// `max_file_size_mb` 是单个文件的最大大小，以 MB 为单位。
    pub fn new(base_path_prefix: &str, max_file_size_mb: u64, batch_size: usize) -> Result<Self, WriterError> {
        let base_file_path = PathBuf::from(base_path_prefix);
        let max_file_size_bytes = max_file_size_mb * 1024 * 1024; // MB 转换为字节

        // 确保父目录存在
        if let Some(parent) = base_file_path.parent() {
            if parent.as_os_str() != "" { // 避免在当前目录创建时创建不存在的父目录
                fs::create_dir_all(parent)?;
            }
        }

        Ok(IncrementalJsonWriter {
            base_file_path,
            current_file_idx: 0, // 从文件索引 0 开始
            max_file_size_bytes,
            current_file_size_bytes: 0,
            current_writer: None,
            batch_size,
            buffer: Vec::with_capacity(batch_size),
        })
    }

    /// 获取完整的文件路径 (e.g., "data/my_log_0.json")
    fn get_full_file_path(&self, index: usize) -> PathBuf {
        let mut path = self.base_file_path.clone(); // 复制基础路径
        let filename = format!("{}_{}.json",
                               path.file_name().unwrap_or_default().to_string_lossy(), // 获取文件名部分
                               index
        );
        path.set_file_name(filename); // 设置新的带索引的文件名
        path
    }

    /// 打开或创建一个新的文件，并设置 self.current_writer。
    /// 会查找最新的文件索引并从该文件继续写入，或者创建新的文件。
    fn open_current_file(&mut self) -> Result<(), WriterError> {
        // 先检查是否有打开的文件，如果有，先刷新并关闭
        if let Some(writer) = self.current_writer.take() {
            writer.into_inner().unwrap().sync_all()?; // 确保所有数据写入磁盘
        }

        // 查找最新的文件索引，以便在程序重启后可以继续写入
        // 这需要遍历所有可能的文件，找到最大的那个索引
        let mut latest_idx = 0;
        loop {
            let candidate_path = self.get_full_file_path(latest_idx);
            if candidate_path.exists() {
                latest_idx += 1;
            } else {
                break;
            }
        }
        // 如果最新索引文件存在，则从该文件继续。否则，使用 latest_idx-1 + 1 作为新文件。
        // 但为了简单和确定性，我们总是从索引0开始，然后简单地找到最后一个文件索引加上1
        // 注意：这种查找方式适用于程序每次启动都重新开始写文件的情况。
        // 如果需要接续之前未写满的文件，逻辑会更复杂：需要read_dir并解析文件名。
        // 为了本例的增量写入和文件切换，我们简单地从 current_file_idx + 1 开始。
        // 这里为了演示“增量写入”，我们默认从0开始写，然后自动递增。
        // 如果要实现“断点续写”，则需要在启动时扫描目录找到最大的 index 文件，并读取其大小。
        // For simplicity, let's just use current_file_idx as the next one to create
        // and always start from 0 if no file exists.

        let mut target_idx = self.current_file_idx;
        let mut target_path = self.get_full_file_path(target_idx);

        // 如果文件已存在，我们应该尝试从它上次的大小继续，或者直接创建新文件。
        // 这里策略是：如果当前文件已满或不存在，就开下一个文件。
        // 简化的策略：总是从 current_file_idx 开始，如果已存在，且未满，就追加；如果已满，就开下一个。
        // 更简单的策略：每次调用 open_current_file 都开一个新文件，或者检查当前文件是否真的需要切换

        // 寻找一个适合写入的新文件
        loop {
            let next_file_path = self.get_full_file_path(self.current_file_idx);
            if next_file_path.exists() {
                let metadata = fs::metadata(&next_file_path)?;
                if metadata.len() >= self.max_file_size_bytes {
                    self.current_file_idx += 1; // 切换到下一个文件
                    self.current_file_size_bytes = 0; // 重置大小
                    continue; // 继续循环，尝试下一个文件
                } else {
                    // 文件存在且未满，尝试追加写入
                    self.current_file_size_bytes = metadata.len();
                    target_path = next_file_path; // 找到目标文件
                    break;
                }
            } else {
                // 文件不存在，这是个新文件
                self.current_file_size_bytes = 0;
                target_path = next_file_path;
                break;
            }
        }

        let file = OpenOptions::new()
            .create(true) // 如果文件不存在则创建
            .append(true) // 以追加模式打开
            .open(&target_path)?;

        self.current_writer = Some(BufWriter::new(file));
        Ok(())
    }

    /// 添加一条日志到缓冲区。如果缓冲区满，会触发一次写入。
    pub fn add_entry(&mut self, entry: T) -> Result<(), Box<dyn std::error::Error>> {
        self.buffer.push(entry);

        // 如果缓冲区达到批处理大小，执行写入
        if self.buffer.len() >= self.batch_size {
            self.flush()?;
        }

        Ok(())
    }

    /// 写入单个结构体数据到文件中。
    pub fn flush(&mut self) -> Result<(), WriterError> {
        // 序列化缓冲区中的每条日志为 JSON 字符串，并用换行符连接
        // 注意：这里每个 entry 都是一个独立的 JSON 对象
        let mut batch_json_lines = String::new();
        for entry in &self.buffer {
            // 将每条日志序列化为 JSON 字符串
            let json_line = serde_json::to_string(entry)?;
            batch_json_lines.push_str(&json_line);
            batch_json_lines.push('\n'); // JSON Lines 格式需要换行符
        }

        let data_len = batch_json_lines.as_bytes().len() as u64;

        // 如果没有打开的文件，或者当前文件已满，则打开新文件
        if self.current_writer.is_none() || self.current_file_size_bytes + data_len > self.max_file_size_bytes {
            // 如果 current_writer.is_some()，说明前一个文件已满，需要刷新并关闭
            if let Some(writer) = self.current_writer.take() {
                writer.into_inner().unwrap().sync_all()?; // 确保所有数据写入磁盘
                self.current_file_idx += 1; // 切换到下一个文件索引
                self.current_file_size_bytes = 0; // 重置计数
            }
            // 打开新的文件
            self.open_current_file()?;
        }

        // 写入数据
        if let Some(writer) = self.current_writer.as_mut() {
            writer.write_all(batch_json_lines.as_bytes())?;
            self.current_file_size_bytes += data_len;

            writer.flush()?; // 确保 BufWriter 的内容写入到文件
            writer.get_ref().sync_all()?; // 确保数据写入磁盘
            
            // 清空缓冲区，准备下一批
            self.buffer.clear();
            self.buffer.reserve(self.batch_size); // 重新预分配容量
        } else {
            // 这应该不会发生，因为前面已经确保了 writer 是 Some
            return Err(WriterError::Io(io::Error::new(io::ErrorKind::Other, "209 Writer not initialized unexpectedly")));
        }

        Ok(())
    }
}

// 当 IncrementalJsonWriter 实例被 Drop 时，确保文件被关闭并数据写入磁盘
impl<T>  Drop for IncrementalJsonWriter<T> where T: Serialize {
    fn drop(&mut self) {
        if let Some(mut writer) = self.current_writer.take() {
            // Try to flush and sync, but ignore errors during shutdown
            let _ = writer.flush();
            let _ = writer.into_inner().and_then(|f| Ok(f.sync_all()));
        }
    }
}
#[derive(Serialize, Debug)]
pub struct OrderLog{
    // 挂单信息
    pub order_info: Vec<Vec<String>>,
    // 记录时间
    pub record_timestamp: i64
}

#[derive(Serialize, Debug)]
pub struct FixPriceLog{
    // 定价
    pub fix_price: Decimal,
    // 记录时间
    pub record_timestamp: i64
}

// --------------------- 示例数据结构 ---------------------
#[derive(Serialize, Debug)]
struct LogEntry {
    timestamp: String,
    level: String,
    message: String,
    session_id: String,
    #[serde(skip_serializing_if = "Option::is_none")] // 如果为None则不序列化
    user_id: Option<u32>,
}

#[cfg(test)]
mod tests {
    use crate::incremental_json_writer::{IncrementalJsonWriter, LogEntry};
    #[test]
    fn test_incremental_json_writer() {
        // 初始化写入器，文件前缀为 "logs/app_log"，最大文件大小为 1MB
        let base_path = "../logs/app_log";
        let max_size_mb = 10; // 1 MB
        let mut writer = IncrementalJsonWriter::<LogEntry>::new(base_path, max_size_mb, 100usize).unwrap();

        println!("Starting to write log entries...");

        let mut user_id_counter = 1000;

        // 写入大量数据，直到触发文件切换
        for i in 0..200000 { // 写入20万条记录，看看文件切换效果
            let timestamp = chrono::Local::now().to_rfc3339();
            let level = if i % 10 == 0 { "ERROR" } else { "INFO" }.to_string();
            let message = format!("Processing record {} for user {}...", i, user_id_counter);
            let session_id = format!("sess_{}", i / 100);

            let user_id = if i % 5 == 0 {
                user_id_counter += 1;
                Some(user_id_counter)
            } else {
                None
            };

            let entry = LogEntry {
                timestamp,
                level,
                message,
                session_id,
                user_id,
            };

            writer.add_entry(entry).unwrap();
        }
        if writer.buffer.len() > 0 {
            writer.flush().unwrap(); // 确保所有数据都已写入
        }
        println!("\nFinished writing records. Final flush...");
        println!("Logs written to files under the '{}' directory.", base_path);
    }
}