"""
Utility for generating comment-based patches for security findings.

Generates Git unified diff patches that add @TODO-pq comment annotations
above vulnerable lines without modifying the actual code.
"""
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Vulnerability hint table
# Each entry: (generic_fix_hint, {ext: language_specific_hint})
# ---------------------------------------------------------------------------
_VULN_HINTS: dict[str, tuple[str, dict[str, str]]] = {
    "md5": (
        "sha256(data)",
        {
            ".py":   "hashlib.sha256(data).hexdigest()",
            ".java": 'MessageDigest.getInstance("SHA-256")',
            ".js":   "crypto.createHash('sha256').update(data).digest('hex')",
            ".ts":   "crypto.createHash('sha256').update(data).digest('hex')",
            ".go":   "sha256.New() / sha256.Sum256(data)",
            ".rb":   "Digest::SHA256.hexdigest(data)",
            ".php":  "hash('sha256', $data)",
            ".cs":   "SHA256.HashData(data)",
            ".kt":   'MessageDigest.getInstance("SHA-256")',
            ".swift":"SHA256.hash(data: data)",
            ".rs":   "sha2::Sha256::digest(data)",
            ".cpp":  "EVP_DigestInit_ex(ctx, EVP_sha256(), NULL)",
            ".scala":'MessageDigest.getInstance("SHA-256")',
        },
    ),
    "sha1": (
        "sha256(data)",
        {
            ".py":   "hashlib.sha256(data).hexdigest()",
            ".java": 'MessageDigest.getInstance("SHA-256")',
            ".js":   "crypto.createHash('sha256').update(data).digest('hex')",
            ".ts":   "crypto.createHash('sha256').update(data).digest('hex')",
            ".go":   "sha256.New() / sha256.Sum256(data)",
            ".rb":   "Digest::SHA256.hexdigest(data)",
            ".php":  "hash('sha256', $data)",
            ".cs":   "SHA256.HashData(data)",
            ".kt":   'MessageDigest.getInstance("SHA-256")',
            ".swift":"SHA256.hash(data: data)",
            ".rs":   "sha2::Sha256::digest(data)",
        },
    ),
    "md4": (
        "sha256(data)",
        {},
    ),
    "des": (
        "AES-256-GCM",
        {
            ".py":   "Cipher(algorithms.AES(key), modes.GCM(iv))",
            ".java": 'Cipher.getInstance("AES/GCM/NoPadding")',
            ".js":   "crypto.createCipheriv('aes-256-gcm', key, iv)",
            ".ts":   "crypto.createCipheriv('aes-256-gcm', key, iv)",
            ".go":   "aes.NewCipher(key) + cipher.NewGCM(block)",
            ".rb":   "OpenSSL::Cipher.new('aes-256-gcm')",
            ".php":  "openssl_encrypt($data, 'aes-256-gcm', $key, 0, $iv, $tag)",
            ".cs":   "AesGcm(key)",
            ".kt":   'Cipher.getInstance("AES/GCM/NoPadding")',
            ".swift":"AES.GCM.seal(data, using: key)",
            ".rs":   "aes_gcm::Aes256Gcm",
        },
    ),
    "3des": (
        "AES-256-GCM",
        {
            ".py":   "Cipher(algorithms.AES(key), modes.GCM(iv))",
            ".java": 'Cipher.getInstance("AES/GCM/NoPadding")',
            ".go":   "aes.NewCipher(key) + cipher.NewGCM(block)",
            ".cs":   "AesGcm(key)",
            ".kt":   'Cipher.getInstance("AES/GCM/NoPadding")',
        },
    ),
    "rc4": (
        "AES-256-GCM or ChaCha20-Poly1305",
        {
            ".py":   "Cipher(algorithms.AES(key), modes.GCM(iv))",
            ".go":   "chacha20poly1305.New(key)",
            ".js":   "crypto.createCipheriv('aes-256-gcm', key, iv)",
            ".ts":   "crypto.createCipheriv('aes-256-gcm', key, iv)",
        },
    ),
    "rsa": (
        "ML-KEM (key exchange) or ML-DSA (signatures)",
        {
            ".py":   "from cryptography.hazmat.primitives.asymmetric import x25519  # interim",
            ".java": "// Migrate to ML-KEM / ML-DSA when available in JCA",
            ".go":   "// Use golang.org/x/crypto/mlkem768 (FIPS 203)",
        },
    ),
    "ecc": (
        "ML-DSA (FIPS 204) for signatures, ML-KEM (FIPS 203) for key exchange",
        {
            ".py":   "# Use liboqs-python for ML-DSA/ML-KEM",
            ".java": "// Migrate to ML-KEM / ML-DSA when available in JCA",
            ".go":   "// Use golang.org/x/crypto/mlkem768 (FIPS 203)",
        },
    ),
    "dh_dsa": (
        "ML-KEM (FIPS 203) for key exchange, ML-DSA (FIPS 204) for signatures",
        {},
    ),
    "aes128": (
        "AES-256 (256-bit key)",
        {
            ".py":   "key = os.urandom(32)  # 256-bit",
            ".java": 'KeyGenerator.getInstance("AES").init(256)',
            ".js":   "crypto.randomBytes(32)  // 256-bit key",
            ".ts":   "crypto.randomBytes(32)  // 256-bit key",
            ".go":   "key := make([]byte, 32)  // 256-bit",
            ".cs":   "new AesGcm(key256)",
            ".kt":   'KeyGenerator.getInstance("AES").init(256)',
        },
    ),
    "ecb": (
        "AES-GCM (authenticated encryption)",
        {
            ".py":   "Cipher(algorithms.AES(key), modes.GCM(iv))",
            ".java": 'Cipher.getInstance("AES/GCM/NoPadding")',
            ".js":   "crypto.createCipheriv('aes-256-gcm', key, iv)",
            ".ts":   "crypto.createCipheriv('aes-256-gcm', key, iv)",
            ".go":   "cipher.NewGCM(block)",
            ".cs":   "AesGcm(key)",
            ".kt":   'Cipher.getInstance("AES/GCM/NoPadding")',
            ".swift":"AES.GCM.seal(data, using: key)",
        },
    ),
    "tls_weak": (
        "TLS 1.3 minimum",
        {
            ".py":   "ssl_context.minimum_version = ssl.TLSVersion.TLSv1_3",
            ".java": 'sslContext.createSSLEngine(); // use TLSv1.3',
            ".go":   "tls.Config{MinVersion: tls.VersionTLS13}",
            ".js":   "tls.connect({ minVersion: 'TLSv1.3' })",
            ".ts":   "tls.connect({ minVersion: 'TLSv1.3' })",
            ".rb":   "OpenSSL::SSL::SSLContext.new.min_version = OpenSSL::SSL::TLS1_3_VERSION",
            ".cs":   "SslProtocols.Tls13",
        },
    ),
    "weak_random": (
        "cryptographically secure RNG",
        {
            ".py":   "import secrets; secrets.token_bytes(32)",
            ".java": "new java.security.SecureRandom()",
            ".js":   "crypto.randomBytes(32)",
            ".ts":   "crypto.randomBytes(32)",
            ".go":   "crypto/rand.Read(buf)",
            ".rb":   "SecureRandom.bytes(32)",
            ".php":  "random_bytes(32)",
            ".cs":   "RandomNumberGenerator.GetBytes(32)",
            ".kt":   "java.security.SecureRandom()",
            ".swift":"SystemRandomNumberGenerator()",
            ".rs":   "rand::rngs::OsRng",
        },
    ),
    "static_iv": (
        "random IV per encryption: os.urandom(16) / SecureRandom",
        {
            ".py":   "iv = os.urandom(16)",
            ".java": "byte[] iv = new byte[12]; new SecureRandom().nextBytes(iv);",
            ".js":   "const iv = crypto.randomBytes(16);",
            ".ts":   "const iv = crypto.randomBytes(16);",
            ".go":   "iv := make([]byte, 12); rand.Read(iv)",
            ".cs":   "byte[] iv = new byte[12]; RandomNumberGenerator.Fill(iv);",
            ".kt":   "val iv = ByteArray(12).also { SecureRandom().nextBytes(it) }",
        },
    ),
    "hardcoded_key": (
        "environment variable or secrets manager",
        {
            ".py":   "key = os.environ['SECRET_KEY']",
            ".java": 'System.getenv("SECRET_KEY")',
            ".js":   "process.env.SECRET_KEY",
            ".ts":   "process.env.SECRET_KEY",
            ".go":   'os.Getenv("SECRET_KEY")',
            ".rb":   "ENV['SECRET_KEY']",
            ".php":  "$_ENV['SECRET_KEY']",
            ".cs":   'Environment.GetEnvironmentVariable("SECRET_KEY")',
            ".kt":   'System.getenv("SECRET_KEY")',
        },
    ),
    "cbc_no_auth": (
        "AES-GCM (authenticated encryption)",
        {
            ".py":   "Cipher(algorithms.AES(key), modes.GCM(iv))",
            ".java": 'Cipher.getInstance("AES/GCM/NoPadding")',
            ".go":   "cipher.NewGCM(block)",
            ".cs":   "AesGcm(key)",
        },
    ),
    "pbkdf2_weak": (
        "310,000+ iterations (NIST SP 800-132) or Argon2",
        {
            ".py":   "argon2.PasswordHasher() or PBKDF2HMAC(iterations=310_000)",
            ".java": "new PBKDF2WithHmacSHA256(310_000, 256)",
            ".js":   "crypto.pbkdf2(pass, salt, 310000, 32, 'sha256', cb)",
            ".ts":   "crypto.pbkdf2(pass, salt, 310000, 32, 'sha256', cb)",
            ".cs":   "new Rfc2898DeriveBytes(pass, salt, 310_000, HashAlgorithmName.SHA256)",
        },
    ),
    "blowfish": (
        "AES-256-GCM",
        {},
    ),
}

