compliance-scanner-agent/compliance-agent/src/scheduler.rs

use mongodb::bson::doc;
use tokio_cron_scheduler::{Job, JobScheduler};

use compliance_core::models::ScanTrigger;

use crate::agent::ComplianceAgent;
use crate::database::Database;
use crate::error::AgentError;

/// Default tenant the scheduler runs against when neither the tenant
/// registry nor `SCHEDULER_TENANT_IDS` are configured. Matches the
/// dev-injector default so a bare `cargo run` has the scheduler
/// scanning whatever lives in `<prefix>_dev`.
const DEFAULT_SCHEDULER_TENANT_ID: &str = "dev";

/// Request timeout when fetching the live tenant list from the
/// registry. Kept short — if the registry is slow we'd rather fall
/// back to env-configured ids and finish the tick than block the
/// scheduler loop.
const REGISTRY_FETCH_TIMEOUT_SECS: u64 = 5;

pub async fn start_scheduler(agent: &ComplianceAgent) -> Result<(), AgentError> {
    let sched = JobScheduler::new()
        .await
        .map_err(|e| AgentError::Scheduler(format!("Failed to create scheduler: {e}")))?;

    // Periodic scan job
    let scan_agent = agent.clone();
    let scan_schedule = agent.config.scan_schedule.clone();
    let scan_job = Job::new_async(scan_schedule.as_str(), move |_uuid, _lock| {
        let agent = scan_agent.clone();
        Box::pin(async move {
            tracing::info!("Scheduled scan triggered");
            let tenants = scheduler_tenants(&agent).await;
            tracing::debug!(
                tenant_count = tenants.len(),
                "Scheduled scan: tenants resolved"
            );
            for tenant_id in tenants {
                scan_all_repos(&agent, &tenant_id).await;
            }
        })
    })
    .map_err(|e| AgentError::Scheduler(format!("Failed to create scan job: {e}")))?;
    sched
        .add(scan_job)
        .await
        .map_err(|e| AgentError::Scheduler(format!("Failed to add scan job: {e}")))?;

    // CVE monitor job (daily)
    let cve_agent = agent.clone();
    let cve_schedule = agent.config.cve_monitor_schedule.clone();
    let cve_job = Job::new_async(cve_schedule.as_str(), move |_uuid, _lock| {
        let agent = cve_agent.clone();
        Box::pin(async move {
            tracing::info!("CVE monitor triggered");
            let tenants = scheduler_tenants(&agent).await;
            tracing::debug!(
                tenant_count = tenants.len(),
                "CVE monitor: tenants resolved"
            );
            for tenant_id in tenants {
                monitor_cves(&agent, &tenant_id).await;
            }
        })
    })
    .map_err(|e| AgentError::Scheduler(format!("Failed to create CVE monitor job: {e}")))?;
    sched
        .add(cve_job)
        .await
        .map_err(|e| AgentError::Scheduler(format!("Failed to add CVE monitor job: {e}")))?;

    sched
        .start()
        .await
        .map_err(|e| AgentError::Scheduler(format!("Failed to start scheduler: {e}")))?;

    let tenants = scheduler_tenants(agent).await;
    let source = if agent.config.tenant_registry_url.is_some() {
        "tenant-registry (env fallback)"
    } else {
        "env (SCHEDULER_TENANT_IDS)"
    };
    tracing::info!(
        "Scheduler started: scans='{}', CVE monitor='{}', tenant source={source}, tenants={tenants:?}",
        agent.config.scan_schedule,
        agent.config.cve_monitor_schedule,
    );

    // Keep scheduler alive
    loop {
        tokio::time::sleep(tokio::time::Duration::from_secs(3600)).await;
    }
}

