use compliance_core::error::CoreError;
use compliance_core::models::dast::{DastEvidence, DastFinding, DastVulnType};
use compliance_core::models::Severity;
use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult};
use serde_json::json;
use tracing::{info, warn};

/// Tool that checks email security configuration (DMARC and SPF records).
pub struct DmarcCheckerTool;

impl DmarcCheckerTool {
    pub fn new() -> Self {
        Self
    }

    /// Query TXT records for a given name using `dig`.
    async fn query_txt(name: &str) -> Result<Vec<String>, CoreError> {
        let output = tokio::process::Command::new("dig")
            .args(["+short", "TXT", name])
            .output()
            .await
            .map_err(|e| CoreError::Dast(format!("dig command failed: {e}")))?;

        let stdout = String::from_utf8_lossy(&output.stdout);
        let lines: Vec<String> = stdout
            .lines()
            .map(|l| l.trim().trim_matches('"').to_string())
            .filter(|l| !l.is_empty())
            .collect();
        Ok(lines)
    }

    /// Parse a DMARC record string and return the policy value.
    fn parse_dmarc_policy(record: &str) -> Option<String> {
        for part in record.split(';') {
            let part = part.trim();
            if let Some(val) = part.strip_prefix("p=") {
                return Some(val.trim().to_lowercase());
            }
        }
        None
    }

    /// Parse DMARC record for sub-domain policy (sp=).
    fn parse_dmarc_subdomain_policy(record: &str) -> Option<String> {
        for part in record.split(';') {
            let part = part.trim();
            if let Some(val) = part.strip_prefix("sp=") {
                return Some(val.trim().to_lowercase());
            }
        }
        None
    }

    /// Parse DMARC record for reporting URI (rua=).
    fn parse_dmarc_rua(record: &str) -> Option<String> {
        for part in record.split(';') {
            let part = part.trim();
            if let Some(val) = part.strip_prefix("rua=") {
                return Some(val.trim().to_string());
            }
        }
        None
    }

    /// Check if an SPF record is present and parse the policy.
    fn is_spf_record(record: &str) -> bool {
        record.starts_with("v=spf1")
    }

    /// Evaluate SPF record strength.
    fn spf_uses_soft_fail(record: &str) -> bool {
        record.contains("~all")
    }

    fn spf_allows_all(record: &str) -> bool {
        record.contains("+all")
    }
}

impl PentestTool for DmarcCheckerTool {
    fn name(&self) -> &str {
        "dmarc_checker"
    }

    fn description(&self) -> &str {
        "Checks email security configuration for a domain. Queries DMARC and SPF records and \
         evaluates policy strength. Reports missing or weak email authentication settings."
    }

