hermes-tool 1.8.44

//! Index management operations: create, index, commit, merge, info, warmup

use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufRead, BufReader};
use std::path::PathBuf;

use anyhow::{Context, Result};
use tracing::info;

use hermes_core::{Document, FsDirectory, IndexConfig, IndexWriter, parse_schema};

use crate::release_memory_to_os;

pub async fn create_index(index_path: PathBuf, schema_path: PathBuf) -> Result<()> {
    let schema_content = fs::read_to_string(&schema_path)
        .with_context(|| format!("Failed to read schema file: {:?}", schema_path))?;

    let schema = parse_schema(&schema_content)
        .map_err(|e| anyhow::anyhow!("Failed to parse schema: {}", e))?;

    info!("Parsed schema with {} fields", schema.fields().count());

    std::fs::create_dir_all(&index_path)
        .with_context(|| format!("Failed to create index directory: {:?}", index_path))?;

    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();

    let _writer = IndexWriter::create(dir, schema, config).await?;

    info!("Created index at {:?}", index_path);

    Ok(())
}

pub async fn init_index_from_sdl(index_path: PathBuf, sdl: String) -> Result<()> {
    let schema =
        parse_schema(&sdl).map_err(|e| anyhow::anyhow!("Failed to parse schema: {}", e))?;

    std::fs::create_dir_all(&index_path)
        .with_context(|| format!("Failed to create index directory: {:?}", index_path))?;

    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();

    let _writer = IndexWriter::create(dir, schema.clone(), config).await?;

    info!("Created index at {:?}", index_path);
    info!("Schema has {} fields", schema.fields().count());

    Ok(())
}

async fn index_from_reader<R: BufRead>(
    writer: &mut IndexWriter<FsDirectory>,
    reader: R,
    progress_interval: usize,
) -> Result<usize> {
    let schema = writer.schema().clone();
    let mut count = 0usize;
    let mut errors = 0usize;
    let start_time = std::time::Instant::now();

    for line in reader.lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }

        let json: serde_json::Value = match sonic_rs::from_str(&line) {
            Ok(v) => v,
            Err(e) => {
                if errors < 10 {
                    tracing::warn!("Failed to parse JSON at line {}: {}", count + 1, e);
                }
                errors += 1;
                continue;
            }
        };

        let doc = match Document::from_json(&json, &schema) {
            Some(d) => d,
            None => {
                if errors < 10 {
                    tracing::warn!("Failed to parse document at line {}", count + 1);
                }
                errors += 1;
                continue;
            }
        };

        writer.add_document(doc)?;
        count += 1;

        if progress_interval > 0 && count.is_multiple_of(progress_interval) {
            let elapsed = start_time.elapsed().as_secs_f64();
            let rate = count as f64 / elapsed;
            info!(
                "Progress: {} documents indexed ({:.0} docs/sec)",
                count, rate
            );
        }
    }

    writer.commit().await?;

    // Wait for any in-flight background merges to complete before returning
    // Otherwise they will be cancelled when the IndexWriter/runtime is dropped
    info!("Waiting for background merges to complete...");
    writer.wait_for_merging_thread().await;

    release_memory_to_os();

    let elapsed = start_time.elapsed();
    let rate = count as f64 / elapsed.as_secs_f64();

    if errors > 0 {
        tracing::warn!("Skipped {} documents due to parse errors", errors);
    }

    info!(
        "Indexed {} documents in {:.2}s ({:.0} docs/sec)",
        count,
        elapsed.as_secs_f64(),
        rate
    );

    Ok(count)
}