_KEYWORD_CLASSIFY = [
    (["ecdlp", "ecdsa", "ecdh", "secp256k1", "curve25519", "elliptic curve"], "ecc"),
    (["rsa key size", "rsa is broken", "crypt::rsa", "rsa public-key"], "rsa"),
    (["diffie-hellman", "dh/dsa", "digital signature algorithm", "dsa is broken"], "dh_dsa"),
    (["md4"], "md4"),
    (["md5"], "md5"),
    (["sha-1", "sha1"], "sha1"),
    (["3des", "triple des", "tripledes", "desede"], "3des"),
    (["des"], "des"),
    (["rc4", "arc4"], "rc4"),
    (["blowfish"], "blowfish"),
    (["aes-128", "aes128", "aes-192", "aes192", "128-bit key"], "aes128"),
    (["ecb", "electronic codebook"], "ecb"),
    (["tls 1.0", "tls 1.1", "sslv3", "old tls", "rfc 8996"], "tls_weak"),
    (["non-cryptographic prng", "math/rand", "math.random", "java.util.random",
      "predictable", "weak random"], "weak_random"),
    (["static iv", "hardcoded iv", "hardcoded_iv", "fresh random iv"], "static_iv"),
    (["hardcoded", "hard-coded", "source code exposes"], "hardcoded_key"),
    (["padding oracle", "cbc without", "unauthenticated"], "cbc_no_auth"),
    (["pbkdf2", "iterations"], "pbkdf2_weak"),
    (["rsa"], "rsa"),
]