    fn input_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "domain": {
                    "type": "string",
                    "description": "The domain to check (e.g., 'example.com')"
                }
            },
            "required": ["domain"]
        })
    }

    fn execute<'a>(
        &'a self,
        input: serde_json::Value,
        context: &'a PentestToolContext,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<PentestToolResult, CoreError>> + Send + 'a>> {
        Box::pin(async move {
        let domain = input
            .get("domain")
            .and_then(|v| v.as_str())
            .ok_or_else(|| CoreError::Dast("Missing required 'domain' parameter".to_string()))?;

        let target_id = context
            .target
            .id
            .map(|oid| oid.to_hex())
            .unwrap_or_else(|| "unknown".to_string());

        let mut findings = Vec::new();
        let mut email_data = json!({});

        // ---- DMARC check ----
        let dmarc_domain = format!("_dmarc.{domain}");
        let dmarc_records = Self::query_txt(&dmarc_domain).await.unwrap_or_default();

        let dmarc_record = dmarc_records.iter().find(|r| r.starts_with("v=DMARC1"));

        match dmarc_record {
            Some(record) => {
                email_data["dmarc_record"] = json!(record);

                let policy = Self::parse_dmarc_policy(record);
                let sp = Self::parse_dmarc_subdomain_policy(record);
                let rua = Self::parse_dmarc_rua(record);

                email_data["dmarc_policy"] = json!(policy);
                email_data["dmarc_subdomain_policy"] = json!(sp);
                email_data["dmarc_rua"] = json!(rua);

                // Warn on weak policy
                if let Some(ref p) = policy {
                    if p == "none" {
                        let evidence = DastEvidence {
                            request_method: "DNS".to_string(),
                            request_url: dmarc_domain.clone(),
                            request_headers: None,
                            request_body: None,
                            response_status: 0,
                            response_headers: None,
                            response_snippet: Some(record.clone()),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                            String::new(),
                            target_id.clone(),
                            DastVulnType::EmailSecurity,
                            format!("Weak DMARC policy for {domain}"),
                            format!(
                                "The DMARC policy for {domain} is set to 'none', which only monitors \
                                 but does not enforce email authentication. Attackers can spoof emails \
                                 from this domain."
                            ),
                            Severity::Medium,
                            domain.to_string(),
                            "DNS".to_string(),
                        );
                        finding.cwe = Some("CWE-290".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                            "Upgrade the DMARC policy from 'p=none' to 'p=quarantine' or \
                             'p=reject' after verifying legitimate email flows are properly \
                             authenticated."
                                .to_string(),
                        );
                        findings.push(finding);
                        warn!(domain, "DMARC policy is 'none'");
                    }
                }

                // Warn if no reporting URI
                if rua.is_none() {
                    let evidence = DastEvidence {
                        request_method: "DNS".to_string(),
                        request_url: dmarc_domain.clone(),
                        request_headers: None,
                        request_body: None,
                        response_status: 0,
                        response_headers: None,
                        response_snippet: Some(record.clone()),
                        screenshot_path: None,
                        payload: None,
                        response_time_ms: None,
                    };

                    let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::EmailSecurity,
                        format!("DMARC without reporting for {domain}"),
                        format!(
                            "The DMARC record for {domain} does not include a reporting URI (rua=). \
                             Without reporting, you will not receive aggregate feedback about email \
                             authentication failures."
                        ),
                        Severity::Info,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                    finding.cwe = Some("CWE-778".to_string());
                    finding.evidence = vec![evidence];
                    finding.remediation = Some(
                        "Add a 'rua=' tag to your DMARC record to receive aggregate reports. \
                         Example: 'rua=mailto:dmarc-reports@example.com'."
                            .to_string(),
                    );
                    findings.push(finding);
                }
            }
            None => {
                email_data["dmarc_record"] = json!(null);

                let evidence = DastEvidence {
                    request_method: "DNS".to_string(),
                    request_url: dmarc_domain.clone(),
                    request_headers: None,
                    request_body: None,
                    response_status: 0,
                    response_headers: None,
                    response_snippet: Some("No DMARC record found".to_string()),
                    screenshot_path: None,
                    payload: None,
                    response_time_ms: None,
                };

                let mut finding = DastFinding::new(
                    String::new(),
                    target_id.clone(),
                    DastVulnType::EmailSecurity,
                    format!("Missing DMARC record for {domain}"),
                    format!(
                        "No DMARC record was found for {domain}. Without DMARC, there is no \
                         policy to prevent email spoofing and phishing attacks using this domain."
                    ),
                    Severity::High,
                    domain.to_string(),
                    "DNS".to_string(),
                );
                finding.cwe = Some("CWE-290".to_string());
                finding.evidence = vec![evidence];
                finding.remediation = Some(
                    "Create a DMARC TXT record at _dmarc.<domain>. Start with 'v=DMARC1; p=none; \
                     rua=mailto:dmarc@example.com' and gradually move to 'p=reject'."
                        .to_string(),
                );
                findings.push(finding);
                warn!(domain, "No DMARC record found");
            }
        }

        // ---- SPF check ----
        let txt_records = Self::query_txt(domain).await.unwrap_or_default();
        let spf_record = txt_records.iter().find(|r| Self::is_spf_record(r));

        match spf_record {
            Some(record) => {
                email_data["spf_record"] = json!(record);

                if Self::spf_allows_all(record) {
                    let evidence = DastEvidence {
                        request_method: "DNS".to_string(),
                        request_url: domain.to_string(),
                        request_headers: None,
                        request_body: None,
                        response_status: 0,
                        response_headers: None,
                        response_snippet: Some(record.clone()),
                        screenshot_path: None,
                        payload: None,
                        response_time_ms: None,
                    };

                    let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::EmailSecurity,
                        format!("SPF allows all senders for {domain}"),
                        format!(
                            "The SPF record for {domain} uses '+all' which allows any server to \
                             send email on behalf of this domain, completely negating SPF protection."
                        ),
                        Severity::Critical,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                    finding.cwe = Some("CWE-290".to_string());
                    finding.evidence = vec![evidence];
                    finding.remediation = Some(
                        "Change '+all' to '-all' (hard fail) or '~all' (soft fail) in your SPF record. \
                         Only list authorized mail servers."
                            .to_string(),
                    );
                    findings.push(finding);
                } else if Self::spf_uses_soft_fail(record) {
                    let evidence = DastEvidence {
                        request_method: "DNS".to_string(),
                        request_url: domain.to_string(),
                        request_headers: None,
                        request_body: None,
                        response_status: 0,
                        response_headers: None,
                        response_snippet: Some(record.clone()),
                        screenshot_path: None,
                        payload: None,
                        response_time_ms: None,
                    };

                    let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::EmailSecurity,
                        format!("SPF soft fail for {domain}"),
                        format!(
                            "The SPF record for {domain} uses '~all' (soft fail) instead of \
                             '-all' (hard fail). Soft fail marks unauthorized emails as suspicious \
                             but does not reject them."
                        ),
                        Severity::Low,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                    finding.cwe = Some("CWE-290".to_string());
                    finding.evidence = vec![evidence];
                    finding.remediation = Some(
                        "Consider changing '~all' to '-all' in your SPF record once you have \
                         confirmed all legitimate mail sources are listed."
                            .to_string(),
                    );
                    findings.push(finding);
                }
            }
            None => {
                email_data["spf_record"] = json!(null);

                let evidence = DastEvidence {
                    request_method: "DNS".to_string(),
                    request_url: domain.to_string(),
                    request_headers: None,
                    request_body: None,
                    response_status: 0,
                    response_headers: None,
                    response_snippet: Some("No SPF record found".to_string()),
                    screenshot_path: None,
                    payload: None,
                    response_time_ms: None,
                };

                let mut finding = DastFinding::new(
                    String::new(),
                    target_id.clone(),
                    DastVulnType::EmailSecurity,
                    format!("Missing SPF record for {domain}"),
                    format!(
                        "No SPF record was found for {domain}. Without SPF, any server can claim \
                         to send email on behalf of this domain."
                    ),
                    Severity::High,
                    domain.to_string(),
                    "DNS".to_string(),
                );
                finding.cwe = Some("CWE-290".to_string());
                finding.evidence = vec![evidence];
                finding.remediation = Some(
                    "Create an SPF TXT record for your domain. Example: \
                     'v=spf1 include:_spf.google.com -all'."
                        .to_string(),
                );
                findings.push(finding);
                warn!(domain, "No SPF record found");
            }
        }

        let count = findings.len();
        info!(domain, findings = count, "DMARC/SPF check complete");

        Ok(PentestToolResult {
            summary: if count > 0 {
                format!("Found {count} email security issues for {domain}.")
            } else {
                format!("Email security configuration looks good for {domain}.")
            },
            findings,
            data: email_data,
        })
        })
    }
}