consciousness/poc-memory/src/bin/memory-search.rs

// memory-search: combined hook for session context loading + ambient memory retrieval
//
// Modes:
//   --hook     Run as Claude Code UserPromptSubmit hook (reads stdin, injects into conversation)
//   --debug    Replay last stashed input, dump every stage to stdout
//   --seen     Show the seen set for current session
//   (default)  No-op (future: manual search modes)

use clap::Parser;
use poc_memory::search::{self, AlgoStage};
use std::collections::{BTreeMap, HashSet};
use std::fs;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

#[derive(Parser)]
#[command(name = "memory-search")]
struct Args {
    /// Run as Claude Code hook (reads stdin, outputs for injection)
    #[arg(long)]
    hook: bool,

    /// Debug mode: replay last stashed input, dump every stage
    #[arg(short, long)]
    debug: bool,

    /// Show session state: seen set info
    #[arg(long)]
    seen: bool,

    /// Search query (bypasses stashed input, uses this as the prompt)
    #[arg(long, short)]
    query: Option<String>,

    /// Algorithm pipeline stages
    pipeline: Vec<String>,
}

const STASH_PATH: &str = "/tmp/claude-memory-search/last-input.json";

fn now_secs() -> u64 {
    SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs()
}
/// Max bytes per context chunk (hook output limit is ~10K chars)
const CHUNK_SIZE: usize = 9000;

struct Session {
    session_id: String,
    transcript_path: String,
    state_dir: PathBuf,
}

impl Session {
    fn load(hook: bool) -> Option<Self> {
        let state_dir = PathBuf::from("/tmp/claude-memory-search");
        fs::create_dir_all(&state_dir).ok();

        let input = if hook {
            let mut buf = String::new();
            io::stdin().read_to_string(&mut buf).unwrap_or_default();
            fs::write(STASH_PATH, &buf).ok();
            buf
        } else {
            fs::read_to_string(STASH_PATH).ok()?
        };

        let json: serde_json::Value = serde_json::from_str(&input).ok()?;
        let session_id = json["session_id"].as_str().unwrap_or("").to_string();
        if session_id.is_empty() { return None; }
        let transcript_path = json["transcript_path"].as_str().unwrap_or("").to_string();

        Some(Session { session_id, transcript_path, state_dir })
    }

    fn path(&self, prefix: &str) -> PathBuf {
        self.state_dir.join(format!("{}-{}", prefix, self.session_id))
    }
}

fn main() {
    // Daemon agent calls set POC_AGENT=1 — skip memory search.
    if std::env::var("POC_AGENT").is_ok() {
        return;
    }

    let args = Args::parse();

    if args.seen {
        show_seen();
        return;
    }

    // --query mode: skip all hook/context machinery, just search
    if let Some(ref query_str) = args.query {
        run_query_mode(query_str, &args);
        return;
    }

    let Some(session) = Session::load(args.hook) else { return };
    let debug = args.debug || !args.hook;

    let is_compaction = poc_memory::transcript::detect_new_compaction(
        &session.state_dir, &session.session_id, &session.transcript_path,
    );
    let cookie_path = session.path("cookie");
    let is_first = !cookie_path.exists();

    if is_first || is_compaction {
        // Rotate seen set on compaction, clear on first
        if is_compaction {
            fs::rename(&session.path("seen"), &session.path("seen-prev")).ok();
        } else {
            fs::remove_file(&session.path("seen")).ok();
            fs::remove_file(&session.path("seen-prev")).ok();
        }
        fs::remove_file(&session.path("returned")).ok();
    }

    if debug {
        println!("[memory-search] session={} is_first={} is_compaction={}",
            session.session_id, is_first, is_compaction);
    }

    if is_first || is_compaction {
        if is_first {
            fs::write(&cookie_path, generate_cookie()).ok();
        }

        if debug { println!("[memory-search] loading full context"); }

        if let Ok(output) = Command::new("poc-memory").args(["admin", "load-context"]).output() {
            if output.status.success() {
                let ctx = String::from_utf8_lossy(&output.stdout).to_string();
                if !ctx.trim().is_empty() {
                    let mut ctx_seen = load_seen(&session.state_dir, &session.session_id);
                    for line in ctx.lines() {
                        if line.starts_with("--- ") && line.ends_with(" ---") {
                            let inner = &line[4..line.len() - 4];
                            if let Some(paren) = inner.rfind(" (") {
                                let key = inner[..paren].trim();
                                mark_seen(&session.state_dir, &session.session_id, key, &mut ctx_seen);
                            }
                        }
                    }

                    let chunks = chunk_context(&ctx, CHUNK_SIZE);
                    if debug {
                        println!("[memory-search] context: {} bytes, {} chunks",
                            ctx.len(), chunks.len());
                    }

                    if let Some(first) = chunks.first() {
                        if args.hook { print!("{}", first); }
                    }
                    save_pending_chunks(&session.state_dir, &session.session_id, &chunks[1..]);
                }
            }
        }
    } else if let Some(chunk) = pop_pending_chunk(&session.state_dir, &session.session_id) {
        if debug {
            println!("[memory-search] drip-feeding pending chunk: {} bytes", chunk.len());
        }
        if args.hook { print!("{}", chunk); }
    }

    // Surface agent: consume previous result, inject memories, spawn next run
    if args.hook {
        surface_agent_cycle(&session);
    }

    cleanup_stale_files(&session.state_dir, Duration::from_secs(86400));
}

