# ---------------------------------------------------------------------------
#  Copyright (c) DevSecCode. All rights reserved.
# ---------------------------------------------------------------------------
"""Prove -> Map -> Fix pipeline — the full loop that Shannon can't do.

Shannon stops at "here's a PoC". Deva does:
1. PROVE  — exploit confirmed the vulnerability is real (already done by runner)
2. MAP    — map the proven CWEs to compliance controls (HIPAA, PCI-DSS, SOC 2, etc.)
3. FIX    — use LLM to generate a code fix based on the CPG taint path
4. VERIFY — re-run the exploit to confirm the fix works

This pipeline operates on a list of ProofOfConcept objects and enriches them
with compliance mapping and fix suggestions, then optionally verifies fixes.
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any, Optional

from dsc_api.services.exploit_engine.poc import ProofOfConcept, PocReport
from dsc_api.services.exploit_engine.runner import ExploitResult, ExploitRunner
from dsc_api.services.exploit_engine.planner import ExploitStrategy

logger = logging.getLogger(__name__)


@dataclass
class FixSuggestion:
    """An AI-generated code fix for a proven vulnerability."""
    file_path: str
    start_line: int
    end_line: int
    original_code: str
    fixed_code: str
    explanation: str
    fix_type: str  # "parameterize", "escape", "validate", "sanitize", "encode", "restrict"


@dataclass
class VerificationResult:
    """Result of re-running an exploit after applying a fix."""
    poc_id: str
    fix_applied: bool
    exploit_still_works: bool  # False = fix is effective
    details: str = ""


@dataclass
class ProveMapFixResult:
    """Complete result of the pipeline for one PoC."""
    poc: ProofOfConcept
    compliance_controls: list[dict[str, str]]
    fix_suggestions: list[FixSuggestion]
    verification: Optional[VerificationResult] = None


class ProveMapFixPipeline:
    """Orchestrates the full Prove -> Map -> Fix -> Verify loop.

    Usage:
        pipeline = ProveMapFixPipeline(workspace_root, llm_client)
        result = await pipeline.process(poc)
        print(f"Controls impacted: {len(result.compliance_controls)}")
        print(f"Fix: {result.fix_suggestions[0].explanation}")
    """

    def __init__(
        self,
        workspace_root: str,
        llm_client: Any = None,
        compliance_engine: Any = None,
    ) -> None:
        self._workspace_root = workspace_root
        self._llm_client = llm_client
        self._compliance_engine = compliance_engine

    async def process(self, poc: ProofOfConcept) -> ProveMapFixResult:
        """Run the full pipeline for a single PoC."""
        # Step 1: MAP — compliance control mapping
        controls = self._map_compliance(poc)
        poc.compliance_controls = controls

        # Step 2: FIX — generate code fix
        fixes = await self._generate_fix(poc)
        if fixes:
            poc.fix_suggestion = fixes[0].explanation
            poc.fix_code = fixes[0].fixed_code

        return ProveMapFixResult(
            poc=poc,
            compliance_controls=controls,
            fix_suggestions=fixes,
        )

    async def process_report(self, report: PocReport) -> list[ProveMapFixResult]:
        """Run the pipeline for all PoCs in a report."""
        results: list[ProveMapFixResult] = []
        for poc in report.pocs:
            result = await self.process(poc)
            results.append(result)
        return results

    # ------------------------------------------------------------------
    # Step 2: MAP — Compliance control mapping
    # ------------------------------------------------------------------

    def _map_compliance(self, poc: ProofOfConcept) -> list[dict[str, str]]:
        """Map proven vulnerability CWEs to compliance framework controls."""
        controls: list[dict[str, str]] = []

        # Use compliance engine if available
        if self._compliance_engine:
            try:
                for cwe_id in poc.cwes:
                    mapped = self._compliance_engine.get_controls_for_cwe(cwe_id)
                    for ctrl in mapped:
                        controls.append({
                            "framework": ctrl.get("framework", ""),
                            "control_id": ctrl.get("control_id", ""),
                            "title": ctrl.get("title", ""),
                        })
                return controls
            except Exception:
                pass

        # Fallback: built-in CWE -> compliance mapping
        for cwe_id in poc.cwes:
            controls.extend(_BUILT_IN_CWE_CONTROLS.get(cwe_id, []))

        return controls

    # ------------------------------------------------------------------
    # Step 3: FIX — AI-generated code fix
    # ------------------------------------------------------------------

    async def _generate_fix(self, poc: ProofOfConcept) -> list[FixSuggestion]:
        """Generate a code fix for the proven vulnerability."""
        if not poc.taint_sink:
            return []

        # Parse sink location
        sink_parts = poc.taint_sink.split(" (")
        if len(sink_parts) < 2:
            return []

        sink_name = sink_parts[0]
        location = sink_parts[1].rstrip(")")
        loc_parts = location.rsplit(":", 1)
        if len(loc_parts) < 2:
            return []

        file_path = loc_parts[0]
        try:
            line = int(loc_parts[1])
        except ValueError:
            return []

        # Read the vulnerable code
        abs_path = os.path.join(self._workspace_root, file_path)
        try:
            with open(abs_path, "r", errors="replace") as f:
                lines = f.readlines()
        except (OSError, IOError):
            return []

        start = max(0, line - 5)
        end = min(len(lines), line + 10)
        original_code = "".join(lines[start:end])

        if self._llm_client:
            return await self._llm_fix(poc, file_path, line, original_code)

        # Fallback: template-based fix suggestions
        return self._template_fix(poc, file_path, line, original_code)

    async def _llm_fix(
        self, poc: ProofOfConcept, file_path: str, line: int, original_code: str,
    ) -> list[FixSuggestion]:
        """Use LLM to generate a context-aware fix."""
        prompt = f"""Fix this {poc.category.value.replace('_', ' ')} vulnerability.