#[allow(clippy::too_many_arguments)]
pub async fn index_documents(
    index_path: PathBuf,
    documents_path: Option<PathBuf>,
    use_stdin: bool,
    progress_interval: usize,
    max_indexing_memory_mb: usize,
    indexing_threads: Option<usize>,
    compression_threads: Option<usize>,
    optimization: hermes_core::structures::IndexOptimization,
) -> Result<()> {
    let optimization_mode = optimization;

    let dir = FsDirectory::new(&index_path);
    let default_config = IndexConfig::default();
    let config = IndexConfig {
        max_indexing_memory_bytes: max_indexing_memory_mb * 1024 * 1024,
        num_indexing_threads: indexing_threads.unwrap_or(default_config.num_indexing_threads),
        num_compression_threads: compression_threads
            .unwrap_or(default_config.num_compression_threads),
        optimization: optimization_mode,
        ..default_config
    };
    let mut writer = IndexWriter::open(dir, config.clone()).await?;

    info!("Opened index at {:?}", index_path);
    info!(
        "Schema fields: {:?}",
        writer
            .schema()
            .fields()
            .map(|(_, e)| &e.name)
            .collect::<Vec<_>>()
    );
    info!(
        "Indexing threads: {}, Compression threads: {}, Optimization: {:?} (zstd level {})",
        config.num_indexing_threads,
        config.num_compression_threads,
        optimization_mode,
        optimization_mode.zstd_level()
    );

    let count = if use_stdin {
        info!("Reading documents from stdin...");
        let stdin = io::stdin();
        let reader = stdin.lock();
        index_from_reader(&mut writer, reader, progress_interval).await?
    } else if let Some(path) = documents_path {
        info!("Reading documents from {:?}", path);
        let file = File::open(&path)
            .with_context(|| format!("Failed to open documents file: {:?}", path))?;
        let reader = BufReader::new(file);
        index_from_reader(&mut writer, reader, progress_interval).await?
    } else {
        anyhow::bail!("Either --documents or --stdin must be specified");
    };

    info!("Successfully indexed {} documents", count);
    Ok(())
}

pub async fn commit_index(index_path: PathBuf) -> Result<()> {
    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();
    let mut writer = IndexWriter::open(dir, config).await?;

    writer.commit().await?;
    info!("Committed index at {:?}", index_path);

    Ok(())
}

pub async fn merge_index(index_path: PathBuf) -> Result<()> {
    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();
    let mut writer = IndexWriter::open(dir, config).await?;

    info!("Starting force merge...");
    writer.force_merge().await?;
    info!("Force merge completed");

    Ok(())
}

pub async fn reorder_index(index_path: PathBuf) -> Result<()> {
    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();
    let mut writer = IndexWriter::open(dir, config).await?;

    info!("Starting BP reorder...");
    writer.reorder().await?;
    info!("Reorder completed");

    Ok(())
}

pub async fn search_index(
    index_path: PathBuf,
    query_str: &str,
    limit: usize,
    offset: usize,
) -> Result<()> {
    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();
    let index = hermes_core::Index::open(dir, config).await?;
    let schema = index.schema().clone();

    let response = index
        .query_offset(query_str, limit, offset)
        .await
        .with_context(|| format!("Search failed for query: {}", query_str))?;

    info!(
        "Found {} results (total: {})",
        response.hits.len(),
        response.total_hits
    );

    for (i, hit) in response.hits.iter().enumerate() {
        println!(
            "--- Result {} (score: {:.4}) ---",
            offset + i + 1,
            hit.score
        );
        if let Some(doc) = index.get_document(&hit.address).await? {
            let json = doc.to_json(&schema);
            println!("{}", serde_json::to_string_pretty(&json)?);
        }
    }

    println!("---");
    println!(
        "Showing {}-{} of {} results",
        offset + 1,
        offset + response.hits.len(),
        response.total_hits
    );

    Ok(())
}

pub async fn show_info(index_path: PathBuf) -> Result<()> {
    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();
    let index = hermes_core::Index::open(dir, config).await?;

    println!("Index: {:?}", index_path);
    println!("Documents: {}", index.num_docs().await?);
    let segments = index.segment_readers().await?;
    println!("Segments: {}", segments.len());
    println!();
    println!("Schema:");
    for (_field, entry) in index.schema().fields() {
        println!(
            "  {} ({:?}) - indexed: {}, stored: {}",
            entry.name, entry.field_type, entry.indexed, entry.stored
        );
    }

    // Per-field sparse vector statistics (postings, avg vector length)
    let mut sparse_vectors: HashMap<u32, u64> = HashMap::new();
    let mut sparse_postings: HashMap<u32, u64> = HashMap::new();
    for segment in &segments {
        for (&field_id, idx) in segment.sparse_indexes() {
            *sparse_vectors.entry(field_id).or_default() += idx.total_vectors as u64;
            *sparse_postings.entry(field_id).or_default() += idx.total_postings();
        }
        for (&field_id, idx) in segment.bmp_indexes() {
            *sparse_vectors.entry(field_id).or_default() += idx.total_vectors as u64;
            *sparse_postings.entry(field_id).or_default() += idx.total_postings();
        }
    }
    if !sparse_vectors.is_empty() {
        println!();
        println!("Sparse vectors:");
        let mut fields: Vec<_> = sparse_vectors.keys().copied().collect();
        fields.sort_unstable();
        for field_id in fields {
            let name = index
                .schema()
                .get_field_name(hermes_core::dsl::Field(field_id))
                .unwrap_or("unknown");
            let vectors = sparse_vectors[&field_id];
            let postings = sparse_postings.get(&field_id).copied().unwrap_or(0);
            let avg = if vectors > 0 {
                postings as f64 / vectors as f64
            } else {
                0.0
            };
            println!(
                "  {}: {} vectors, {} postings, avg length {:.1} terms/vector",
                name, vectors, postings, avg
            );
        }
    }

    Ok(())
}

