use anyhow::Result;
use std::path::Path;
use tree_sitter::Node;

use crate::analyzers::ast_analyzer::AstAnalyzer;
use crate::types::{Severity, Vulnerability, VulnerabilityType};

use super::LanguageAnalyzer;

pub struct PythonAnalyzer;

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

    fn check_python_import(
        &self,
        node_text: &str,
        node: &Node,
        file_path: &Path,
        source: &str,
        analyzer: &AstAnalyzer,
    ) -> Option<Vulnerability> {
        let lower = node_text.to_lowercase();
        
        // Cryptography library (contains both safe and unsafe algorithms)
        if lower.contains("from cryptography") || lower.contains("import cryptography") {
            if lower.contains("rsa") {
                return Some(analyzer.create_vulnerability(
                    node,
                    file_path,
                    VulnerabilityType::Rsa2048,
                    Severity::Medium,
                    "cryptography library RSA import detected",
                    "Migrate to post-quantum alternatives",
                    source,
                ));
            }
            if lower.contains("ec") || lower.contains("ecdsa") || lower.contains("ecdh") {
                return Some(analyzer.create_vulnerability(
                    node,
                    file_path,
                    VulnerabilityType::Ecdsa,
                    Severity::Medium,
                    "cryptography library EC import detected",
                    "Migrate to post-quantum alternatives",
                    source,
                ));
            }
            if lower.contains("dsa") {
                return Some(analyzer.create_vulnerability(
                    node,
                    file_path,
                    VulnerabilityType::Dsa,
                    Severity::Medium,
                    "cryptography library DSA import detected",
                    "Migrate to post-quantum alternatives",
                    source,
                ));
            }
            if lower.contains("dh") {
                return Some(analyzer.create_vulnerability(
                    node,
                    file_path,
                    VulnerabilityType::DhKeyExchange,
                    Severity::Medium,
                    "cryptography library DH import detected",
                    "Migrate to post-quantum alternatives",
                    source,
                ));
            }
        }
        
        // PyCrypto/PyCryptodome (legacy, quantum-unsafe)
        if lower.contains("crypto.publickey") || lower.contains("from crypto import") {
            if lower.contains("rsa") {
                return Some(analyzer.create_vulnerability(
                    node,
                    file_path,
                    VulnerabilityType::Rsa2048,
                    Severity::High,
                    "PyCrypto/PyCryptodome RSA import detected",
                    "Migrate to post-quantum cryptography",
                    source,
                ));
            }
            if lower.contains("dsa") {
                return Some(analyzer.create_vulnerability(
                    node,
                    file_path,
                    VulnerabilityType::Dsa,
                    Severity::High,
                    "PyCrypto/PyCryptodome DSA import detected",
                    "Migrate to post-quantum cryptography",
                    source,
                ));
            }
            if lower.contains("elgamal") {
                return Some(analyzer.create_vulnerability(
                    node,
                    file_path,
                    VulnerabilityType::RsaOther,
                    Severity::High,
                    "PyCrypto/PyCryptodome ElGamal import detected",
                    "Migrate to post-quantum cryptography",
                    source,
                ));
            }
        }
        
        // Paramiko (SSH library, uses RSA/ECDSA)
        if lower.contains("paramiko") {
            return Some(analyzer.create_vulnerability(
                node,
                file_path,
                VulnerabilityType::Rsa2048,
                Severity::Medium,
                "Paramiko import detected - SSH library using RSA/ECDSA",
                "Verify SSH configuration uses post-quantum algorithms when available",
                source,
            ));
        }
        
        // pyOpenSSL
        if lower.contains("openssl") && lower.contains("import") {
            return Some(analyzer.create_vulnerability(
                node,
                file_path,
                VulnerabilityType::Rsa2048,
                Severity::Medium,
                "pyOpenSSL import detected - may use quantum-unsafe algorithms",
                "Audit OpenSSL usage for quantum-unsafe cryptography",
                source,
            ));
        }
        
        // python-jose, PyJWT (JWT libraries using RSA/ECDSA)
        if lower.contains("jose") || lower.contains("jwt") {
            return Some(analyzer.create_vulnerability(
                node,
                file_path,
                VulnerabilityType::Rsa2048,
                Severity::Medium,
                "JWT library import detected - typically uses RSA/ECDSA",
                "Verify JWT signing algorithm is quantum-safe",
                source,
            ));
        }

        // ecdsa library
        if lower.contains("import ecdsa") || lower.contains("from ecdsa") {
            return Some(analyzer.create_vulnerability(
                node,
                file_path,
                VulnerabilityType::Ecdsa,
                Severity::High,
                "ecdsa library import detected",
                "Migrate to CRYSTALS-Dilithium for signatures",
                source,
            ));
        }
        
        // rsa library
        if lower.contains("import rsa") && !lower.contains("paramiko") {
            return Some(analyzer.create_vulnerability(
                node,
                file_path,
                VulnerabilityType::Rsa2048,
                Severity::High,
                "rsa library import detected",
                "Migrate to post-quantum cryptography",
                source,
            ));
        }
        
        None
    }
}