def _classify_message(message: str) -> Optional[str]:
    lower = message.lower()
    for keywords, key in _KEYWORD_CLASSIFY:
        if any(kw in lower for kw in keywords):
            return key
    return None


def _get_fix_hint(message: str, file_ext: str) -> Optional[str]:
    """Returns a language-specific fix hint string, or None."""
    key = _classify_message(message)
    if not key or key not in _VULN_HINTS:
        return None
    generic, lang_map = _VULN_HINTS[key]
    return lang_map.get(file_ext) or generic


def _build_comment_lines(
    comment_prefix: str,
    message: str,
    fix_hint: Optional[str],
    indentation: str,
) -> list[str]:
    """Build the full comment block as a list of lines (with newlines)."""
    # Wrap message at ~100 chars
    max_w = 100 - len(indentation) - len(comment_prefix) - 1
    words, lines, cur = message.split(), [], ""
    for w in words:
        if cur and len(cur) + 1 + len(w) > max_w:
            lines.append(cur)
            cur = w
        else:
            cur = (cur + " " + w).strip()
    if cur:
        lines.append(cur)

    result = []
    if comment_prefix == "<!--":
        body = f"@TODO-pq {message}"
        if fix_hint:
            body += f" | Fix: {fix_hint}"
        result.append(f"{indentation}<!-- {body} -->\n")
    elif comment_prefix == "/*":
        result.append(f"{indentation}/* @TODO-pq {lines[0]}\n")
        for l in lines[1:]:
            result.append(f"{indentation} *  {l}\n")
        if fix_hint:
            result.append(f"{indentation} *  Fix: {fix_hint}\n")
        result.append(f"{indentation} */\n")
    else:
        p = comment_prefix
        result.append(f"{indentation}{p} @TODO-pq {lines[0]}\n")
        for l in lines[1:]:
            result.append(f"{indentation}{p}          {l}\n")
        if fix_hint:
            result.append(f"{indentation}{p}          Fix: {fix_hint}\n")

    return result


