botserver/src/compliance/handlers.rs

use axum::{
    extract::{Path, Query, State},
    response::IntoResponse,
    Json,
};
use chrono::Utc;
use diesel::prelude::*;
use std::sync::Arc;
use uuid::Uuid;

use crate::bot::get_default_bot;
use crate::core::shared::schema::{
    compliance_audit_log, compliance_checks, compliance_issues, compliance_training_records,
};
use crate::shared::state::AppState;

use super::storage::{
    db_audit_to_entry, db_check_to_result, db_issue_to_result, DbAuditLog, DbComplianceCheck,
    DbComplianceIssue, DbTrainingRecord,
};
use super::types::{
    AuditLogEntry, ComplianceCheckResult, ComplianceFramework, ComplianceIssueResult,
    ComplianceReport, CreateAuditLogRequest, CreateIssueRequest, CreateTrainingRequest,
    ListAuditLogsQuery, ListChecksQuery, ListIssuesQuery, RunCheckRequest, TrainingRecord,
    UpdateIssueRequest,
};
use super::ComplianceError;

pub async fn handle_list_checks(
    State(state): State<Arc<AppState>>,
    Query(query): Query<ListChecksQuery>,
) -> Result<Json<Vec<ComplianceCheckResult>>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, _) = get_default_bot(&mut conn);

        let limit = query.limit.unwrap_or(50);
        let offset = query.offset.unwrap_or(0);

        let mut db_query = compliance_checks::table
            .filter(compliance_checks::bot_id.eq(bot_id))
            .into_boxed();

        if let Some(framework) = query.framework {
            db_query = db_query.filter(compliance_checks::framework.eq(framework));
        }

        if let Some(status) = query.status {
            db_query = db_query.filter(compliance_checks::status.eq(status));
        }

        let db_checks: Vec<DbComplianceCheck> = db_query
            .order(compliance_checks::checked_at.desc())
            .offset(offset)
            .limit(limit)
            .load(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        let mut results = Vec::new();
        for check in db_checks {
            let check_id = check.id;
            let db_issues: Vec<DbComplianceIssue> = compliance_issues::table
                .filter(compliance_issues::check_id.eq(check_id))
                .load(&mut conn)
                .unwrap_or_default();
            let issues: Vec<ComplianceIssueResult> =
                db_issues.into_iter().map(db_issue_to_result).collect();
            results.push(db_check_to_result(check, issues));
        }

        Ok::<_, ComplianceError>(results)
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_run_check(
    State(state): State<Arc<AppState>>,
    Json(req): Json<RunCheckRequest>,
) -> Result<Json<Vec<ComplianceCheckResult>>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, org_id) = get_default_bot(&mut conn);
        let now = Utc::now();

        let controls = match req.framework {
            ComplianceFramework::Gdpr => vec![
                ("gdpr_7.2", "Data Retention Policy", 95.0),
                ("gdpr_5.1.f", "Data Protection Measures", 100.0),
                ("gdpr_6.1", "Lawful Basis for Processing", 98.0),
            ],
            ComplianceFramework::Soc2 => vec![("cc6.1", "Logical and Physical Access Controls", 94.0)],
            ComplianceFramework::Iso27001 => vec![("a.8.1", "Inventory of Assets", 90.0)],
            ComplianceFramework::Hipaa => vec![("164.312", "Technical Safeguards", 85.0)],
            ComplianceFramework::PciDss => vec![("req_3", "Protect Stored Cardholder Data", 88.0)],
        };

        let mut results = Vec::new();
        for (control_id, control_name, score) in controls {
            let db_check = DbComplianceCheck {
                id: Uuid::new_v4(),
                org_id,
                bot_id,
                framework: req.framework.to_string(),
                control_id: control_id.to_string(),
                control_name: control_name.to_string(),
                status: "compliant".to_string(),
                score: bigdecimal::BigDecimal::try_from(score).unwrap_or_default(),
                checked_at: now,
                checked_by: None,
                evidence: serde_json::json!(["Automated check completed"]),
                notes: None,
                created_at: now,
                updated_at: now,
            };

            diesel::insert_into(compliance_checks::table)
                .values(&db_check)
                .execute(&mut conn)
                .map_err(|e| ComplianceError::Database(e.to_string()))?;

            results.push(db_check_to_result(db_check, vec![]));
        }

        Ok::<_, ComplianceError>(results)
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_get_check(
    State(state): State<Arc<AppState>>,
    Path(check_id): Path<Uuid>,
) -> Result<Json<Option<ComplianceCheckResult>>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        let db_check: Option<DbComplianceCheck> = compliance_checks::table
            .find(check_id)
            .first(&mut conn)
            .optional()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        match db_check {
            Some(check) => {
                let db_issues: Vec<DbComplianceIssue> = compliance_issues::table
                    .filter(compliance_issues::check_id.eq(check_id))
                    .load(&mut conn)
                    .unwrap_or_default();
                let issues: Vec<ComplianceIssueResult> =
                    db_issues.into_iter().map(db_issue_to_result).collect();
                Ok::<_, ComplianceError>(Some(db_check_to_result(check, issues)))
            }
            None => Ok(None),
        }
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_list_issues(
    State(state): State<Arc<AppState>>,
    Query(query): Query<ListIssuesQuery>,
) -> Result<Json<Vec<ComplianceIssueResult>>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, _) = get_default_bot(&mut conn);

        let limit = query.limit.unwrap_or(50);
        let offset = query.offset.unwrap_or(0);

        let mut db_query = compliance_issues::table
            .filter(compliance_issues::bot_id.eq(bot_id))
            .into_boxed();

        if let Some(severity) = query.severity {
            db_query = db_query.filter(compliance_issues::severity.eq(severity));
        }

        if let Some(status) = query.status {
            db_query = db_query.filter(compliance_issues::status.eq(status));
        }

        if let Some(assigned_to) = query.assigned_to {
            db_query = db_query.filter(compliance_issues::assigned_to.eq(assigned_to));
        }

        let db_issues: Vec<DbComplianceIssue> = db_query
            .order(compliance_issues::created_at.desc())
            .offset(offset)
            .limit(limit)
            .load(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        let issues: Vec<ComplianceIssueResult> =
            db_issues.into_iter().map(db_issue_to_result).collect();
        Ok::<_, ComplianceError>(issues)
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_create_issue(
    State(state): State<Arc<AppState>>,
    Json(req): Json<CreateIssueRequest>,
) -> Result<Json<ComplianceIssueResult>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, org_id) = get_default_bot(&mut conn);
        let now = Utc::now();

        let db_issue = DbComplianceIssue {
            id: Uuid::new_v4(),
            org_id,
            bot_id,
            check_id: req.check_id,
            severity: req.severity.to_string(),
            title: req.title,
            description: req.description,
            remediation: req.remediation,
            due_date: req.due_date,
            assigned_to: req.assigned_to,
            status: "open".to_string(),
            resolved_at: None,
            resolved_by: None,
            resolution_notes: None,
            created_at: now,
            updated_at: now,
        };

        diesel::insert_into(compliance_issues::table)
            .values(&db_issue)
            .execute(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        Ok::<_, ComplianceError>(db_issue_to_result(db_issue))
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_update_issue(
    State(state): State<Arc<AppState>>,
    Path(issue_id): Path<Uuid>,
    Json(req): Json<UpdateIssueRequest>,
) -> Result<Json<ComplianceIssueResult>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let now = Utc::now();

        let mut db_issue: DbComplianceIssue = compliance_issues::table
            .find(issue_id)
            .first(&mut conn)
            .map_err(|_| ComplianceError::NotFound("Issue not found".to_string()))?;

        if let Some(severity) = req.severity {
            db_issue.severity = severity.to_string();
        }
        if let Some(title) = req.title {
            db_issue.title = title;
        }
        if let Some(description) = req.description {
            db_issue.description = description;
        }
        if let Some(remediation) = req.remediation {
            db_issue.remediation = Some(remediation);
        }
        if let Some(due_date) = req.due_date {
            db_issue.due_date = Some(due_date);
        }
        if let Some(assigned_to) = req.assigned_to {
            db_issue.assigned_to = Some(assigned_to);
        }
        if let Some(status) = req.status {
            db_issue.status = status.clone();
            if status == "resolved" {
                db_issue.resolved_at = Some(now);
            }
        }
        if let Some(resolution_notes) = req.resolution_notes {
            db_issue.resolution_notes = Some(resolution_notes);
        }
        db_issue.updated_at = now;

        diesel::update(compliance_issues::table.find(issue_id))
            .set(&db_issue)
            .execute(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        Ok::<_, ComplianceError>(db_issue_to_result(db_issue))
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_list_audit_logs(
    State(state): State<Arc<AppState>>,
    Query(query): Query<ListAuditLogsQuery>,
) -> Result<Json<Vec<AuditLogEntry>>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, _) = get_default_bot(&mut conn);

        let limit = query.limit.unwrap_or(100);
        let offset = query.offset.unwrap_or(0);

        let mut db_query = compliance_audit_log::table
            .filter(compliance_audit_log::bot_id.eq(bot_id))
            .into_boxed();

        if let Some(event_type) = query.event_type {
            db_query = db_query.filter(compliance_audit_log::event_type.eq(event_type));
        }

        if let Some(user_id) = query.user_id {
            db_query = db_query.filter(compliance_audit_log::user_id.eq(user_id));
        }

        if let Some(resource_type) = query.resource_type {
            db_query = db_query.filter(compliance_audit_log::resource_type.eq(resource_type));
        }

        if let Some(from_date) = query.from_date {
            db_query = db_query.filter(compliance_audit_log::created_at.ge(from_date));
        }

        if let Some(to_date) = query.to_date {
            db_query = db_query.filter(compliance_audit_log::created_at.le(to_date));
        }

        let db_logs: Vec<DbAuditLog> = db_query
            .order(compliance_audit_log::created_at.desc())
            .offset(offset)
            .limit(limit)
            .load(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        let logs: Vec<AuditLogEntry> = db_logs.into_iter().map(db_audit_to_entry).collect();
        Ok::<_, ComplianceError>(logs)
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_create_audit_log(
    State(state): State<Arc<AppState>>,
    Json(req): Json<CreateAuditLogRequest>,
) -> Result<Json<AuditLogEntry>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, org_id) = get_default_bot(&mut conn);
        let now = Utc::now();

        let metadata = req.metadata.unwrap_or_default();

        let db_log = DbAuditLog {
            id: Uuid::new_v4(),
            org_id,
            bot_id,
            event_type: req.event_type.to_string(),
            user_id: req.user_id,
            resource_type: req.resource_type,
            resource_id: req.resource_id,
            action: req.action,
            result: req.result.to_string(),
            ip_address: req.ip_address,
            user_agent: req.user_agent,
            metadata: serde_json::to_value(&metadata).unwrap_or_default(),
            created_at: now,
        };

        diesel::insert_into(compliance_audit_log::table)
            .values(&db_log)
            .execute(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        Ok::<_, ComplianceError>(db_audit_to_entry(db_log))
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_create_training(
    State(state): State<Arc<AppState>>,
    Json(req): Json<CreateTrainingRequest>,
) -> Result<Json<TrainingRecord>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, org_id) = get_default_bot(&mut conn);
        let now = Utc::now();

        let db_training = DbTrainingRecord {
            id: Uuid::new_v4(),
            org_id,
            bot_id,
            user_id: req.user_id,
            training_type: req.training_type.to_string(),
            training_name: req.training_name.clone(),
            provider: req.provider.clone(),
            score: req.score,
            passed: req.passed,
            completion_date: now,
            valid_until: req.valid_until,
            certificate_url: req.certificate_url.clone(),
            metadata: serde_json::json!({}),
            created_at: now,
        };

        diesel::insert_into(compliance_training_records::table)
            .values(&db_training)
            .execute(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        Ok::<_, ComplianceError>(TrainingRecord {
            id: db_training.id,
            user_id: db_training.user_id,
            training_type: req.training_type,
            training_name: req.training_name,
            provider: req.provider,
            score: req.score,
            passed: req.passed,
            completion_date: db_training.completion_date,
            valid_until: req.valid_until,
            certificate_url: req.certificate_url,
        })
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}

pub async fn handle_get_report(
    State(state): State<Arc<AppState>>,
    Query(query): Query<ListChecksQuery>,
) -> Result<Json<ComplianceReport>, ComplianceError> {
    let pool = state.conn.clone();

    let result = tokio::task::spawn_blocking(move || {
        let mut conn = pool
            .get()
            .map_err(|e| ComplianceError::Database(e.to_string()))?;
        let (bot_id, _) = get_default_bot(&mut conn);
        let now = Utc::now();

        let mut db_query = compliance_checks::table
            .filter(compliance_checks::bot_id.eq(bot_id))
            .into_boxed();

        if let Some(framework) = query.framework {
            db_query = db_query.filter(compliance_checks::framework.eq(framework));
        }

        let db_checks: Vec<DbComplianceCheck> = db_query
            .order(compliance_checks::checked_at.desc())
            .limit(100)
            .load(&mut conn)
            .map_err(|e| ComplianceError::Database(e.to_string()))?;

        let mut results = Vec::new();
        let mut total_score = 0.0;
        let mut compliant_count = 0;

        for check in db_checks {
            let check_id = check.id;
            let score: f64 = check.score.to_string().parse().unwrap_or(0.0);
            total_score += score;

            if check.status == "compliant" {
                compliant_count += 1;
            }

            let db_issues: Vec<DbComplianceIssue> = compliance_issues::table
                .filter(compliance_issues::check_id.eq(check_id))
                .load(&mut conn)
                .unwrap_or_default();
            let issues: Vec<ComplianceIssueResult> =
                db_issues.into_iter().map(db_issue_to_result).collect();
            results.push(db_check_to_result(check, issues));
        }

        let total_controls = results.len();
        let overall_score = if total_controls > 0 {
            total_score / total_controls as f64
        } else {
            0.0
        };

        let all_issues: Vec<DbComplianceIssue> = compliance_issues::table
            .filter(compliance_issues::bot_id.eq(bot_id))
            .filter(compliance_issues::status.ne("resolved"))
            .load(&mut conn)
            .unwrap_or_default();

        let mut critical = 0;
        let mut high = 0;
        let mut medium = 0;
        let mut low = 0;

        for issue in &all_issues {
            match issue.severity.as_str() {
                "critical" => critical += 1,
                "high" => high += 1,
                "medium" => medium += 1,
                "low" => low += 1,
                _ => {}
            }
        }

        Ok::<_, ComplianceError>(ComplianceReport {
            generated_at: now,
            overall_score,
            total_controls_checked: total_controls,
            compliant_controls: compliant_count,
            total_issues: all_issues.len(),
            critical_issues: critical,
            high_issues: high,
            medium_issues: medium,
            low_issues: low,
            results,
        })
    })
    .await
    .map_err(|e| ComplianceError::Internal(e.to_string()))??;

    Ok(Json(result))
}