use mongodb::{Client, Collection};

use compliance_core::models::*;

#[derive(Clone, Debug)]
pub struct Database {
    inner: mongodb::Database,
}

impl Database {
    pub async fn connect(uri: &str, db_name: &str) -> Result<Self, mongodb::error::Error> {
        let client = Client::with_uri_str(uri).await?;
        let db = client.database(db_name);
        db.run_command(mongodb::bson::doc! { "ping": 1 }).await?;
        tracing::info!("MCP server connected to MongoDB '{db_name}'");
        Ok(Self { inner: db })
    }

    pub fn findings(&self) -> Collection<Finding> {
        self.inner.collection("findings")
    }

    pub fn sbom_entries(&self) -> Collection<SbomEntry> {
        self.inner.collection("sbom_entries")
    }

    pub fn dast_findings(&self) -> Collection<DastFinding> {
        self.inner.collection("dast_findings")
    }

    pub fn dast_scan_runs(&self) -> Collection<DastScanRun> {
        self.inner.collection("dast_scan_runs")
    }
}