## Vulnerability
- Category: {poc.category.value}
- CWEs: {', '.join(poc.cwes)}
- Endpoint: {poc.endpoint}
- Proven payload: {poc.payload[:200]}

## Data Flow
Source: {poc.taint_source}
{"-> ".join(poc.taint_hops) + " -> " if poc.taint_hops else ""}Sink: {poc.taint_sink}

## Vulnerable Code ({file_path}:{line})
```
{original_code}
```

Generate a fixed version of this code that:
1. Prevents the vulnerability while preserving functionality
2. Uses the language's standard security practices (parameterized queries, escaping, validation)
3. Is minimal — only change what's necessary

Return JSON:
{{
  "fixed_code": "the corrected code",
  "explanation": "what was changed and why",
  "fix_type": "parameterize|escape|validate|sanitize|encode|restrict"
}}"""

        try:
            response = await self._llm_client.generate_text(
                prompt=prompt,
                system="You are a security engineer fixing proven vulnerabilities. "
                       "Return ONLY valid JSON with the fix.",
                max_tokens=1500,
                temperature=0.2,
            )

            import re
            # Try to extract JSON
            data = None
            try:
                data = json.loads(response)
            except json.JSONDecodeError:
                match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', response, re.DOTALL)
                if match:
                    data = json.loads(match.group(1))

            if data:
                return [FixSuggestion(
                    file_path=file_path,
                    start_line=max(1, line - 5),
                    end_line=line + 10,
                    original_code=original_code,
                    fixed_code=data.get("fixed_code", ""),
                    explanation=data.get("explanation", ""),
                    fix_type=data.get("fix_type", "sanitize"),
                )]
        except Exception:
            logger.debug("LLM fix generation failed", exc_info=True)

        return self._template_fix(poc, file_path, line, original_code)

    @staticmethod
    def _template_fix(
        poc: ProofOfConcept, file_path: str, line: int, original_code: str,
    ) -> list[FixSuggestion]:
        """Template-based fix suggestions when LLM is unavailable."""
        from dsc_api.services.exploit_engine.planner import VulnCategory

        fix_templates: dict[VulnCategory, tuple[str, str]] = {
            VulnCategory.SQL_INJECTION: (
                "Use parameterized queries instead of string interpolation. "
                "Replace raw SQL string construction with prepared statements using placeholders (?, $1, %s).",
                "parameterize",
            ),
            VulnCategory.COMMAND_INJECTION: (
                "Avoid passing user input to shell commands. Use subprocess with argument lists "
                "instead of shell=True. Validate input against an allowlist.",
                "validate",
            ),
            VulnCategory.XSS: (
                "Escape user input before rendering in HTML. Use framework's auto-escaping "
                "(React JSX, template literals with encoding). Never use innerHTML with user data.",
                "escape",
            ),
            VulnCategory.SSRF: (
                "Validate and restrict URLs to an allowlist of safe domains. Block requests to "
                "private IPs (10.x, 172.16.x, 192.168.x, 127.x, 169.254.x). Use URL parsing.",
                "restrict",
            ),
            VulnCategory.PATH_TRAVERSAL: (
                "Resolve the path with path.resolve() and verify it stays within the expected "
                "base directory. Reject paths containing '..' segments.",
                "validate",
            ),
            VulnCategory.NOSQL_INJECTION: (
                "Sanitize user input before using in MongoDB queries. Cast to expected types "
                "(String, Number). Reject objects as query parameter values.",
                "sanitize",
            ),
            VulnCategory.XXE: (
                "Disable external entity processing in XML parser. Set noent: false, "
                "expand_entities: false, or use defusedxml.",
                "restrict",
            ),
            VulnCategory.DESERIALIZATION: (
                "Never deserialize untrusted data with pickle/marshal. Use JSON for data exchange. "
                "If YAML is required, use yaml.safe_load() instead of yaml.load().",
                "restrict",
            ),
            VulnCategory.OPEN_REDIRECT: (
                "Validate redirect URLs against an allowlist of safe domains. Use relative paths "
                "for internal redirects. Reject URLs with external hostnames.",
                "validate",
            ),
        }

        explanation, fix_type = fix_templates.get(
            poc.category,
            ("Review the data flow path and add appropriate input validation.", "validate"),
        )

        return [FixSuggestion(
            file_path=file_path,
            start_line=max(1, line - 5),
            end_line=line + 10,
            original_code=original_code,
            fixed_code="",  # Template fixes don't generate code
            explanation=explanation,
            fix_type=fix_type,
        )]


# ---------------------------------------------------------------------------
# Built-in CWE -> Compliance control mapping (used when compliance engine unavailable)
# ---------------------------------------------------------------------------

_BUILT_IN_CWE_CONTROLS: dict[str, list[dict[str, str]]] = {
    "CWE-89": [
        {"framework": "PCI-DSS", "control_id": "6.5.1", "title": "Injection flaws"},
        {"framework": "HIPAA", "control_id": "164.312(a)(2)(iv)", "title": "Encryption and decryption"},
        {"framework": "SOC2", "control_id": "CC6.1", "title": "Logical and physical access controls"},
        {"framework": "NIST 800-53", "control_id": "SI-10", "title": "Information Input Validation"},
        {"framework": "OWASP", "control_id": "A03:2021", "title": "Injection"},
    ],
    "CWE-78": [
        {"framework": "PCI-DSS", "control_id": "6.5.1", "title": "Injection flaws"},
        {"framework": "NIST 800-53", "control_id": "SI-10", "title": "Information Input Validation"},
        {"framework": "OWASP", "control_id": "A03:2021", "title": "Injection"},
    ],
    "CWE-79": [
        {"framework": "PCI-DSS", "control_id": "6.5.7", "title": "Cross-site scripting (XSS)"},
        {"framework": "NIST 800-53", "control_id": "SI-10", "title": "Information Input Validation"},
        {"framework": "OWASP", "control_id": "A03:2021", "title": "Injection"},
    ],
    "CWE-918": [
        {"framework": "NIST 800-53", "control_id": "SC-7", "title": "Boundary Protection"},
        {"framework": "OWASP", "control_id": "A10:2021", "title": "Server-Side Request Forgery (SSRF)"},
        {"framework": "SOC2", "control_id": "CC6.6", "title": "Restricting access"},
    ],
    "CWE-22": [
        {"framework": "PCI-DSS", "control_id": "6.5.8", "title": "Improper access control"},
        {"framework": "NIST 800-53", "control_id": "AC-6", "title": "Least Privilege"},
        {"framework": "OWASP", "control_id": "A01:2021", "title": "Broken Access Control"},
    ],
    "CWE-502": [
        {"framework": "NIST 800-53", "control_id": "SI-10", "title": "Information Input Validation"},
        {"framework": "OWASP", "control_id": "A08:2021", "title": "Software and Data Integrity Failures"},
    ],
    "CWE-611": [
        {"framework": "NIST 800-53", "control_id": "SI-10", "title": "Information Input Validation"},
        {"framework": "OWASP", "control_id": "A05:2021", "title": "Security Misconfiguration"},
    ],
    "CWE-943": [
        {"framework": "PCI-DSS", "control_id": "6.5.1", "title": "Injection flaws"},
        {"framework": "OWASP", "control_id": "A03:2021", "title": "Injection"},
    ],
    "CWE-601": [
        {"framework": "OWASP", "control_id": "A01:2021", "title": "Broken Access Control"},
        {"framework": "NIST 800-53", "control_id": "SI-10", "title": "Information Input Validation"},
    ],
    "CWE-95": [
        {"framework": "PCI-DSS", "control_id": "6.5.1", "title": "Injection flaws"},
        {"framework": "OWASP", "control_id": "A03:2021", "title": "Injection"},
    ],
    "CWE-117": [
        {"framework": "NIST 800-53", "control_id": "AU-9", "title": "Protection of Audit Information"},
    ],
}
