neiam /dms

use sea_orm::DatabaseConnection;
use std::env;

#[derive(Debug, Clone)]
pub struct Config {
    pub database_url: String,
    pub session_duration_hours: i64,
    pub registration_token_duration_days: i64,
    pub bind_address: String,
    pub storage_root: String,
    pub amqp_url: String,
    pub worker_prefetch: u16,
    pub worker_max_attempts: i32,
    pub max_upload_bytes: usize,
    pub auto_migrate: bool,
    pub registration_enabled: bool,
    pub describe: DescribeConfig,
}

#[derive(Debug, Clone)]
pub enum DescribeBackendKind {
    Noop,
    Rknn,
}

#[derive(Debug, Clone)]
pub struct DescribeConfig {
    pub backend: DescribeBackendKind,
    pub endpoint: String,
    pub model: String,
    pub prompt: String,
    pub timeout_secs: u64,
    pub max_new_tokens: u32,
}

impl Config {
    pub fn from_env() -> Result<Self, env::VarError> {
        Ok(Config {
            database_url: env::var("DATABASE_URL")
                .unwrap_or_else(|_| "postgres://dms_user:dms_password@localhost:5433/dms".to_string()),
            session_duration_hours: env::var("SESSION_DURATION_HOURS")
                .unwrap_or_else(|_| "24".to_string())
                .parse()
                .unwrap_or(24),
            registration_token_duration_days: env::var("REGISTRATION_TOKEN_DURATION_DAYS")
                .unwrap_or_else(|_| "7".to_string())
                .parse()
                .unwrap_or(7),
            bind_address: env::var("BIND_ADDRESS")
                .unwrap_or_else(|_| "0.0.0.0:3007".to_string()),
            storage_root: env::var("STORAGE_ROOT")
                .unwrap_or_else(|_| "./storage".to_string()),
            amqp_url: env::var("AMQP_URL")
                .unwrap_or_else(|_| "amqp://dms_user:dms_password@localhost:5674/%2f".to_string()),
            worker_prefetch: env::var("WORKER_PREFETCH")
                .unwrap_or_else(|_| "4".to_string())
                .parse()
                .unwrap_or(4),
            worker_max_attempts: env::var("WORKER_MAX_ATTEMPTS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(100),
            max_upload_bytes: env::var("MAX_UPLOAD_BYTES")
                .unwrap_or_else(|_| "104857600".to_string()) // 100 MiB
                .parse()
                .unwrap_or(100 * 1024 * 1024),
            auto_migrate: env::var("AUTO_MIGRATE")
                .map(|v| matches!(v.to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on"))
                .unwrap_or(true),
            registration_enabled: env::var("REGISTRATION_ENABLED")
                .map(|v| matches!(v.to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on"))
                .unwrap_or(false),
            describe: DescribeConfig::from_env(),
        })
    }
}

impl DescribeConfig {
    fn from_env() -> Self {
        let backend = match env::var("DESCRIBE_BACKEND")
            .unwrap_or_else(|_| "noop".to_string())
            .to_ascii_lowercase()
            .as_str()
        {
            "rknn" | "rknn-llm" | "rknnllm" => DescribeBackendKind::Rknn,
            _ => DescribeBackendKind::Noop,
        };
        DescribeConfig {
            backend,
            endpoint: env::var("RKNN_LLM_ENDPOINT")
                .unwrap_or_else(|_| "http://r5bp-0:8080/describe".to_string()),
            model: env::var("RKNN_LLM_MODEL").unwrap_or_else(|_| "qwen2-vl-2b".to_string()),
            prompt: env::var("RKNN_LLM_PROMPT").unwrap_or_else(|_| {
                "Describe this image in one or two concise sentences. \
                 List the salient objects, setting, and any visible text."
                    .to_string()
            }),
            timeout_secs: env::var("RKNN_LLM_TIMEOUT_SECS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(120),
            max_new_tokens: env::var("RKNN_LLM_MAX_NEW_TOKENS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(256),
        }
    }
}

pub async fn create_db_connection(database_url: &str) -> Result<DatabaseConnection, sea_orm::DbErr> {
    sea_orm::Database::connect(database_url).await
}