ImageApi/src/state.rs

use crate::ai::insight_chat::{ChatLockMap, InsightChatService};
use crate::ai::openrouter::OpenRouterClient;
use crate::ai::{InsightGenerator, OllamaClient, SmsApiClient};
use crate::database::{
    CalendarEventDao, DailySummaryDao, ExifDao, InsightDao, KnowledgeDao, LocationHistoryDao,
    SearchHistoryDao, SqliteCalendarEventDao, SqliteDailySummaryDao, SqliteExifDao,
    SqliteInsightDao, SqliteKnowledgeDao, SqliteLocationHistoryDao, SqliteSearchHistoryDao,
    connect,
};
use crate::database::{PreviewDao, SqlitePreviewDao};
use crate::libraries::{self, Library};
use crate::tags::{SqliteTagDao, TagDao};
use crate::video::actors::{
    PlaylistGenerator, PreviewClipGenerator, StreamActor, VideoPlaylistManager,
};
use actix::{Actor, Addr};
use std::env;
use std::sync::{Arc, Mutex};

pub struct AppState {
    pub stream_manager: Arc<Addr<StreamActor>>,
    pub playlist_manager: Arc<Addr<VideoPlaylistManager>>,
    pub preview_clip_generator: Arc<Addr<PreviewClipGenerator>>,
    /// All configured media libraries. Ordered by `id` ascending; the first
    /// entry is the primary library.
    pub libraries: Vec<Library>,
    /// Legacy shim equal to `libraries[0].root_path`. Phase 2 transitional —
    /// new code should go through `primary_library()`.
    pub base_path: String,
    pub thumbnail_path: String,
    pub video_path: String,
    pub gif_path: String,
    pub preview_clips_path: String,
    pub excluded_dirs: Vec<String>,
    pub ollama: OllamaClient,
    /// `None` when `OPENROUTER_API_KEY` is not configured. Consulted only
    /// when a request explicitly opts into `backend=hybrid`. Currently
    /// reached via `insight_generator`; kept here so future handlers
    /// (insight_chat) can route to it without threading it through the
    /// generator.
    #[allow(dead_code)]
    pub openrouter: Option<Arc<OpenRouterClient>>,
    /// Curated list of OpenRouter model ids exposed to clients. Sourced from
    /// `OPENROUTER_ALLOWED_MODELS` (comma-separated). Empty when unset.
    pub openrouter_allowed_models: Vec<String>,
    pub sms_client: SmsApiClient,
    pub insight_generator: InsightGenerator,
    /// Chat continuation service. Hold an Arc so handlers can clone cheaply.
    pub insight_chat: Arc<InsightChatService>,
}

impl AppState {
    pub fn primary_library(&self) -> &Library {
        self.libraries
            .first()
            .expect("AppState constructed without any libraries")
    }

    pub fn library_by_id(&self, id: i32) -> Option<&Library> {
        self.libraries.iter().find(|l| l.id == id)
    }

    pub fn library_by_name(&self, name: &str) -> Option<&Library> {
        self.libraries.iter().find(|l| l.name == name)
    }
}

impl AppState {
    pub fn new(
        stream_manager: Arc<Addr<StreamActor>>,
        libraries_vec: Vec<Library>,
        thumbnail_path: String,
        video_path: String,
        gif_path: String,
        preview_clips_path: String,
        excluded_dirs: Vec<String>,
        ollama: OllamaClient,
        openrouter: Option<Arc<OpenRouterClient>>,
        openrouter_allowed_models: Vec<String>,
        sms_client: SmsApiClient,
        insight_generator: InsightGenerator,
        insight_chat: Arc<InsightChatService>,
        preview_dao: Arc<Mutex<Box<dyn PreviewDao>>>,
    ) -> Self {
        assert!(
            !libraries_vec.is_empty(),
            "AppState::new requires at least one library"
        );
        let base_path = libraries_vec[0].root_path.clone();
        let playlist_generator = PlaylistGenerator::new();
        let video_playlist_manager =
            VideoPlaylistManager::new(video_path.clone(), playlist_generator.start());

        let preview_clip_generator = PreviewClipGenerator::new(
            preview_clips_path.clone(),
            libraries_vec.clone(),
            preview_dao,
        );

        Self {
            stream_manager,
            playlist_manager: Arc::new(video_playlist_manager.start()),
            preview_clip_generator: Arc::new(preview_clip_generator.start()),
            libraries: libraries_vec,
            base_path,
            thumbnail_path,
            video_path,
            gif_path,
            preview_clips_path,
            excluded_dirs,
            ollama,
            openrouter,
            openrouter_allowed_models,
            sms_client,
            insight_generator,
            insight_chat,
        }
    }