# Language-specific comment syntax mapping
COMMENT_SYNTAX = {
    # Hash-based comments
    '.py': '#',
    '.rb': '#',
    '.sh': '#',
    '.bash': '#',
    '.yml': '#',
    '.yaml': '#',
    '.env': '#',
    '.envrc': '#',
    '.bashrc': '#',
    '.zshrc': '#',
    '.profile': '#',
    '.conf': '#',
    '.cfg': '#',
    '.ini': '#',
    '.toml': '#',
    '.properties': '#',
    '.tf': '#',       # Terraform
    '.tfvars': '#',   # Terraform variables
    '.hcl': '#',      # HashiCorp config (Vault, Consul, Nomad)
    '.ps1': '#',      # PowerShell
    '.psm1': '#',     # PowerShell module
    '.psd1': '#',     # PowerShell data
    '.pl': '#',       # Perl
    '.pm': '#',       # Perl module
    '.r': '#',        # R language
    '.R': '#',        # R language
    '.tcl': '#',      # Tcl
    '.awk': '#',      # AWK
    '.cmake': '#',    # CMake
    '.mk': '#',       # Makefile include

    # Double-slash comments
    '.js': '//',
    '.ts': '//',
    '.jsx': '//',
    '.tsx': '//',
    '.java': '//',
    '.c': '//',
    '.cpp': '//',
    '.cc': '//',
    '.h': '//',
    '.hpp': '//',
    '.cs': '//',
    '.go': '//',
    '.swift': '//',
    '.kt': '//',
    '.rs': '//',
    '.php': '//',
    '.dart': '//',    # Flutter/Dart
    '.groovy': '//',  # Gradle/Jenkins
    '.gradle': '//',  # Gradle build files
    '.scala': '//',   # Scala
    '.m': '//',       # Objective-C (not MATLAB)
    '.mm': '//',      # Objective-C++

    # Other formats
    '.html': '<!--',
    '.xml': '<!--',
    '.xaml': '<!--',  # WPF/MAUI
    '.css': '/*',
    '.scss': '//',
    '.sql': '--',
    '.lua': '--',     # Lua
    '.vb': "'",       # Visual Basic
    '.vbs': "'",      # VBScript
    '.matlab': '%',   # MATLAB
}


def get_comment_syntax(file_path: str) -> str:
    """Get the comment syntax for a file based on its extension.

    Args:
        file_path: Path to the file

    Returns:
        Comment prefix string (e.g., '#', '//', '<!--')
        Defaults to '//' if extension not recognized
    """
    ext = Path(file_path).suffix.lower()
    # Also handle extensionless files by name (e.g. .envrc, Makefile)
    name = Path(file_path).name.lower()
    _hash_names = {'.envrc', '.bashrc', '.zshrc', '.profile', '.npmrc', '.gemrc',
                   'makefile', 'dockerfile', 'vagrantfile', 'brewfile', 'procfile',
                   'jenkinsfile', 'rakefile', 'gemfile', 'podfile'}
    if not ext and name in _hash_names:
        return '#'
    return COMMENT_SYNTAX.get(ext, '#')