pub async fn heatmap_bmp_grid(
    index_path: PathBuf,
    field_name: Option<String>,
    width: Option<usize>,
    height: Option<usize>,
    segment_idx: usize,
) -> Result<()> {
    use hermes_core::{FieldType, Index};

    let dir = FsDirectory::new(&index_path);
    let config = IndexConfig::default();
    let index = Index::open(dir, config).await?;
    let schema = index.schema().clone();
    let segments = index.segment_readers().await?;

    anyhow::ensure!(!segments.is_empty(), "Index has no segments");
    anyhow::ensure!(
        segment_idx < segments.len(),
        "Segment index {} out of range (have {} segments)",
        segment_idx,
        segments.len()
    );

    let segment = &segments[segment_idx];

    // Find BMP field
    let (field, field_name) = if let Some(name) = &field_name {
        let f = schema
            .get_field(name)
            .ok_or_else(|| anyhow::anyhow!("Field '{}' not found in schema", name))?;
        (f, name.clone())
    } else {
        // Auto-detect first BMP field
        schema
            .fields()
            .find(|(_, entry)| {
                entry.field_type == FieldType::SparseVector
                    && entry
                        .sparse_vector_config
                        .as_ref()
                        .is_some_and(|cfg| cfg.format == hermes_core::structures::SparseFormat::Bmp)
            })
            .map(|(f, entry)| (f, entry.name.clone()))
            .ok_or_else(|| anyhow::anyhow!("No BMP sparse vector field found in schema"))?
    };

    let bmp = segment.bmp_index(field).ok_or_else(|| {
        anyhow::anyhow!(
            "No BMP index for field '{}' in segment {}",
            field_name,
            segment_idx
        )
    })?;

    let dims = bmp.dims() as usize;
    let num_blocks = bmp.num_blocks as usize;
    let packed_row_size = bmp.packed_row_size();
    let grid = bmp.grid_slice();

    // Get terminal size
    let (term_cols, term_rows) = terminal_size();
    let out_cols = width.unwrap_or(term_cols.saturating_sub(8)).max(10);
    // ×2 for half-block packing (each char row = 2 pixel rows)
    let out_pixel_rows = height.unwrap_or(term_rows.saturating_sub(6)).max(4) * 2;

    let dim_bin_size = dims.div_ceil(out_pixel_rows).max(1);
    let block_bin_size = num_blocks.div_ceil(out_cols).max(1);
    let actual_rows = dims.div_ceil(dim_bin_size);
    let actual_cols = num_blocks.div_ceil(block_bin_size);

    // Downsample grid via max aggregation
    let mut heatmap = vec![vec![0u8; actual_cols]; actual_rows];
    let mut nonzero_cells = 0u64;
    let mut total_cells = 0u64;

    for d in 0..dims {
        let row = (d / dim_bin_size).min(actual_rows - 1);
        let row_base = d * packed_row_size;
        for b in 0..num_blocks {
            let col = (b / block_bin_size).min(actual_cols - 1);
            let byte = grid[row_base + b / 2];
            let val = if b % 2 == 0 { byte & 0x0F } else { byte >> 4 };
            if val > 0 {
                nonzero_cells += 1;
            }
            total_cells += 1;
            heatmap[row][col] = heatmap[row][col].max(val);
        }
    }

    let sparsity = if total_cells > 0 {
        100.0 * (1.0 - nonzero_cells as f64 / total_cells as f64)
    } else {
        100.0
    };

    // Print header
    println!(
        "BMP Grid Heatmap: field={}, segment={}",
        field_name, segment_idx
    );
    println!(
        "  dims={}, blocks={}, real_docs={}, total_terms={}, total_postings={}",
        dims,
        num_blocks,
        bmp.num_real_docs(),
        bmp.total_terms(),
        bmp.total_postings()
    );
    println!(
        "  grid: {}x{} -> {}x{} (bin: {}d x {}b), sparsity: {:.1}%",
        dims, num_blocks, actual_rows, actual_cols, dim_bin_size, block_bin_size, sparsity
    );
    println!();

    // Render heatmap using Unicode half-blocks
    for row_pair in (0..actual_rows).step_by(2) {
        // Dim axis label
        let dim_start = row_pair * dim_bin_size;
        print!("{:>6} ", dim_start);

        let top_row = &heatmap[row_pair];
        let bot_row = if row_pair + 1 < actual_rows {
            Some(&heatmap[row_pair + 1])
        } else {
            None
        };
        for (col, &top) in top_row.iter().enumerate() {
            let bot = bot_row.map_or(0, |r| r[col]);
            let fg = nibble_to_color(top);
            let bg = nibble_to_color(bot);
            print!("\x1b[38;5;{}m\x1b[48;5;{}m\u{2580}", fg, bg);
        }
        println!("\x1b[0m");
    }

    // Block axis labels
    print!("       ");
    let label_step = actual_cols / 5;
    if label_step > 0 {
        for i in 0..5 {
            let block_start = i * label_step * block_bin_size;
            let padding = label_step;
            print!("{:<width$}", block_start, width = padding);
        }
    }
    println!();

    // Color scale legend
    print!("       ");
    for v in 0..=15u8 {
        let c = nibble_to_color(v);
        print!("\x1b[38;5;{}m\u{2588}", c);
    }
    println!("\x1b[0m  0..............15");

    Ok(())
}

