poimen-memory/crates/mem-ingest/src/pi_session.rs

use anyhow::anyhow;
use futures::stream::{Stream, StreamExt};
use mem_chunk::RecordSource;
use mem_core::{Provenance, Record, Role};
use serde_json::Value;
use std::path::{Path, PathBuf};
use tokio::fs::File;
use tokio::io::{AsyncBufReadExt, BufReader};

/// A source that reads pi session JSONL files.
pub struct PiSessionSource {
    file_path: PathBuf,
}

impl PiSessionSource {
    /// Create a new Pi session source from a file path.
    pub fn new(file_path: PathBuf) -> Self {
        PiSessionSource { file_path }
    }

    /// Extract the project key from the session file's cwd.
    pub async fn read_project_key(&self) -> anyhow::Result<String> {
        let file = File::open(&self.file_path).await?;
        let mut reader = BufReader::new(file);
        let mut first_line = String::new();
        reader.read_line(&mut first_line).await?;

        let value: Value = serde_json::from_str(&first_line)?;
        let session_type = value
            .get("type")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow!("Missing 'type' field in session header"))?;

        if session_type != "session" {
            return Err(anyhow!(
                "Expected first line to have type='session', got '{}'",
                session_type
            ));
        }

        let cwd = value
            .get("cwd")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow!("Missing 'cwd' field in session header"))?;

        Ok(cwd.to_string())
    }
}

impl RecordSource for PiSessionSource {
    fn records(self) -> Box<dyn Stream<Item = Result<Record, String>> + Unpin> {
        let file_path = self.file_path.clone();
        let session_id = file_path
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("unknown")
            .to_string();

        // Create a stream that reads the file asynchronously
        let stream = futures::stream::once(async move {
            match read_pi_session_lines(&file_path, &session_id).await {
                Ok(records) => records,
                Err(e) => vec![Err(e.to_string())],
            }
        })
        .flat_map(futures::stream::iter);

        Box::new(stream.boxed())
    }
}

/// Read lines from a pi session file and convert them to Records.
pub async fn read_pi_session_lines(
    path: &Path,
    session_id: &str,
) -> anyhow::Result<Vec<Result<Record, String>>> {
    let file = File::open(path).await?;
    let reader = BufReader::new(file);
    let mut lines = reader.lines();
    let mut records = Vec::new();
    let mut line_num = 0u64;
    let mut skip_first = true;

    while let Some(line_result) = lines.next_line().await? {
        line_num += 1;
        
        // Skip the first line (session header)
        if skip_first {
            skip_first = false;
            continue;
        }

        match serde_json::from_str::<Value>(&line_result) {
            Ok(value) => {
                let record_type = value
                    .get("type")
                    .and_then(|v| v.as_str())
                    .unwrap_or("unknown");

                match record_type {
                    "message" => {
                        if let Ok(Some(record)) = parse_message(&value, session_id, line_num) {
                            records.push(Ok(record));
                        }
                    }
                    "compaction" => {
                        let msg = value
                            .get("message")
                            .and_then(|v| v.as_str())
                            .unwrap_or("Context compacted");
                        
                        let timestamp = parse_timestamp(&value);
                        let record = Record {
                            role: Role::System,
                            text: msg.to_string(),
                            timestamp,
                            provenance: Provenance {
                                source_id: format!("pi:{}", session_id),
                                offset: line_num,
                            },
                        };
                        records.push(Ok(record));
                    }
                    "model_change" | "thinking_level_change" => {
                        // Skip these record types
                    }
                    _ => {}
                }
            }
            Err(_) => {
                // Malformed line - skip with warning
                tracing::warn!(
                    "Skipping malformed JSON at line {}: {}",
                    line_num,
                    line_result.trim()
                );
            }
        }
    }

    Ok(records)
}

fn parse_message(
    value: &Value,
    session_id: &str,
    line_num: u64,
) -> Result<Option<Record>, String> {
    let message = value
        .get("message")
        .ok_or_else(|| "Missing 'message' field".to_string())?;

    let role_str = message
        .get("role")
        .and_then(|v| v.as_str())
        .ok_or_else(|| "Missing 'role' field in message".to_string())?;

    let role = match role_str {
        "user" => Role::User,
        "assistant" => Role::Assistant,
        "toolResult" => Role::ToolResult,
        _ => return Ok(None),
    };

    let content = message
        .get("content")
        .ok_or_else(|| "Missing 'content'".to_string())?;
    let text = flatten_content(content)?;

    let timestamp = parse_timestamp(message);
    let record_id = value
        .get("id")
        .and_then(|v| v.as_str())
        .unwrap_or("unknown");

    Ok(Some(Record {
        role,
        text,
        timestamp,
        provenance: Provenance {
            source_id: format!("pi:{}:{}", session_id, record_id),
            offset: line_num,
        },
    }))
}

/// Shared content flattening logic for all sources.
/// Handles: string, array of blocks, and structured objects.
pub fn flatten_content(content: &Value) -> Result<String, String> {
    if let Some(s) = content.as_str() {
        // Simple string content
        return Ok(s.to_string());
    }

    if let Some(arr) = content.as_array() {
        // Array of content blocks
        let mut text_parts = Vec::new();
        for block in arr {
            if let Some(block_type) = block.get("type").and_then(|v| v.as_str()) {
                match block_type {
                    "text" => {
                        if let Some(text) = block.get("text").and_then(|v| v.as_str()) {
                            text_parts.push(text.to_string());
                        }
                    }
                    "tool_use" => {
                        if let Some(name) = block.get("tool_name").and_then(|v| v.as_str()) {
                            text_parts.push(format!("[tool_use: {}]", name));
                        }
                    }
                    _ => {}
                }
            }
        }
        if !text_parts.is_empty() {
            return Ok(text_parts.join(" "));
        }
    }

    if let Some(obj) = content.as_object() {
        // Structured content (e.g., tool result)
        if let Some(text) = obj.get("text").and_then(|v| v.as_str()) {
            return Ok(text.to_string());
        }
        // Fallback: serialize the whole object
        return Ok(content.to_string());
    }

    Err("Could not flatten content".to_string())
}

fn parse_timestamp(value: &Value) -> time::OffsetDateTime {
    value
        .get("timestamp")
        .and_then(|v| v.as_str())
        .and_then(|s| {
            time::OffsetDateTime::parse(s, &time::format_description::well_known::Rfc3339).ok()
        })
        .unwrap_or_else(time::OffsetDateTime::now_utc)
}

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

    #[test]
    fn test_pi_session_source_creation() {
        let source = PiSessionSource::new(PathBuf::from("test.jsonl"));
        assert_eq!(source.file_path.to_string_lossy(), "test.jsonl");
    }
}