use rmcp::{
    handler::server::wrapper::Parameters, model::*, tool, tool_handler, tool_router, ServerHandler,
};

use crate::database::{Database, DatabasePool};
use crate::tools::{dast, findings, oscal, pentest, sbom};

pub struct ComplianceMcpServer {
    pool: DatabasePool,
    /// Tenant this session serves. Bound once at session creation (the HTTP
    /// factory reads the bearer-set tenant while still in the request scope;
    /// stdio passes a synthetic id) — NOT a per-request `task_local`, which is
    /// lost across the `tokio::spawn` that runs the Streamable-HTTP session.
    tenant_id: String,
    #[allow(dead_code)]
    tool_router: rmcp::handler::server::router::tool::ToolRouter<Self>,
}

impl ComplianceMcpServer {
    /// The per-tenant `Database` for this session.
    fn tenant_db(&self) -> Result<Database, rmcp::ErrorData> {
        Ok(self.pool.for_tenant_id(&self.tenant_id))
    }
}

#[tool_router]
impl ComplianceMcpServer {
    pub fn new(pool: DatabasePool, tenant_id: String) -> Self {
        Self {
            pool,
            tenant_id,
            tool_router: Self::tool_router(),
        }
    }

    // ── Findings ──────────────────────────────────────────

    #[tool(
        description = "List security findings with optional filters for repo, severity, status, and scan type"
    )]
    async fn list_findings(
        &self,
        Parameters(params): Parameters<findings::ListFindingsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        findings::list_findings(&db, params).await
    }

    #[tool(description = "Get a single finding by its ID")]
    async fn get_finding(
        &self,
        Parameters(params): Parameters<findings::GetFindingParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        findings::get_finding(&db, params).await
    }

    #[tool(description = "Get a summary of findings counts grouped by severity and status")]
    async fn findings_summary(
        &self,
        Parameters(params): Parameters<findings::FindingsSummaryParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        findings::findings_summary(&db, params).await
    }

    #[tool(
        description = "Emit an OSCAL 1.1 assessment-results document for a repo's findings (mapped findings target their compliance controls; unmapped findings are reported as-is)"
    )]
    async fn oscal_assessment(
        &self,
        Parameters(params): Parameters<oscal::OscalAssessmentParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        oscal::oscal_assessment(&db, params).await
    }

    // ── SBOM ──────────────────────────────────────────────

    #[tool(
        description = "List SBOM packages with optional filters for repo, vulnerabilities, package manager, and license"
    )]
    async fn list_sbom_packages(
        &self,
        Parameters(params): Parameters<sbom::ListSbomPackagesParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        sbom::list_sbom_packages(&db, params).await
    }

    #[tool(
        description = "Generate a vulnerability report for a repository showing all packages with known CVEs"
    )]
    async fn sbom_vuln_report(
        &self,
        Parameters(params): Parameters<sbom::SbomVulnReportParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        sbom::sbom_vuln_report(&db, params).await
    }

    // ── DAST ──────────────────────────────────────────────

    #[tool(
        description = "List DAST findings with optional filters for target, scan run, severity, exploitability, and vulnerability type"
    )]
    async fn list_dast_findings(
        &self,
        Parameters(params): Parameters<dast::ListDastFindingsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        dast::list_dast_findings(&db, params).await
    }

    #[tool(description = "Get a summary of recent DAST scan runs and finding counts")]
    async fn dast_scan_summary(
        &self,
        Parameters(params): Parameters<dast::DastScanSummaryParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        dast::dast_scan_summary(&db, params).await
    }

    // ── Pentest ─────────────────────────────────────────────

    #[tool(
        description = "List AI pentest sessions with optional filters for target, status, and strategy"
    )]
    async fn list_pentest_sessions(
        &self,
        Parameters(params): Parameters<pentest::ListPentestSessionsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::list_pentest_sessions(&db, params).await
    }

    #[tool(description = "Get a single AI pentest session by its ID")]
    async fn get_pentest_session(
        &self,
        Parameters(params): Parameters<pentest::GetPentestSessionParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::get_pentest_session(&db, params).await
    }

    #[tool(
        description = "Get the attack chain DAG for a pentest session showing each tool invocation, its reasoning, and results"
    )]
    async fn get_attack_chain(
        &self,
        Parameters(params): Parameters<pentest::GetAttackChainParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::get_attack_chain(&db, params).await
    }

    #[tool(description = "Get chat messages from a pentest session")]
    async fn get_pentest_messages(
        &self,
        Parameters(params): Parameters<pentest::GetPentestMessagesParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::get_pentest_messages(&db, params).await
    }

    #[tool(
        description = "Get aggregated pentest statistics including running sessions, vulnerability counts, and severity distribution"
    )]
    async fn pentest_stats(
        &self,
        Parameters(params): Parameters<pentest::PentestStatsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::pentest_stats(&db, params).await
    }
}

#[tool_handler]
impl ServerHandler for ComplianceMcpServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            protocol_version: ProtocolVersion::V_2024_11_05,
            capabilities: ServerCapabilities::builder()
                .enable_tools()
                .build(),
            server_info: Implementation::from_build_env(),
            instructions: Some(
                "Compliance Scanner MCP server. Query security findings, SBOM data, DAST results, and AI pentest sessions for your tenant."
                    .to_string(),
            ),
        }
    }
}