/// Tenants the scheduler iterates each tick.
///
/// Resolution order:
/// 1. **Tenant registry** at `agent.config.tenant_registry_url`
///    (`GET /v1/tenants`). Fresh on every tick — picks up newly
///    provisioned tenants without an agent restart.
/// 2. **`SCHEDULER_TENANT_IDS`** env (comma-separated) — fallback when
///    the registry is unreachable, the response is malformed, or no
///    registry URL is configured.
/// 3. **`DEFAULT_SCHEDULER_TENANT_ID`** (`"dev"`) — last-ditch fallback
///    so the scheduler keeps doing something useful in dev.
///
/// We never panic out of this function — the scheduler must keep
/// firing even if the registry is offline.
async fn scheduler_tenants(agent: &ComplianceAgent) -> Vec<String> {
    if let Some(url) = agent.config.tenant_registry_url.as_deref() {
        match fetch_tenants_from_registry(&agent.http, url).await {
            Ok(v) if !v.is_empty() => return v,
            Ok(_) => {
                tracing::warn!("tenant-registry returned empty list; falling back to env");
            }
            Err(e) => {
                tracing::warn!(
                    url = %url,
                    error = %e,
                    "tenant-registry fetch failed; falling back to env"
                );
            }
        }
    }
    tenants_from_env()
}

fn tenants_from_env() -> Vec<String> {
    std::env::var("SCHEDULER_TENANT_IDS")
        .ok()
        .map(|s| {
            s.split(',')
                .map(str::trim)
                .filter(|s| !s.is_empty())
                .map(String::from)
                .collect::<Vec<_>>()
        })
        .filter(|v| !v.is_empty())
        .unwrap_or_else(|| vec![DEFAULT_SCHEDULER_TENANT_ID.to_string()])
}

/// Shape we accept from the registry. Liberal in what we accept:
/// the registry can return any field shape as long as either `id` or
/// `tenant_id` is present. Other fields are ignored.
#[derive(serde::Deserialize)]
struct RegistryTenant {
    #[serde(alias = "tenant_id")]
    id: String,
    /// Filter out non-running tenants if status is present. Missing
    /// status defaults to "active" so older registry deployments keep
    /// working.
    #[serde(default = "default_status")]
    status: String,
}

fn default_status() -> String {
    "active".to_string()
}

#[derive(serde::Deserialize)]
struct RegistryListResponse {
    data: Vec<RegistryTenant>,
}

async fn fetch_tenants_from_registry(
    http: &reqwest::Client,
    base_url: &str,
) -> Result<Vec<String>, String> {
    let url = format!("{}/v1/tenants", base_url.trim_end_matches('/'));
    let resp = http
        .get(&url)
        .timeout(std::time::Duration::from_secs(REGISTRY_FETCH_TIMEOUT_SECS))
        .send()
        .await
        .map_err(|e| format!("request failed: {e}"))?;
    if !resp.status().is_success() {
        return Err(format!("registry returned {}", resp.status()));
    }
    let body: RegistryListResponse = resp
        .json()
        .await
        .map_err(|e| format!("invalid JSON: {e}"))?;
    Ok(filter_active(body.data))
}

/// Frozen/Archived tenants don't need scheduled scans; the M7.1
/// status gate would 402/410 anyway. Skip them so we don't waste
/// cycles. Active / trial / demo / anything-else-unknown all run.
fn filter_active(rows: Vec<RegistryTenant>) -> Vec<String> {
    rows.into_iter()
        .filter(|t| !matches!(t.status.as_str(), "frozen" | "archived"))
        .map(|t| t.id)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tenant(id: &str, status: &str) -> RegistryTenant {
        RegistryTenant {
            id: id.to_string(),
            status: status.to_string(),
        }
    }

    #[test]
    fn filter_active_keeps_running_skips_frozen_archived() {
        let rows = vec![
            tenant("a", "active"),
            tenant("b", "trial"),
            tenant("c", "demo"),
            tenant("d", "frozen"),
            tenant("e", "archived"),
            tenant("f", "weird-but-not-known-dead"),
        ];
        let out = filter_active(rows);
        assert_eq!(out, vec!["a", "b", "c", "f"]);
    }

    #[test]
    fn deserialize_registry_response_accepts_id_or_tenant_id() {
        let body = r#"{"data":[
            {"id":"a","status":"active"},
            {"tenant_id":"b","status":"trial"},
            {"id":"c"}
        ]}"#;
        let parsed: RegistryListResponse = serde_json::from_str(body).unwrap();
        assert_eq!(parsed.data.len(), 3);
        assert_eq!(parsed.data[0].id, "a");
        assert_eq!(parsed.data[1].id, "b");
        assert_eq!(parsed.data[2].id, "c");
        // Default status for the third entry should be "active"
        assert_eq!(parsed.data[2].status, "active");
    }

    /// Combined into a single test: cargo runs tests in parallel and
    /// env vars are process-global, so two separate tests touching
    /// `SCHEDULER_TENANT_IDS` race each other. Doing both checks in
    /// one test keeps them in a deterministic order.
    #[test]
    fn tenants_from_env_resolution() {
        std::env::remove_var("SCHEDULER_TENANT_IDS");
        assert_eq!(
            tenants_from_env(),
            vec![DEFAULT_SCHEDULER_TENANT_ID.to_string()],
            "unset → default"
        );

        std::env::set_var("SCHEDULER_TENANT_IDS", "acme, globex ,,hello");
        let out = tenants_from_env();
        std::env::remove_var("SCHEDULER_TENANT_IDS");
        assert_eq!(
            out,
            vec!["acme", "globex", "hello"],
            "splits + trims + drops empty"
        );

        std::env::set_var("SCHEDULER_TENANT_IDS", "");
        let out = tenants_from_env();
        std::env::remove_var("SCHEDULER_TENANT_IDS");
        assert_eq!(
            out,
            vec![DEFAULT_SCHEDULER_TENANT_ID.to_string()],
            "empty → default"
        );
    }
}

