ImageApi/src/bin/populate_knowledge.rs

use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use clap::Parser;
use walkdir::WalkDir;

use image_api::ai::{InsightGenerator, OllamaClient, SmsApiClient};
use image_api::database::{
    CalendarEventDao, DailySummaryDao, ExifDao, InsightDao, KnowledgeDao, LocationHistoryDao,
    SearchHistoryDao, SqliteCalendarEventDao, SqliteDailySummaryDao, SqliteExifDao,
    SqliteInsightDao, SqliteKnowledgeDao, SqliteLocationHistoryDao, SqliteSearchHistoryDao,
};
use image_api::file_types::{IMAGE_EXTENSIONS, VIDEO_EXTENSIONS};
use image_api::libraries::{self, Library};
use image_api::tags::{SqliteTagDao, TagDao};

#[derive(Parser, Debug)]
#[command(name = "populate_knowledge")]
#[command(
    about = "Batch populate the knowledge base by running the agentic insight loop over a folder"
)]
struct Args {
    /// Directory to scan. Defaults to BASE_PATH from .env
    #[arg(long)]
    path: Option<String>,

    /// Ollama model override. Defaults to OLLAMA_PRIMARY_MODEL from .env
    #[arg(long)]
    model: Option<String>,

    /// Maximum agentic loop iterations per file
    #[arg(long, default_value_t = 12)]
    max_iterations: usize,

    /// HTTP request timeout in seconds. Increase for large/slow models
    #[arg(long, default_value_t = 120)]
    timeout_secs: u64,

    /// Context window size (num_ctx) passed to the model
    #[arg(long)]
    num_ctx: Option<i32>,

    /// Sampling temperature (e.g. 0.8). Omit for model default
    #[arg(long)]
    temperature: Option<f32>,

    /// Top-p (nucleus) sampling (e.g. 0.9). Omit for model default
    #[arg(long)]
    top_p: Option<f32>,

    /// Top-k sampling (e.g. 40). Omit for model default
    #[arg(long)]
    top_k: Option<i32>,

    /// Min-p sampling (e.g. 0.05). Omit for model default
    #[arg(long)]
    min_p: Option<f32>,