/// Direct query mode: search for a term without hook/stash machinery.
fn run_query_mode(query: &str, args: &Args) {
    let store = match poc_memory::store::Store::load() {
        Ok(s) => s,
        Err(e) => { eprintln!("failed to load store: {}", e); return; }
    };

    // Build terms from the query string
    let mut terms: BTreeMap<String, f64> = BTreeMap::new();
    let prompt_terms = search::extract_query_terms(query, 8);
    for word in prompt_terms.split_whitespace() {
        terms.entry(word.to_lowercase()).or_insert(1.0);
    }

    // Also check for exact node key match (the query itself, lowercased)
    let query_lower = query.to_lowercase();
    for (key, node) in &store.nodes {
        if node.deleted { continue; }
        if key.to_lowercase() == query_lower {
            terms.insert(query_lower.clone(), 10.0);
            break;
        }
    }

    println!("[query] terms: {:?}", terms);

    if terms.is_empty() {
        println!("[query] no terms extracted");
        return;
    }

    let graph = poc_memory::graph::build_graph_fast(&store);
    let (seeds, direct_hits) = search::match_seeds(&terms, &store);

    println!("[query] {} seeds", seeds.len());
    let mut sorted = seeds.clone();
    sorted.sort_by(|a, b| b.1.total_cmp(&a.1));
    for (key, score) in sorted.iter().take(20) {
        let marker = if direct_hits.contains(key) { "→" } else { " " };
        println!("  {} {:.4}  {}", marker, score, key);
    }

    let pipeline: Vec<AlgoStage> = if args.pipeline.is_empty() {
        vec![AlgoStage::parse("spread").unwrap()]
    } else {
        args.pipeline.iter()
            .filter_map(|a| AlgoStage::parse(a).ok())
            .collect()
    };

    let max_results = 50;
    let results = search::run_pipeline(&pipeline, seeds, &graph, &store, true, max_results);

    println!("\n[query] top {} results:", results.len().min(25));
    for (i, (key, score)) in results.iter().take(25).enumerate() {
        let marker = if direct_hits.contains(key) { "→" } else { " " };
        println!("  {:2}. {} [{:.4}] {}", i + 1, marker, score, key);
    }
}

/// Split context output into chunks of approximately `max_bytes`, breaking
/// at section boundaries ("--- KEY (group) ---" lines).
fn chunk_context(ctx: &str, max_bytes: usize) -> Vec<String> {
    // Split into sections at group boundaries, then merge small adjacent
    // sections into chunks up to max_bytes.
    let mut sections: Vec<String> = Vec::new();
    let mut current = String::new();

    for line in ctx.lines() {
        // Group headers start new sections
        if line.starts_with("--- ") && line.ends_with(" ---") && !current.is_empty() {
            sections.push(std::mem::take(&mut current));
        }
        if !current.is_empty() {
            current.push('\n');
        }
        current.push_str(line);
    }
    if !current.is_empty() {
        sections.push(current);
    }

    // Merge small sections into chunks, respecting max_bytes
    let mut chunks: Vec<String> = Vec::new();
    let mut chunk = String::new();
    for section in sections {
        if !chunk.is_empty() && chunk.len() + section.len() + 1 > max_bytes {
            chunks.push(std::mem::take(&mut chunk));
        }
        if !chunk.is_empty() {
            chunk.push('\n');
        }
        chunk.push_str(&section);
    }
    if !chunk.is_empty() {
        chunks.push(chunk);
    }
    chunks
}