    /// Parse excluded directories from environment variable
    fn parse_excluded_dirs() -> Vec<String> {
        env::var("EXCLUDED_DIRS")
            .unwrap_or_default()
            .split(',')
            .filter(|dir| !dir.trim().is_empty())
            .map(|dir| dir.trim().to_string())
            .collect()
    }
}

impl Default for AppState {
    fn default() -> Self {
        // Initialize AI clients
        let ollama_primary_url = env::var("OLLAMA_PRIMARY_URL").unwrap_or_else(|_| {
            env::var("OLLAMA_URL").unwrap_or_else(|_| "http://localhost:11434".to_string())
        });
        let ollama_fallback_url = env::var("OLLAMA_FALLBACK_URL").ok();
        let ollama_primary_model = env::var("OLLAMA_PRIMARY_MODEL")
            .or_else(|_| env::var("OLLAMA_MODEL"))
            .unwrap_or_else(|_| "nemotron-3-nano:30b".to_string());
        let ollama_fallback_model = env::var("OLLAMA_FALLBACK_MODEL").ok();

        let ollama = OllamaClient::new(
            ollama_primary_url,
            ollama_fallback_url,
            ollama_primary_model,
            ollama_fallback_model,
        );

        let openrouter = build_openrouter_from_env();
        let openrouter_allowed_models = parse_openrouter_allowed_models();

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

        // Initialize 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 preview_dao: Arc<Mutex<Box<dyn PreviewDao>>> =
            Arc::new(Mutex::new(Box::new(SqlitePreviewDao::new())));

        // Initialize Google Takeout DAOs
        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())));

        // Load base path and ensure the primary library row reflects it.
        let base_path = env::var("BASE_PATH").expect("BASE_PATH was not set in the env");
        let mut seed_conn = connect();
        libraries::seed_or_patch_from_env(&mut seed_conn, &base_path);
        let libraries_vec = libraries::load_all(&mut seed_conn);
        assert!(
            !libraries_vec.is_empty(),
            "libraries table is empty after seed_or_patch_from_env"
        );
        drop(seed_conn);

        // Initialize InsightGenerator with all data sources
        let insight_generator = InsightGenerator::new(
            ollama.clone(),
            openrouter.clone(),
            sms_client.clone(),
            insight_dao.clone(),
            exif_dao.clone(),
            daily_summary_dao.clone(),
            calendar_dao.clone(),
            location_dao.clone(),
            search_dao.clone(),
            tag_dao.clone(),
            knowledge_dao,
            libraries_vec.clone(),
        );

        // Chat continuation reuses the generator for tool dispatch + image
        // loading. The lock map starts empty and grows lazily per file.
        let chat_locks: ChatLockMap =
            Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new()));
        let insight_chat = Arc::new(InsightChatService::new(
            Arc::new(insight_generator.clone()),
            ollama.clone(),
            openrouter.clone(),
            insight_dao.clone(),
            chat_locks,
        ));

        // Ensure preview clips directory exists
        let preview_clips_path =
            env::var("PREVIEW_CLIPS_DIRECTORY").unwrap_or_else(|_| "preview_clips".to_string());
        std::fs::create_dir_all(&preview_clips_path)
            .expect("Failed to create PREVIEW_CLIPS_DIRECTORY");

        Self::new(
            Arc::new(StreamActor {}.start()),
            libraries_vec,
            env::var("THUMBNAILS").expect("THUMBNAILS was not set in the env"),
            env::var("VIDEO_PATH").expect("VIDEO_PATH was not set in the env"),
            env::var("GIFS_DIRECTORY").expect("GIFS_DIRECTORY was not set in the env"),
            preview_clips_path,
            Self::parse_excluded_dirs(),
            ollama,
            openrouter,
            openrouter_allowed_models,
            sms_client,
            insight_generator,
            insight_chat,
            preview_dao,
        )
    }
}