/// Map a 4-bit nibble value (0-15) to an ANSI 256-color index.
fn nibble_to_color(val: u8) -> u8 {
    match val {
        0 => 232, // near-black
        1 => 17,  // dark blue
        2 => 18,
        3 => 19,
        4 => 20, // blue
        5 => 27,
        6 => 33, // cyan
        7 => 39,
        8 => 44,
        9 => 40, // green
        10 => 46,
        11 => 118, // yellow-green
        12 => 226, // yellow
        13 => 220,
        14 => 196, // red
        _ => 231,  // white
    }
}

/// Detect terminal size. Tries `stty size`, then COLUMNS/LINES env vars, falls back to 80x24.
fn terminal_size() -> (usize, usize) {
    // Try `stty size` which returns "rows cols"
    if let Ok(output) = std::process::Command::new("stty")
        .arg("size")
        .arg("-F")
        .arg("/dev/tty")
        .output()
        && let Ok(s) = std::str::from_utf8(&output.stdout)
    {
        let parts: Vec<&str> = s.split_whitespace().collect();
        if parts.len() == 2
            && let (Ok(rows), Ok(cols)) = (parts[0].parse(), parts[1].parse())
        {
            return (cols, rows);
        }
    }
    // Fallback to env vars
    let cols = std::env::var("COLUMNS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(80);
    let rows = std::env::var("LINES")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(24);
    (cols, rows)
}

pub async fn warmup_cache(index_path: PathBuf, cache_size: usize) -> Result<()> {
    use hermes_core::{DirectoryWriter, SLICE_CACHE_FILENAME, SliceCachingDirectory};

    info!(
        "Opening index with slice caching (max {} bytes)...",
        cache_size
    );

    let dir = FsDirectory::new(&index_path);
    let caching_dir = SliceCachingDirectory::new(dir.clone(), cache_size);
    let config = IndexConfig::default();

    let index = hermes_core::Index::open(caching_dir, config).await?;

    info!(
        "Index opened: {} documents, {} segments",
        index.num_docs().await?,
        index.segment_readers().await?.len()
    );

    let stats = index.directory().stats();
    info!(
        "Cache populated: {} bytes in {} slices across {} files",
        stats.total_bytes, stats.total_slices, stats.files_cached
    );

    // Serialize cache data
    let cache_data = index.directory().serialize();
    let cache_file = index_path.join(SLICE_CACHE_FILENAME);
    dir.write(cache_file.as_path(), &cache_data).await?;

    let cache_file_size = std::fs::metadata(&cache_file).map(|m| m.len()).unwrap_or(0);

    info!(
        "Slice cache saved to {:?} ({} bytes)",
        cache_file, cache_file_size
    );

    Ok(())
}