def generate_comment_patch(
    file_path: str,
    line_number: int,
    message: str,
    file_content: Optional[str] = None
) -> Optional[str]:
    """Generate a Git unified diff patch that adds a @TODO-pq comment.

    Args:
        file_path: Absolute path to the file
        line_number: Line number where the vulnerability was detected (1-indexed)
        message: Security message to include in the comment
        file_content: Optional file content (if not provided, will read from file_path)

    Returns:
        Git unified diff patch string, or None if generation fails

    Example:
        >>> generate_comment_patch('/path/to/file.py', 42, 'Use SHA256 instead of MD5')
        # Returns a patch that adds: # @TODO-pq Use SHA256 instead of MD5
    """
    try:
        # Read file content if not provided
        if file_content is None:
            with open(file_path, 'r', encoding='utf-8') as f:
                file_content = f.read()

        lines = file_content.splitlines(keepends=True)

        # Validate line number
        if line_number < 1 or line_number > len(lines):
            logger.error(f"Line number {line_number} out of range for {file_path}")
            return None

        # Get comment syntax and fix hint for this file type
        comment_prefix = get_comment_syntax(file_path)
        file_ext = Path(file_path).suffix.lower()
        fix_hint = _get_fix_hint(message, file_ext)

        # Determine indentation from the vulnerable line
        vulnerable_line = lines[line_number - 1]
        indentation = ' ' * (len(vulnerable_line) - len(vulnerable_line.lstrip()))

        # Build multi-line comment block
        comment_lines = _build_comment_lines(comment_prefix, message, fix_hint, indentation)

        # Build the patch
        # Context: 3 lines before and 3 lines after the insertion point
        context_before = max(0, line_number - 4)
        context_after = min(len(lines), line_number + 3)

        old_start = context_before + 1
        old_count = context_after - context_before
        new_start = old_start
        new_count = old_count + len(comment_lines)

        # Sanitize file_path for use in diff headers (prevent injection)
        safe_path = file_path.replace('\n', '').replace('\r', '').replace('\t', '')

        diff_lines = [
            f"--- a/{safe_path}\n",
            f"+++ b/{safe_path}\n",
            f"@@ -{old_start},{old_count} +{new_start},{new_count} @@\n"
        ]

        # Add context lines before
        for i in range(context_before, line_number - 1):
            diff_lines.append(' ' + lines[i])

        # Add the new comment lines
        for cl in comment_lines:
            diff_lines.append('+' + cl)

        # Add the vulnerable line and context after
        for i in range(line_number - 1, context_after):
            diff_lines.append(' ' + lines[i])

        patch = ''.join(diff_lines)
        return patch

    except FileNotFoundError:
        logger.error(f"File not found: {file_path}")
        return None
    except Exception as e:
        logger.exception(f"Failed to generate comment patch for {file_path}:{line_number}")
        return None


def generate_ai_comment_patch(
    file_path: str,
    line_number: int,
    ai_analysis: str,
    file_content: Optional[str] = None
) -> Optional[str]:
    """Generate a Git unified diff patch with AI analysis as a @TODO-pq comment.

    This is similar to generate_comment_patch but specifically for AI-generated
    analysis/recommendations.

    Args:
        file_path: Absolute path to the file
        line_number: Line number where the issue was detected (1-indexed)
        ai_analysis: AI-generated analysis or recommendation
        file_content: Optional file content (if not provided, will read from file_path)

    Returns:
        Git unified diff patch string, or None if generation fails
    """
    # For now, AI comments use the same format as scanner comments
    # The distinction is just in the message content
    return generate_comment_patch(file_path, line_number, ai_analysis, file_content)