/// Save remaining chunks to disk for drip-feeding on subsequent hook calls.
fn save_pending_chunks(dir: &Path, session_id: &str, chunks: &[String]) {
    let chunks_dir = dir.join(format!("chunks-{}", session_id));
    // Clear any old chunks
    let _ = fs::remove_dir_all(&chunks_dir);
    if chunks.is_empty() { return; }
    fs::create_dir_all(&chunks_dir).ok();
    for (i, chunk) in chunks.iter().enumerate() {
        let path = chunks_dir.join(format!("{:04}", i));
        fs::write(path, chunk).ok();
    }
}

/// Pop the next pending chunk (lowest numbered file). Returns None if no chunks remain.
fn pop_pending_chunk(dir: &Path, session_id: &str) -> Option<String> {
    let chunks_dir = dir.join(format!("chunks-{}", session_id));
    if !chunks_dir.exists() { return None; }

    let mut entries: Vec<_> = fs::read_dir(&chunks_dir).ok()?
        .flatten()
        .filter(|e| e.file_type().map(|t| t.is_file()).unwrap_or(false))
        .collect();
    entries.sort_by_key(|e| e.file_name());

    let first = entries.first()?;
    let content = fs::read_to_string(first.path()).ok()?;
    fs::remove_file(first.path()).ok();

    // Clean up directory if empty
    if fs::read_dir(&chunks_dir).ok().map(|mut d| d.next().is_none()).unwrap_or(true) {
        fs::remove_dir(&chunks_dir).ok();
    }

    Some(content)
}

fn generate_cookie() -> String {
    uuid::Uuid::new_v4().as_simple().to_string()[..12].to_string()
}


/// Parse a seen-file line: "TIMESTAMP\tKEY" or legacy "KEY"
fn parse_seen_line(line: &str) -> &str {
    line.split_once('\t').map(|(_, key)| key).unwrap_or(line)
}

fn load_seen(dir: &Path, session_id: &str) -> HashSet<String> {
    let path = dir.join(format!("seen-{}", session_id));
    if path.exists() {
        let set: HashSet<String> = fs::read_to_string(&path)
            .unwrap_or_default()
            .lines()
            .filter(|s| !s.is_empty())
            .map(|s| parse_seen_line(s).to_string())
            .collect();
        set
    } else {
        HashSet::new()
    }
}

fn mark_seen(dir: &Path, session_id: &str, key: &str, seen: &mut HashSet<String>) {
    if !seen.insert(key.to_string()) { return; }
    let path = dir.join(format!("seen-{}", session_id));
    if let Ok(mut f) = fs::OpenOptions::new().create(true).append(true).open(path) {
        let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M:%S");
        writeln!(f, "{}\t{}", ts, key).ok();
    }
}