/// Build an `OpenRouterClient` from environment variables. Returns `None`
/// when `OPENROUTER_API_KEY` is unset (the hybrid backend is then
/// unavailable and requests for it return a clear error).
fn build_openrouter_from_env() -> Option<Arc<OpenRouterClient>> {
    let api_key = env::var("OPENROUTER_API_KEY").ok()?;
    let base_url = env::var("OPENROUTER_BASE_URL").ok();
    let default_model = env::var("OPENROUTER_DEFAULT_MODEL")
        .unwrap_or_else(|_| "anthropic/claude-sonnet-4".to_string());
    let mut client = OpenRouterClient::new(api_key, base_url, default_model);
    client.set_attribution(
        env::var("OPENROUTER_HTTP_REFERER").ok(),
        env::var("OPENROUTER_APP_TITLE").ok(),
    );
    if let Ok(model) = env::var("OPENROUTER_EMBEDDING_MODEL") {
        client.set_embedding_model(model);
    }
    Some(Arc::new(client))
}

/// Parse `OPENROUTER_ALLOWED_MODELS` (comma-separated) into a vec. Returns
/// empty when unset, in which case `/insights/openrouter/models` reports no
/// curated picks and the server falls back to `OPENROUTER_DEFAULT_MODEL`.
fn parse_openrouter_allowed_models() -> Vec<String> {
    env::var("OPENROUTER_ALLOWED_MODELS")
        .unwrap_or_default()
        .split(',')
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect()
}

#[cfg(test)]
impl AppState {
    /// Creates an AppState instance for testing with temporary directories
    pub fn test_state() -> Self {
        use actix::Actor;
        // Create a base temporary directory
        let temp_dir = tempfile::tempdir().expect("Failed to create temp directory");
        let base_path = temp_dir.path().to_path_buf();

        // Create subdirectories for thumbnails, videos, gifs, and preview clips
        let thumbnail_path = create_test_subdir(&base_path, "thumbnails");
        let video_path = create_test_subdir(&base_path, "videos");
        let gif_path = create_test_subdir(&base_path, "gifs");
        let preview_clips_path = create_test_subdir(&base_path, "preview_clips");

        // Initialize test AI clients
        let ollama = OllamaClient::new(
            "http://localhost:11434".to_string(),
            None,
            "llama3.2".to_string(),
            None,
        );
        let sms_client = SmsApiClient::new("http://localhost:8000".to_string(), None);

        // Initialize test 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())));

        // Initialize test Google Takeout DAOs
        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())));

        // Initialize test InsightGenerator with all data sources
        let base_path_str = base_path.to_string_lossy().to_string();
        let test_lib = Library {
            id: crate::libraries::PRIMARY_LIBRARY_ID,
            name: "main".to_string(),
            root_path: base_path_str.clone(),
        };
        let insight_generator = InsightGenerator::new(
            ollama.clone(),
            None,
            sms_client.clone(),
            insight_dao.clone(),
            exif_dao.clone(),
            daily_summary_dao.clone(),
            calendar_dao.clone(),
            location_dao.clone(),
            search_dao.clone(),
            tag_dao.clone(),
            knowledge_dao,
            vec![test_lib],
        );

        let chat_locks: ChatLockMap =
            Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new()));
        let insight_chat = Arc::new(InsightChatService::new(
            Arc::new(insight_generator.clone()),
            ollama.clone(),
            None,
            insight_dao.clone(),
            chat_locks,
        ));

        // Initialize test preview DAO
        let preview_dao: Arc<Mutex<Box<dyn PreviewDao>>> =
            Arc::new(Mutex::new(Box::new(SqlitePreviewDao::new())));

        // Create the AppState with the temporary paths
        let test_libraries = vec![Library {
            id: crate::libraries::PRIMARY_LIBRARY_ID,
            name: "main".to_string(),
            root_path: base_path_str.clone(),
        }];
        AppState::new(
            Arc::new(StreamActor {}.start()),
            test_libraries,
            thumbnail_path.to_string_lossy().to_string(),
            video_path.to_string_lossy().to_string(),
            gif_path.to_string_lossy().to_string(),
            preview_clips_path.to_string_lossy().to_string(),
            Vec::new(), // No excluded directories for test state
            ollama,
            None,
            Vec::new(),
            sms_client,
            insight_generator,
            insight_chat,
            preview_dao,
        )
    }
}

/// Helper function to create a subdirectory inside the base directory for testing
#[cfg(test)]
fn create_test_subdir(base_path: &std::path::Path, name: &str) -> std::path::PathBuf {
    let dir_path = base_path.join(name);
    std::fs::create_dir_all(&dir_path)
        .unwrap_or_else(|_| panic!("Failed to create {} directory", name));
    dir_path
}