"""
Unified Security Fix Prompt Builder

Single prompt that handles ALL security vulnerabilities - LLM decides strategy.
Replaces separate PQC and general security prompt builders.
"""

import os
from typing import Any, Dict


def build_unified_fix_prompt(
    file_path: str,
    file_content: str,
    finding: Dict[str, Any],
    context_lines: int = 10
) -> str:
    """
    Build a unified prompt for generating security fix guidance.

    The LLM decides:
    - If it's quantum-vulnerable (RSA, ECDSA, MD5, etc.)
    - What the appropriate fix/replacement is
    - Whether to recommend PQC alternatives

    Args:
        file_path: Path to the file being analyzed
        file_content: Full content of the file
        finding: Scanner finding dictionary
        context_lines: Number of lines of context around vulnerability

    Returns:
        Formatted prompt string for LLM
    """
    # Extract finding location
    start_line = (
        (finding.get("start", {}) or {}).get("line")
        or finding.get("start_line")
        or 1
    )
    try:
        start_line = int(start_line)
    except Exception:
        start_line = 1

    # Extract context around vulnerable code
    lines = file_content.splitlines()
    start_idx = max(0, start_line - 1)
    lo = max(0, start_idx - context_lines)
    hi = min(len(lines), start_idx + context_lines)
    context_snippet = "\n".join(lines[lo:hi])

    # Extract finding metadata
    rule = finding.get("check_id") or (finding.get("extra", {}) or {}).get("rule_id") or "unknown_rule"
    severity = (finding.get("extra", {}) or {}).get("severity") or finding.get("severity") or "UNKNOWN"
    message = (finding.get("extra", {}) or {}).get("message") or finding.get("message") or ""

    # Detect language for comment syntax - only the extension is needed, not the full path
    ext = os.path.splitext(file_path)[1].lower()
    language_comment_map = {
        ".py": "#",
        ".js": "//",
        ".ts": "//",
        ".java": "//",
        ".c": "//",
        ".cpp": "//",
        ".go": "//",
        ".rs": "//",
        ".rb": "#",
        ".php": "//",
        ".sh": "#",
    }
    comment_syntax = language_comment_map.get(ext, "#")

    # Use a generic placeholder - the real path is never sent to the AI.
    # Diff headers are rewritten to the correct path by rewrite_diff_headers() after generation.
    placeholder = f"file{ext}" if ext else "no-extension-file"

    prompt = f"""You are a security expert. Generate a code audit comment for this vulnerability.

## CRITICAL REQUIREMENTS:

1. **Output Format**: Generate ONLY a unified diff for: {placeholder}
   - Use headers: `--- a/{placeholder}` and `+++ b/{placeholder}`
   - Must apply cleanly with `git apply`

2. **What to Generate**: Add a SINGLE-LINE @TODO-pq comment above the vulnerable code
   - DO NOT modify the actual vulnerable code - only add ONE comment line above it
   - Use `{comment_syntax}` for comments (detected from file extension)
   - Tag MUST start with `@TODO-pq` followed by your analysis

3. **Comment Format Example** (using `{placeholder}` as the filename):
   ```
   {comment_syntax} @TODO-pq Replace MD5 with SHA-256. Current code uses MD5 which is quantum-vulnerable. Recommendation: import hashlib; hashlib.sha256(data)
   ```

   IMPORTANT: This must be ONE line, not a multi-line comment block.

4. **Quantum Cryptography Detection** (IMPORTANT):
   - Check if this vulnerability involves quantum-vulnerable algorithms:
     * Asymmetric: RSA, ECDSA, DSA, DH, ECDH, ECC
     * Hash: MD5, SHA-1
     * Symmetric: AES-128, DES, 3DES
   - If quantum-vulnerable:
     * Recommend NIST post-quantum alternatives:
       - RSA/ECDSA/DSA -> ML-DSA (FIPS 204) for signatures
       - RSA/DH/ECDH -> ML-KEM (FIPS 203) for key exchange
       - MD5/SHA-1 -> SHA-3-512 or SHA-512
       - AES-128 -> AES-256
     * Explain why it's quantum-vulnerable (Shor's/Grover's algorithm)
   - Otherwise:
     * Provide standard security best practice fix

5. **DO NOT**:
   - Modify or comment out existing code
   - Add multi-line comment blocks
   - Include anything outside unified diff format
   - Add explanations outside the single @TODO-pq comment

## Vulnerability Details:

- **Rule**: {rule}
- **Severity**: {severity}
- **Message**: {message}
- **Location**: Line {start_line}

## Code Context (+/-{context_lines} lines around vulnerability at line {start_line}):

```code
{context_snippet}
```

## Your Task:

Generate a unified diff that adds ONE single-line @TODO-pq comment above the vulnerable code.
Use `{placeholder}` as the filename in the diff headers.

Format: {comment_syntax} @TODO-pq [concise description of issue + recommended fix]

Include in the comment:
1. What the issue is (e.g., "MD5 is quantum-vulnerable")
2. What to use instead (e.g., "Use SHA-256")
3. Brief code example if helpful

Keep it ONE LINE, SHORT, and ACTIONABLE.
"""

    return prompt