/// Surface agent lifecycle: check if previous agent finished, consume results,
/// render and inject unseen memories, spawn next agent run.
fn surface_agent_cycle(session: &Session) {
    let result_path = session.state_dir.join(format!("surface-result-{}", session.session_id));
    let pid_path = session.state_dir.join(format!("surface-pid-{}", session.session_id));

    let surface_timeout = poc_memory::config::get()
        .surface_timeout_secs
        .unwrap_or(120) as u64;

    // Check if previous agent is done
    let agent_done = match fs::read_to_string(&pid_path) {
        Ok(content) => {
            let parts: Vec<&str> = content.split('\t').collect();
            let pid: u32 = parts.first().and_then(|s| s.trim().parse().ok()).unwrap_or(0);
            let start_ts: u64 = parts.get(1).and_then(|s| s.trim().parse().ok()).unwrap_or(0);
            if pid == 0 { true }
            else {
                let alive = unsafe { libc::kill(pid as i32, 0) == 0 };
                if !alive { true }
                else if now_secs().saturating_sub(start_ts) > surface_timeout {
                    unsafe { libc::kill(pid as i32, libc::SIGTERM); }
                    true
                } else { false }
            }
        }
        Err(_) => true,
    };

    if !agent_done { return; }

    // Consume result
    if let Ok(result) = fs::read_to_string(&result_path) {
        if !result.trim().is_empty() {
            let tail_lines: Vec<&str> = result.lines().rev()
                .filter(|l| !l.trim().is_empty()).take(8).collect();
            let has_new = tail_lines.iter().any(|l| l.starts_with("NEW RELEVANT MEMORIES:"));
            let has_none = tail_lines.iter().any(|l| l.starts_with("NO NEW RELEVANT MEMORIES"));

            if has_new {
                let after_marker = result.rsplit_once("NEW RELEVANT MEMORIES:")
                    .map(|(_, rest)| rest).unwrap_or("");
                let keys: Vec<String> = after_marker.lines()
                    .map(|l| l.trim().trim_start_matches("- ").trim().to_string())
                    .filter(|l| !l.is_empty()).collect();

                // Render and inject unseen keys
                let mut seen = load_seen(&session.state_dir, &session.session_id);
                let seen_path = session.path("seen");
                for key in &keys {
                    if !seen.insert(key.clone()) { continue; }
                    if let Ok(output) = Command::new("poc-memory").args(["render", key]).output() {
                        if output.status.success() {
                            let content = String::from_utf8_lossy(&output.stdout);
                            if !content.trim().is_empty() {
                                println!("--- {} (surfaced) ---", key);
                                print!("{}", content);
                                if let Ok(mut f) = fs::OpenOptions::new()
                                    .create(true).append(true).open(&seen_path) {
                                    let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M:%S");
                                    writeln!(f, "{}\t{}", ts, key).ok();
                                }
                            }
                        }
                    }
                }
            } else if !has_none {
                let log_dir = poc_memory::store::memory_dir().join("logs");
                fs::create_dir_all(&log_dir).ok();
                let log_path = log_dir.join("surface-errors.log");
                if let Ok(mut f) = fs::OpenOptions::new().create(true).append(true).open(&log_path) {
                    let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M:%S");
                    let last = tail_lines.first().unwrap_or(&"");
                    let _ = writeln!(f, "[{}] unexpected surface output: {}", ts, last);
                }
            }
        }
    }
    fs::remove_file(&result_path).ok();
    fs::remove_file(&pid_path).ok();

    // Spawn next surface agent
    if let Ok(output_file) = fs::File::create(&result_path) {
        if let Ok(child) = Command::new("poc-memory")
            .args(["agent", "run", "surface", "--count", "1", "--local"])
            .env("POC_SESSION_ID", &session.session_id)
            .stdout(output_file)
            .stderr(std::process::Stdio::null())
            .spawn()
        {
            let pid = child.id();
            let ts = now_secs();
            if let Ok(mut f) = fs::File::create(&pid_path) {
                write!(f, "{}\t{}", pid, ts).ok();
            }
        }
    }
}

fn show_seen() {
    let Some(session) = Session::load(false) else {
        eprintln!("No session state available");
        return;
    };

    println!("Session: {}", session.session_id);

    if let Ok(cookie) = fs::read_to_string(&session.path("cookie")) {
        println!("Cookie:  {}", cookie.trim());
    }

    match fs::read_to_string(&session.path("compaction")) {
        Ok(s) => {
            let offset: u64 = s.trim().parse().unwrap_or(0);
            let ts = poc_memory::transcript::compaction_timestamp(&session.transcript_path, offset);
            match ts {
                Some(t) => println!("Last compaction: offset {} ({})", offset, t),
                None => println!("Last compaction: offset {}", offset),
            }
        }
        Err(_) => println!("Last compaction: none detected"),
    }

    let pending = fs::read_dir(&session.path("chunks")).ok()
        .map(|d| d.flatten().count()).unwrap_or(0);
    if pending > 0 {
        println!("Pending chunks: {}", pending);
    }

    for (label, suffix) in [("Current seen set", ""), ("Previous seen set (pre-compaction)", "-prev")] {
        let path = session.state_dir.join(format!("seen{}-{}", suffix, session.session_id));
        let content = fs::read_to_string(&path).unwrap_or_default();
        let lines: Vec<&str> = content.lines().filter(|s| !s.is_empty()).collect();
        if lines.is_empty() { continue; }

        println!("\n{} ({}):", label, lines.len());
        for line in &lines { println!("  {}", line); }
    }
}

fn cleanup_stale_files(dir: &Path, max_age: Duration) {
    let entries = match fs::read_dir(dir) {
        Ok(e) => e,
        Err(_) => return,
    };
    let cutoff = SystemTime::now() - max_age;
    for entry in entries.flatten() {
        if let Ok(meta) = entry.metadata() {
            if let Ok(modified) = meta.modified() {
                if modified < cutoff {
                    fs::remove_file(entry.path()).ok();
                }
            }
        }
    }
}