ImageApi/src/bin/import_search_history.rs

use anyhow::{Context, Result};
use chrono::Utc;
use clap::Parser;
use image_api::ai::ollama::OllamaClient;
use image_api::bin_progress;
use image_api::database::search_dao::{InsertSearchRecord, SqliteSearchHistoryDao};
use image_api::parsers::search_html_parser::parse_search_html;
use log::{error, info};

// Import the trait to use its methods
use image_api::database::SearchHistoryDao;

#[derive(Parser, Debug)]
#[command(author, version, about = "Import Google Takeout Search History data", long_about = None)]
struct Args {
    /// Path to the search history HTML file
    #[arg(short, long)]
    path: String,

    /// Skip searches that already exist in the database
    #[arg(long, default_value = "true")]
    skip_existing: bool,

    /// Batch size for embedding generation (max 128 recommended)
    #[arg(long, default_value = "64")]
    batch_size: usize,
}

#[tokio::main]
async fn main() -> Result<()> {
    dotenv::dotenv().ok();
    env_logger::init();

    let args = Args::parse();

    info!("Parsing search history file: {}", args.path);
    let searches = parse_search_html(&args.path).context("Failed to parse search history HTML")?;

    info!("Found {} search records", searches.len());

    let primary_url = dotenv::var("OLLAMA_PRIMARY_URL")
        .or_else(|_| dotenv::var("OLLAMA_URL"))
        .unwrap_or_else(|_| "http://localhost:11434".to_string());
    let fallback_url = dotenv::var("OLLAMA_FALLBACK_URL").ok();
    let primary_model = dotenv::var("OLLAMA_PRIMARY_MODEL")
        .or_else(|_| dotenv::var("OLLAMA_MODEL"))
        .unwrap_or_else(|_| "nomic-embed-text:v1.5".to_string());
    let fallback_model = dotenv::var("OLLAMA_FALLBACK_MODEL").ok();

    let ollama = OllamaClient::new(primary_url, fallback_url, primary_model, fallback_model);
    let context = opentelemetry::Context::current();

    let mut inserted_count = 0usize;
    let mut skipped_count = 0usize;
    let mut error_count = 0usize;

    let mut dao_instance = SqliteSearchHistoryDao::new();
    let created_at = Utc::now().timestamp();

    let pb = bin_progress::determinate(searches.len() as u64, "importing");
    let total_batches = searches.len().div_ceil(args.batch_size);

    // Process searches in batches (embeddings are REQUIRED for searches)
    for (batch_idx, chunk) in searches.chunks(args.batch_size).enumerate() {
        // Generate embeddings for this batch
        let queries: Vec<String> = chunk.iter().map(|s| s.query.clone()).collect();

        let pb_for_warn = pb.clone();
        let embeddings_result = tokio::task::spawn({
            let ollama_client = ollama.clone();
            async move {
                // Generate embeddings in parallel for the batch
                let mut embeddings = Vec::new();
                for query in &queries {
                    match ollama_client.generate_embedding(query).await {
                        Ok(emb) => embeddings.push(Some(emb)),
                        Err(e) => {
                            pb_for_warn.println(format!("embedding failed for '{}': {}", query, e));
                            embeddings.push(None);
                        }
                    }
                }
                embeddings
            }
        })
        .await
        .context("Failed to generate embeddings for batch")?;

        // Build batch of searches with embeddings
        let mut batch_inserts = Vec::new();

        for (search, embedding_opt) in chunk.iter().zip(embeddings_result.iter()) {
            // Check if search exists (optional for speed)
            if args.skip_existing
                && let Ok(exists) =
                    dao_instance.search_exists(&context, search.timestamp, &search.query)
                && exists
            {
                skipped_count += 1;
                continue;
            }

            // Only insert if we have an embedding
            if let Some(embedding) = embedding_opt {
                batch_inserts.push(InsertSearchRecord {
                    timestamp: search.timestamp,
                    query: search.query.clone(),
                    search_engine: search.search_engine.clone(),
                    embedding: embedding.clone(),
                    created_at,
                    source_file: Some(args.path.clone()),
                });
            } else {
                pb.println(format!("skipping '{}' — missing embedding", search.query));
                error_count += 1;
            }
        }

        // Batch insert entire chunk in single transaction
        if !batch_inserts.is_empty() {
            match dao_instance.store_searches_batch(&context, batch_inserts) {
                Ok(count) => inserted_count += count,
                Err(e) => {
                    pb.println(format!("batch insert failed: {:?}", e));
                    error_count += chunk.len();
                }
            }
        }

        pb.set_message(format!(
            "inserted={} skipped={} errors={}",
            inserted_count, skipped_count, error_count
        ));
        pb.inc(chunk.len() as u64);

        // Rate limiting between batches
        if batch_idx + 1 < total_batches {
            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
        }
    }

    pb.finish_and_clear();

    info!("=== Import Summary ===");
    info!("Total searches found: {}", searches.len());
    info!("Successfully inserted: {}", inserted_count);
    info!("Skipped (already exist): {}", skipped_count);
    info!("Errors: {}", error_count);
    info!("All imported searches have embeddings for semantic search");

    if error_count > 0 {
        error!(
            "Completed with {} errors — review log output above",
            error_count
        );
    }

    Ok(())
}