/// Resolve the per-tenant database. Logs and returns `None` on failure
/// so the loop in the caller can continue with other tenants.
async fn tenant_db(agent: &ComplianceAgent, tenant_id: &str) -> Option<Database> {
    match agent.db_pool.for_tenant_id(tenant_id).await {
        Ok(db) => Some(db),
        Err(e) => {
            tracing::error!("Scheduler: cannot open tenant database '{tenant_id}': {e}");
            None
        }
    }
}

async fn scan_all_repos(agent: &ComplianceAgent, tenant_id: &str) {
    use futures_util::StreamExt;

    let db = match tenant_db(agent, tenant_id).await {
        Some(db) => db,
        None => return,
    };

    let cursor = match db.onboarded_targets().find(doc! {}).await {
        Ok(c) => c,
        Err(e) => {
            tracing::error!("Failed to list targets for tenant '{tenant_id}': {e}");
            return;
        }
    };

    let targets: Vec<_> = cursor.filter_map(|r| async { r.ok() }).collect().await;

    for target in targets {
        let target_id = target.id.map(|id| id.to_hex()).unwrap_or_default();
        if let Err(e) = agent
            .run_target_scan(tenant_id, &target_id, ScanTrigger::Scheduled)
            .await
        {
            tracing::error!(
                "Scheduled scan failed for {} (tenant '{tenant_id}'): {e}",
                target.name
            );
        }
    }
}