    /// Re-process files that already have an insight stored
    #[arg(long, default_value_t = false)]
    reprocess: bool,
}

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

    let args = Args::parse();

    let base_path = dotenv::var("BASE_PATH")?;
    let scan_path = args.path.as_deref().unwrap_or(&base_path).to_string();

    // Ollama config from env with CLI overrides
    let primary_url = std::env::var("OLLAMA_PRIMARY_URL")
        .or_else(|_| std::env::var("OLLAMA_URL"))
        .unwrap_or_else(|_| "http://localhost:11434".to_string());
    let fallback_url = std::env::var("OLLAMA_FALLBACK_URL").ok();
    let primary_model = args
        .model
        .clone()
        .or_else(|| std::env::var("OLLAMA_PRIMARY_MODEL").ok())
        .or_else(|| std::env::var("OLLAMA_MODEL").ok())
        .unwrap_or_else(|| "nemotron-3-nano:30b".to_string());
    let fallback_model = std::env::var("OLLAMA_FALLBACK_MODEL").ok();

    let mut ollama = OllamaClient::new(
        primary_url.clone(),
        fallback_url,
        primary_model.clone(),
        fallback_model,
    )
    .with_request_timeout(args.timeout_secs);

    if let Some(ctx) = args.num_ctx {
        ollama.set_num_ctx(Some(ctx));
    }
    if args.temperature.is_some()
        || args.top_p.is_some()
        || args.top_k.is_some()
        || args.min_p.is_some()
    {
        ollama.set_sampling_params(args.temperature, args.top_p, args.top_k, args.min_p);
    }

    let sms_api_url =
        std::env::var("SMS_API_URL").unwrap_or_else(|_| "http://localhost:8000".to_string());
    let sms_api_token = std::env::var("SMS_API_TOKEN").ok();
    let sms_client = SmsApiClient::new(sms_api_url, sms_api_token);

    // Wire up all DAOs
    let insight_dao: Arc<Mutex<Box<dyn InsightDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteInsightDao::new())));
    let exif_dao: Arc<Mutex<Box<dyn ExifDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteExifDao::new())));
    let daily_summary_dao: Arc<Mutex<Box<dyn DailySummaryDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteDailySummaryDao::new())));
    let calendar_dao: Arc<Mutex<Box<dyn CalendarEventDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteCalendarEventDao::new())));
    let location_dao: Arc<Mutex<Box<dyn LocationHistoryDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteLocationHistoryDao::new())));
    let search_dao: Arc<Mutex<Box<dyn SearchHistoryDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteSearchHistoryDao::new())));
    let tag_dao: Arc<Mutex<Box<dyn TagDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteTagDao::default())));
    let knowledge_dao: Arc<Mutex<Box<dyn KnowledgeDao>>> =
        Arc::new(Mutex::new(Box::new(SqliteKnowledgeDao::new())));

    let populate_lib = Library {
        id: libraries::PRIMARY_LIBRARY_ID,
        name: "main".to_string(),
        root_path: base_path.clone(),
    };

    let generator = InsightGenerator::new(
        ollama,
        sms_client,
        insight_dao.clone(),
        exif_dao,
        daily_summary_dao,
        calendar_dao,
        location_dao,
        search_dao,
        tag_dao,
        knowledge_dao,
        vec![populate_lib],
    );

    println!("Knowledge Base Population");
    println!("=========================");
    println!("Scan path:      {}", scan_path);
    println!("Model:          {}", primary_model);
    println!("Max iterations: {}", args.max_iterations);
    println!("Timeout:        {}s", args.timeout_secs);
    if let Some(ctx) = args.num_ctx {
        println!("Num ctx:        {}", ctx);
    }
    if let Some(t) = args.temperature {
        println!("Temperature:    {}", t);
    }
    if let Some(p) = args.top_p {
        println!("Top P:          {}", p);
    }
    if let Some(k) = args.top_k {
        println!("Top K:          {}", k);
    }
    if let Some(m) = args.min_p {
        println!("Min P:          {}", m);
    }
    println!(
        "Mode:           {}",
        if args.reprocess {
            "reprocess all"
        } else {
            "skip existing"
        }
    );
    println!();

    // Collect all image and video files
    let all_extensions: Vec<&str> = IMAGE_EXTENSIONS
        .iter()
        .chain(VIDEO_EXTENSIONS.iter())
        .copied()
        .collect();

    println!("Scanning {}...", scan_path);
    let files: Vec<PathBuf> = WalkDir::new(&scan_path)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|e| e.file_type().is_file())
        .filter(|e| {
            e.path()
                .extension()
                .and_then(|ext| ext.to_str())
                .map(|ext| all_extensions.contains(&ext.to_lowercase().as_str()))
                .unwrap_or(false)
        })
        .map(|e| e.path().to_path_buf())
        .collect();

    let total = files.len();
    println!("Found {} files\n", total);

    if total == 0 {
        println!("Nothing to process.");
        return Ok(());
    }

    let cx = opentelemetry::Context::new();
    let mut processed = 0usize;
    let mut skipped = 0usize;
    let mut errors = 0usize;

    for (i, path) in files.iter().enumerate() {
        let relative = match path.strip_prefix(&base_path) {
            Ok(p) => p.to_string_lossy().replace('\\', "/"),
            Err(_) => path.to_string_lossy().replace('\\', "/"),
        };

        let prefix = format!("[{}/{}]", i + 1, total);

        // Check for existing insight unless --reprocess
        if !args.reprocess {
            let has_insight = insight_dao
                .lock()
                .unwrap()
                .get_insight(&cx, &relative)
                .unwrap_or(None)
                .is_some();

            if has_insight {
                println!("{} skip  {}", prefix, relative);
                skipped += 1;
                continue;
            }
        }

        println!("{} start {}", prefix, relative);

        match generator
            .generate_agentic_insight_for_photo(
                &relative,
                args.model.clone(),
                None,
                args.num_ctx,
                args.temperature,
                args.top_p,
                args.top_k,
                args.min_p,
                args.max_iterations,
            )
            .await
        {
            Ok(_) => {
                println!("{} done  {}", prefix, relative);
                processed += 1;
            }
            Err(e) => {
                eprintln!("{} error {} — {:?}", prefix, relative, e);
                errors += 1;
            }
        }
    }

    println!();
    println!("=========================");
    println!("Complete");
    println!("  Processed: {}", processed);
    println!("  Skipped:   {}", skipped);
    println!("  Errors:    {}", errors);

    Ok(())
}