use serde::{Deserialize, Serialize};
use regex::Regex;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Vulnerability {
    pub id: String,
    pub severity: Severity,
    pub category: Category,
    pub title: String,
    pub description: String,
    pub affected_tools: Vec<String>,
    pub evidence_pattern: String,
    pub remediation: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Severity {
    Critical,
    High,
    Medium,
    Low,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Category {
    CommandInjection,
    PathTraversal,
    Authentication,
    NetworkExposure,
    InformationDisclosure,
    CredentialExposure,
    SSRF,
    ToolPoisoning,
}

pub struct VulnDb {
    vulnerabilities: Vec<Vulnerability>,
}

impl VulnDb {
    pub fn new() -> Self {
        eprintln!("[VULNDB] Loading {} vulnerabilities", 10);
        let mut db = Self {
            vulnerabilities: Vec::new(),
        };
        db.load_default_vulnerabilities();
        db
    }

    fn load_default_vulnerabilities(&mut self) {
        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-001".to_string(),
            severity: Severity::Critical,
            category: Category::CommandInjection,
            title: "Command injection via tool description".to_string(),
            description: "MCP server allows arbitrary command execution through tool inputSchema".to_string(),
            affected_tools: vec!["execute_command".to_string(), "run_shell".to_string(), "exec".to_string()],
            evidence_pattern: r#"("command"|"exec"|"shell").*?"type":\s*"string""#.to_string(),
            remediation: "Remove execute_command tool or sandbox with strict allowlist".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-002".to_string(),
            severity: Severity::High,
            category: Category::PathTraversal,
            title: "Path traversal in file system tools".to_string(),
            description: "File read/write tools lack path validation, allowing access outside workspace".to_string(),
            affected_tools: vec!["read_file".to_string(), "write_file".to_string(), "list_directory".to_string()],
            evidence_pattern: r#""path"\s*:\s*\{[^}]*"type"\s*:\s*"string""#.to_string(),
            remediation: "Add path pattern validation to inputSchema".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-003".to_string(),
            severity: Severity::High,
            category: Category::SSRF,
            title: "SSRF via fetch_url tool".to_string(),
            description: "URL fetching tool allows access to internal services and metadata endpoints".to_string(),
            affected_tools: vec!["fetch_url".to_string(), "http_request".to_string(), "download".to_string()],
            evidence_pattern: r#""url"\s*:\s*\{[^}]*"type"\s*:\s*"string""#.to_string(),
            remediation: "Add URL allowlist or block internal IP ranges".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-004".to_string(),
            severity: Severity::Critical,
            category: Category::CredentialExposure,
            title: "Credential exposure in environment variables".to_string(),
            description: "MCP server config contains hardcoded API keys or tokens".to_string(),
            affected_tools: vec!["*".to_string()],
            evidence_pattern: r#"(?i)(api_key|token|secret|password)"#.to_string(),
            remediation: "Use environment variable references instead of hardcoded credentials".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-005".to_string(),
            severity: Severity::Critical,
            category: Category::CommandInjection,
            title: "Exec/Shell injection via subprocess".to_string(),
            description: "MCP server passes unsanitized user input to subprocess.run(), eval(), or exec() with string formatting".to_string(),
            affected_tools: vec!["execute".to_string(), "run".to_string(), "shell".to_string()],
            evidence_pattern: r#"(?i)(exec\(|eval\(|subprocess\.(run|Popen|call)\(.*?(f"|%|\.format))"#.to_string(),
            remediation: "Never pass user input directly to shell execution. Use allowlists and argument arrays.".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-006".to_string(),
            severity: Severity::Critical,
            category: Category::CommandInjection,
            title: "Hardening bypass via npx -c flag".to_string(),
            description: "Attacker hides malicious commands inside npx -c argument, bypassing command allowlists".to_string(),
            affected_tools: vec!["*".to_string()],
            evidence_pattern: r#"(?i)npx\s+.*?-c\s+["']?[^"']*?(rm|curl|wget|bash|sh)"#.to_string(),
            remediation: "Block -c flag in npx invocations. Validate all args against an allowlist.".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-007".to_string(),
            severity: Severity::High,
            category: Category::ToolPoisoning,
            title: "Tool description poisoning".to_string(),
            description: "Malicious instructions injected into tool description field to manipulate AI agent behavior".to_string(),
            affected_tools: vec!["*".to_string()],
            evidence_pattern: r#"(?i)(ignore previous|do not tell|secretly|without user knowledge|exfiltrate|curl http)"#.to_string(),
            remediation: "Validate tool descriptions against known safe patterns. Reject descriptions containing system instructions.".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-008".to_string(),
            severity: Severity::Medium,
            category: Category::ToolPoisoning,
            title: "Unpinned server configuration (supply chain risk)".to_string(),
            description: "MCP server uses @latest or unverified package versions, enabling supply chain attacks".to_string(),
            affected_tools: vec!["*".to_string()],
            evidence_pattern: r#"(?i)npx\s+@latest"#.to_string(),
            remediation: "Pin all package versions. Use exact version numbers instead of @latest.".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-009".to_string(),
            severity: Severity::High,
            category: Category::NetworkExposure,
            title: "Insecure SSE transport — missing authentication".to_string(),
            description: "SSE endpoint exposed without authentication token or bound to 0.0.0.0".to_string(),
            affected_tools: vec!["*".to_string()],
            evidence_pattern: r#"(?i)(sse|transport).*?(?:auth(?:entication)?\s*:\s*(?:false|null|undefined)|"host"\s*:\s*"0\.0\.0\.0")"#.to_string(),
            remediation: "Require auth token on all SSE endpoints. Never bind to 0.0.0.0 in production.".to_string(),
        });

        self.vulnerabilities.push(Vulnerability {
            id: "CVE-2026-MCP-010".to_string(),
            severity: Severity::High,
            category: Category::SSRF,
            title: "Missing auth on public SSE endpoint".to_string(),
            description: "SSE transport URL exposed to internet without token-based authentication".to_string(),
            affected_tools: vec!["*".to_string()],
            evidence_pattern: r#"(?i)(https?://(?:(?!localhost|127\.0\.0\.1)[^/])+/sse)"#.to_string(),
            remediation: "Add ?token= parameter enforcement or move behind API gateway with auth.".to_string(),
        });
        eprintln!("[VULNDB] Loaded {} total vulns", self.vulnerabilities.len());
    }

    pub fn check_tool(&self, tool_name: &str, tool_schema: &str) -> Vec<Vulnerability> {
        eprintln!("[CHECK] tool={} checking against {} vulns", tool_name, self.vulnerabilities.len());
        let mut matches = Vec::new();
        
        for vuln in &self.vulnerabilities {
            let tool_matches = vuln.affected_tools.iter().any(|t| {
                t == "*" || tool_name.to_lowercase().contains(&t.to_lowercase())
            });

            if tool_matches {
                let regex = Regex::new(&vuln.evidence_pattern);
                if let Ok(re) = regex {
                    if re.is_match(tool_schema) {
                        matches.push(vuln.clone());
                    }
                }
            }
        }
        
        matches
    }

    pub fn all_vulnerabilities(&self) -> &[Vulnerability] {
        &self.vulnerabilities
    }

    pub fn get_rules_for_policy(&self) -> Vec<PolicyRule> {
        self.vulnerabilities.iter().map(|v| PolicyRule {
            name: format!("block-{}", v.id.to_lowercase().replace(".", "-")),
            action: if v.severity == Severity::Critical { "block".to_string() } else { "log".to_string() },
            severity: format!("{:?}", v.severity),
            condition: Condition {
                tool_pattern: Some(v.affected_tools.join("|")),
                tool_name: None,
                argument_contains: None,
                path_matches: None,
            },
        }).collect()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PolicyRule {
    pub name: String,
    pub action: String,
    pub severity: String,
    pub condition: Condition,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Condition {
    pub tool_name: Option<String>,
    pub tool_pattern: Option<String>,
    pub argument_contains: Option<Vec<String>>,
    pub path_matches: Option<String>,
}