async fn monitor_cves(agent: &ComplianceAgent, tenant_id: &str) {
    use compliance_core::models::notification::{parse_severity, CveNotification};
    use compliance_core::models::SbomEntry;
    use futures_util::StreamExt;

    let db = match tenant_db(agent, tenant_id).await {
        Some(db) => db,
        None => return,
    };

    // Fetch all SBOM entries grouped by repo
    let cursor = match db.sbom_entries().find(doc! {}).await {
        Ok(c) => c,
        Err(e) => {
            tracing::error!("CVE monitor: failed to list SBOM entries for '{tenant_id}': {e}");
            return;
        }
    };
    let entries: Vec<SbomEntry> = cursor.filter_map(|r| async { r.ok() }).collect().await;
    if entries.is_empty() {
        tracing::debug!("CVE monitor: no SBOM entries for tenant '{tenant_id}', skipping");
        return;
    }

    tracing::info!(
        "CVE monitor: checking {} dependencies for new CVEs (tenant '{tenant_id}')",
        entries.len()
    );

    // Build a repo_id → repo_name lookup
    let repo_ids: std::collections::HashSet<String> =
        entries.iter().map(|e| e.repo_id.clone()).collect();
    let mut repo_names: std::collections::HashMap<String, String> =
        std::collections::HashMap::new();
    for rid in &repo_ids {
        if let Ok(oid) = mongodb::bson::oid::ObjectId::parse_str(rid) {
            if let Ok(Some(repo)) = db.onboarded_targets().find_one(doc! { "_id": oid }).await {
                repo_names.insert(rid.clone(), repo.name.clone());
            }
        }
    }

    // Use the existing CveScanner to query OSV.dev
    let nvd_key = agent.config.nvd_api_key.as_ref().map(|k| {
        use secrecy::ExposeSecret;
        k.expose_secret().to_string()
    });
    let scanner = crate::pipeline::cve::CveScanner::new(
        agent.http.clone(),
        agent.config.searxng_url.clone(),
        nvd_key,
    );

    // Group entries by repo for scanning
    let mut entries_by_repo: std::collections::HashMap<String, Vec<SbomEntry>> =
        std::collections::HashMap::new();
    for entry in entries {
        entries_by_repo
            .entry(entry.repo_id.clone())
            .or_default()
            .push(entry);
    }

    let mut new_notifications = 0u32;

    for (repo_id, mut repo_entries) in entries_by_repo {
        let repo_name = repo_names
            .get(&repo_id)
            .cloned()
            .unwrap_or_else(|| repo_id.clone());

        // Scan dependencies for CVEs
        let alerts = match scanner.scan_dependencies(&repo_id, &mut repo_entries).await {
            Ok(a) => a,
            Err(e) => {
                tracing::warn!("CVE monitor: scan failed for {repo_name}: {e}");
                continue;
            }
        };

        // Upsert CVE alerts (existing logic)
        for alert in &alerts {
            let filter = doc! { "cve_id": &alert.cve_id, "repo_id": &alert.repo_id };
            let update = doc! { "$setOnInsert": mongodb::bson::to_bson(alert).unwrap_or_default() };
            let _ = db
                .cve_alerts()
                .update_one(filter, update)
                .upsert(true)
                .await;
        }

        // Update SBOM entries with discovered vulnerabilities
        for entry in &repo_entries {
            if entry.known_vulnerabilities.is_empty() {
                continue;
            }
            if let Some(entry_id) = &entry.id {
                let _ = db
                    .sbom_entries()
                    .update_one(
                        doc! { "_id": entry_id },
                        doc! { "$set": {
                            "known_vulnerabilities": mongodb::bson::to_bson(&entry.known_vulnerabilities).unwrap_or_default(),
                            "updated_at": mongodb::bson::DateTime::now(),
                        }},
                    )
                    .await;
            }
        }

        // Create notifications for NEW CVEs (dedup against existing notifications)
        for alert in &alerts {
            let filter = doc! {
                "cve_id": &alert.cve_id,
                "repo_id": &alert.repo_id,
                "package_name": &alert.affected_package,
                "package_version": &alert.affected_version,
            };
            // Only insert if not already exists (upsert with $setOnInsert)
            let severity = parse_severity(alert.severity.as_deref(), alert.cvss_score);
            let mut notification = CveNotification::new(
                alert.cve_id.clone(),
                repo_id.clone(),
                repo_name.clone(),
                alert.affected_package.clone(),
                alert.affected_version.clone(),
                severity,
            );
            notification.cvss_score = alert.cvss_score;
            notification.summary = alert.summary.clone();
            notification.url = Some(format!("https://osv.dev/vulnerability/{}", alert.cve_id));

            let update = doc! {
                "$setOnInsert": mongodb::bson::to_bson(&notification).unwrap_or_default()
            };
            match db
                .cve_notifications()
                .update_one(filter, update)
                .upsert(true)
                .await
            {
                Ok(result) if result.upserted_id.is_some() => {
                    new_notifications += 1;
                }
                Err(e) => {
                    tracing::warn!("CVE monitor: failed to create notification: {e}");
                }
                _ => {} // Already exists
            }
        }
    }

    if new_notifications > 0 {
        tracing::info!(
            "CVE monitor: created {new_notifications} new notification(s) for tenant '{tenant_id}'"
        );
    } else {
        tracing::info!("CVE monitor: no new CVEs found for tenant '{tenant_id}'");
    }
}