impl LanguageAnalyzer for PythonAnalyzer {
    fn check_node(
        &self,
        node: &Node,
        source: &str,
        file_path: &Path,
        analyzer: &AstAnalyzer,
    ) -> Option<Vulnerability> {
        let node_type = node.kind();
        let node_text = match node.utf8_text(source.as_bytes()) {
            Ok(text) => text,
            Err(_) => return None,
        };

        match node_type {
            "call" => {
                let func_name = analyzer.get_python_func_name(node, source)?;

                // RSA operations
                if func_name.contains("RSA.generate") || func_name.contains("rsa.generate") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::Rsa2048,
                        Severity::Critical,
                        "RSA key generation detected",
                        "Migrate to CRYSTALS-Kyber for key encapsulation",
                        source,
                    ));
                }

                // ECDSA operations
                if func_name.contains("ECDSA") || func_name.contains("ECC") || 
                   func_name.contains("ec.generate") || func_name.contains("EC.generate") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::Ecdsa,
                        Severity::Critical,
                        "ECDSA/ECC operation detected",
                        "Migrate to CRYSTALS-Dilithium for signatures",
                        source,
                    ));
                }

                // DSA operations
                if func_name.contains("DSA.generate") || func_name.contains("dsa.generate") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::Dsa,
                        Severity::Critical,
                        "DSA key generation detected",
                        "Migrate to CRYSTALS-Dilithium for signatures",
                        source,
                    ));
                }

                // DH operations
                if func_name.contains("DH.generate") || func_name.contains("dh.generate") ||
                   func_name.contains("DHE") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::DhKeyExchange,
                        Severity::Critical,
                        "Diffie-Hellman key exchange detected",
                        "Migrate to CRYSTALS-Kyber for key exchange",
                        source,
                    ));
                }

                // Weak hash functions
                if func_name.contains("MD5") || func_name.contains("md5") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::Md5,
                        Severity::High,
                        "MD5 hash function detected",
                        "Replace with SHA-256 or SHA-3",
                        source,
                    ));
                }

                if func_name.contains("SHA1") || func_name.contains("sha1") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::Sha1,
                        Severity::High,
                        "SHA-1 hash function detected",
                        "Replace with SHA-256 or SHA-3",
                        source,
                    ));
                }

                // Weak ciphers
                if func_name.contains("DES") || func_name.contains("des.new") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::Des,
                        Severity::Critical,
                        "DES cipher detected",
                        "Replace with AES-256-GCM",
                        source,
                    ));
                }

                if func_name.contains("RC4") || func_name.contains("ARC4") {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::Rc4,
                        Severity::Critical,
                        "RC4 cipher detected",
                        "Replace with AES-256-GCM or ChaCha20-Poly1305",
                        source,
                    ));
                }

                // Weak random
                if func_name == "random.random" || func_name == "random.randint" {
                    return Some(analyzer.create_vulnerability(
                        node,
                        file_path,
                        VulnerabilityType::WeakRandom,
                        Severity::Medium,
                        "Non-cryptographic RNG detected",
                        "Use secrets module for cryptographic operations",
                        source,
                    ));
                }

                None
            }
            // CRITICAL FIX #3: Improved Python import detection
            "import_statement" | "import_from_statement" => {
                self.check_python_import(node_text, node, file_path, source, analyzer)
            }
            "string" => {
                analyzer.check_string_literal(node_text, node, file_path, source, node)
            }
            "integer" => {
                analyzer.check_key_size_literal(node_text, node, file_path, source)
            }
            _ => None,
        }
    }
}