def build_code_gen_fix_prompt(
    file_path: str,
    file_content: str,
    finding: Dict[str, Any],
    context_lines: int = 30
) -> str:
    """
    Build a prompt for generating actual code replacement (not just a comment annotation).

    The LLM generates a unified diff that REPLACES the vulnerable code with a
    PQC-safe or security-hardened implementation.
    """
    start_line = (
        (finding.get("start", {}) or {}).get("line")
        or finding.get("start_line")
        or 1
    )
    try:
        start_line = int(start_line)
    except Exception:
        start_line = 1

    lines = file_content.splitlines()
    start_idx = max(0, start_line - 1)
    lo = max(0, start_idx - context_lines)
    hi = min(len(lines), start_idx + context_lines)
    context_snippet = "\n".join(lines[lo:hi])

    rule = finding.get("check_id") or (finding.get("extra", {}) or {}).get("rule_id") or "unknown_rule"
    severity = (finding.get("extra", {}) or {}).get("severity") or finding.get("severity") or "UNKNOWN"
    message = (finding.get("extra", {}) or {}).get("message") or finding.get("message") or ""

    ext = os.path.splitext(file_path)[1].lower()
    language_comment_map = {
        ".py": "#",
        ".js": "//",
        ".ts": "//",
        ".java": "//",
        ".c": "//",
        ".cpp": "//",
        ".go": "//",
        ".rs": "//",
        ".rb": "#",
        ".php": "//",
        ".sh": "#",
    }
    comment_syntax = language_comment_map.get(ext, "#")
    placeholder = f"file{ext}" if ext else "no-extension-file"

    prompt = f"""You are a senior security engineer specialising in post-quantum cryptography migration.
Generate a CODE FIX (actual replacement code) for the vulnerability below.

## CRITICAL REQUIREMENTS:

1. **Output Format**: Generate ONLY a unified diff for file: {placeholder}
   - Use headers: `--- a/{placeholder}` and `+++ b/{placeholder}`
   - Must apply cleanly with `git apply`

2. **What to Generate**: REPLACE the vulnerable code with a secure, PQC-safe implementation
   - First add ONE single-line `@TODO-pq REPLACED:` comment above the new code (see format below)
   - Then add the replacement code lines
   - Remove the original insecure lines
   - Add any required import statements at the top of the file if missing
   - Keep all surrounding code intact

3. **Migration comment format** (add this line immediately above the new code):
   ```
   {comment_syntax} @TODO-pq REPLACED: [old algorithm] -> [new algorithm] - [one-sentence reason]
   ```
   Example: `{comment_syntax} @TODO-pq REPLACED: RSA-2048 -> ML-KEM-768 (FIPS 203) - quantum-safe key encapsulation`

4. **Quantum-vulnerable algorithms** - use these replacements:
   - RSA / ECDSA / DSA / DH / ECDH / ECC -> ML-KEM (FIPS 203) for key exchange; ML-DSA (FIPS 204) for signatures
   - MD5 / SHA-1 -> SHA-3-512 or SHA-512
   - AES-128 / DES / 3DES -> AES-256
   - If no PQC library is available, use the strongest classical alternative

5. **DO NOT**:
   - Add any comments other than the single @TODO-pq REPLACED line
   - Include explanations outside the diff
   - Change unrelated code

## Vulnerability Details:

- **Rule**: {rule}
- **Severity**: {severity}
- **Message**: {message}
- **Location**: Line {start_line}

## Code Context (+/-{context_lines} lines around vulnerability at line {start_line}):

```code
{context_snippet}
```

## Your Task:

Generate a unified diff that:
1. Adds ONE `{comment_syntax} @TODO-pq REPLACED: ...` comment above the fixed code
2. Replaces the vulnerable code at line {start_line} with a PQC-safe implementation

Use `{placeholder}` in the diff headers. Output ONLY the unified diff, nothing else.
"""
    return prompt
