"""
Exploit Synthesizer Adapter - Foundry PoC Generator.

Generates Foundry test contracts to confirm/deny vulnerability findings.
Based on SolidiFI approach to automated exploit generation.

Architecture:
1. Takes a finding + contract source
2. Generates a Foundry test that attempts to exploit the vulnerability
3. Runs forge test to verify exploitability
4. Returns confirmation status with PoC code

Author: Fernando Boiero <fboiero@frvm.utn.edu.ar>
Date: 2026-01-31
License: AGPL-3.0
"""

import hashlib
import logging
import re
import shutil
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from src.core.tool_protocol import (
    ToolAdapter,
    ToolCapability,
    ToolCategory,
    ToolMetadata,
    ToolStatus,
)

logger = logging.getLogger(__name__)


POC_TEMPLATES = {
    "reentrancy": """// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";

interface ITarget {{
    {interface_functions}
}}

contract ReentrancyAttacker {{
    ITarget public target;
    uint256 public attackCount;

    constructor(address _target) {{
        target = ITarget(_target);
    }}

    function attack() external payable {{
        target.{vulnerable_function}{{value: msg.value}}({attack_args});
    }}

    receive() external payable {{
        if (attackCount < 3) {{
            attackCount++;
            target.{withdraw_function}({withdraw_args});
        }}
    }}
}}

contract ReentrancyPoCTest is Test {{
    ITarget public target;
    ReentrancyAttacker public attacker;

    function setUp() public {{
        // Deploy target contract
        // target = ITarget(deployCode("Target.sol"));
    }}

    function test_reentrancy_exploit() public {{
        // This test verifies the reentrancy vulnerability
        // If it passes, the vulnerability is confirmed
        vm.deal(address(attacker), 1 ether);
        attacker.attack{{value: 0.1 ether}}();
        assertTrue(address(attacker).balance > 0.1 ether, "Reentrancy exploit failed");
    }}
}}
""",
    "access_control": """// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";

interface ITarget {{
    {interface_functions}
}}

contract AccessControlPoCTest is Test {{
    ITarget public target;
    address public attacker = address(0xBAD);

    function setUp() public {{
        // Deploy target contract
    }}

    function test_unauthorized_access() public {{
        vm.prank(attacker);
        // Call privileged function as non-owner
        // target.{vulnerable_function}({attack_args});
        // If this doesn't revert, access control is broken
    }}
}}
""",
    "integer_overflow": """// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";

interface ITarget {{
    {interface_functions}
}}

contract OverflowPoCTest is Test {{
    ITarget public target;

    function setUp() public {{
        // Deploy target contract
    }}

    function test_overflow_exploit() public {{
        // Test with max uint values
        // target.{vulnerable_function}(type(uint256).max);
    }}
}}
""",
    "default": """// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";

contract GenericPoCTest is Test {{
    function setUp() public {{
        // Deploy target contract at: {contract_path}
    }}

    function test_exploit() public {{
        // Vulnerability: {vuln_type}
        // Location: {file}:{line}
        // Description: {description}
        // TODO: Implement specific exploit for this vulnerability
    }}
}}
""",
}


class ExploitSynthesizerAdapter(ToolAdapter):
    """
    Foundry PoC generator for vulnerability confirmation.

    Generates and optionally runs Foundry tests to confirm findings.
    """

    def __init__(self):
        super().__init__()
        self._poc_dir = Path(tempfile.gettempdir()) / "miesc" / "poc"
        self._poc_dir.mkdir(parents=True, exist_ok=True)

    def get_metadata(self) -> ToolMetadata:
        return ToolMetadata(
            name="exploit_synthesizer",
            version="1.0.0",
            category=ToolCategory.ADVANCED_AI_ENSEMBLE,
            author="Fernando Boiero (based on SolidiFI approach)",
            license="AGPL-3.0",
            homepage="https://github.com/fboiero/MIESC",
            repository="https://github.com/fboiero/MIESC",
            documentation="https://book.getfoundry.sh/",
            installation_cmd="curl -L https://foundry.paradigm.xyz | bash && foundryup",
            capabilities=[
                ToolCapability(
                    name="poc_generation",
                    description="Generate Foundry PoC tests to confirm vulnerabilities",
                    supported_languages=["solidity"],
                    detection_types=[
                        "exploit_confirmed",
                        "exploit_failed",
                        "poc_generated",
                    ],
                ),
            ],
            cost=0.0,
            requires_api_key=False,
            is_optional=True,
        )

    def is_available(self) -> ToolStatus:
        if shutil.which("forge"):
            return ToolStatus.AVAILABLE
        return ToolStatus.AVAILABLE  # Can still generate PoC without running

    def analyze(self, contract_path: str, **kwargs) -> Dict[str, Any]:
        start_time = time.time()
        contract_path = str(Path(contract_path).resolve())
        findings = kwargs.get("findings", [])
        run_poc = kwargs.get("run_poc", False)

        if not findings:
            return {
                "tool": "exploit_synthesizer",
                "version": "1.0.0",
                "status": "success",
                "findings": [],
                "metadata": {"reason": "no findings to synthesize PoC for"},
                "execution_time": time.time() - start_time,
                "error": None,
            }

        try:
            with open(contract_path, "r", errors="ignore") as f:
                source_code = f.read()
        except Exception as e:
            return {
                "tool": "exploit_synthesizer",
                "version": "1.0.0",
                "status": "error",
                "findings": [],
                "metadata": {},
                "execution_time": time.time() - start_time,
                "error": str(e),
            }

        poc_results = []
        for finding in findings[:5]:  # Limit to 5 findings
            poc = self._generate_poc(finding, contract_path, source_code)
            if poc:
                if run_poc and shutil.which("forge"):
                    result = self._run_poc(poc, contract_path)
                    poc["execution_result"] = result
                    poc["confirmed"] = result.get("passed", False)
                poc_results.append(poc)

        synthesized_findings = self.normalize_findings(poc_results)

        return {
            "tool": "exploit_synthesizer",
            "version": "1.0.0",
            "status": "success",
            "findings": synthesized_findings,
            "metadata": {
                "contract": contract_path,
                "findings_processed": len(findings),
                "pocs_generated": len(poc_results),
                "pocs_confirmed": sum(1 for p in poc_results if p.get("confirmed")),
                "poc_directory": str(self._poc_dir),
            },
            "execution_time": time.time() - start_time,
            "error": None,
        }

    def _generate_poc(self, finding: Dict, contract_path: str, source_code: str) -> Optional[Dict]:
        """Generate a PoC test for a finding."""
        vuln_type = self._classify_vuln(finding.get("type", ""))
        template = POC_TEMPLATES.get(vuln_type, POC_TEMPLATES["default"])

        # Extract function info from finding
        func_name = finding.get("location", {}).get("function", "")
        file_path = finding.get("location", {}).get("file", contract_path)
        line = finding.get("location", {}).get("line", 0)
        description = finding.get("description", finding.get("message", ""))

        # Try to extract interface from source
        interface_functions = self._extract_interface(source_code, func_name)

        poc_code = template.format(
            interface_functions=interface_functions or "function placeholder() external;",
            vulnerable_function=func_name or "vulnerableFunction",
            withdraw_function="withdraw",
            attack_args="",
            withdraw_args="",
            contract_path=contract_path,
            vuln_type=vuln_type,
            file=file_path,
            line=line,
            description=description[:200],
        )

        # Save PoC to file
        poc_filename = (
            f"poc_{vuln_type}_{hashlib.md5(f'{file_path}:{line}'.encode()).hexdigest()[:8]}.t.sol"
        )
        poc_path = self._poc_dir / poc_filename

        try:
            poc_path.write_text(poc_code)
        except Exception as e:
            logger.debug(f"Failed to write PoC: {e}")

        return {
            "vuln_type": vuln_type,
            "original_finding": finding,
            "poc_code": poc_code,
            "poc_path": str(poc_path),
            "confirmed": None,  # Unknown until executed
            "file": file_path,
            "line": line,
        }

    def _classify_vuln(self, vuln_type: str) -> str:
        """Classify vulnerability type for template selection."""
        vuln_lower = vuln_type.lower()
        if "reentrancy" in vuln_lower:
            return "reentrancy"
        elif "access" in vuln_lower or "control" in vuln_lower or "auth" in vuln_lower:
            return "access_control"
        elif "overflow" in vuln_lower or "underflow" in vuln_lower:
            return "integer_overflow"
        return "default"

    def _extract_interface(self, source_code: str, func_name: str) -> str:
        """Extract function signatures for interface generation."""
        interfaces = []
        pattern = r"function\s+(\w+)\s*\(([^)]*)\)\s*(external|public)[^;{]*"
        for match in re.finditer(pattern, source_code):
            name = match.group(1)
            params = match.group(2).strip()
            match.group(3)
            # Get return type
            full_match = match.group(0)
            returns = ""
            ret_match = re.search(r"returns\s*\(([^)]+)\)", full_match)
            if ret_match:
                returns = f" returns ({ret_match.group(1)})"
            interfaces.append(f"function {name}({params}) external{returns};")

        return "\n    ".join(interfaces[:10]) if interfaces else ""

    def _run_poc(self, poc: Dict, contract_path: str) -> Dict[str, Any]:
        """Run PoC test with forge."""
        poc_path = poc.get("poc_path", "")
        if not poc_path or not Path(poc_path).exists():
            return {"passed": False, "error": "PoC file not found"}

        try:
            result = subprocess.run(
                ["forge", "test", "--match-path", poc_path, "-vvv"],
                capture_output=True,
                text=True,
                timeout=120,
                cwd=str(Path(contract_path).parent),
            )
            passed = result.returncode == 0
            return {
                "passed": passed,
                "stdout": result.stdout[-1000:] if result.stdout else "",
                "stderr": result.stderr[-500:] if result.stderr else "",
            }
        except subprocess.TimeoutExpired:
            return {"passed": False, "error": "PoC execution timeout"}
        except Exception as e:
            return {"passed": False, "error": str(e)}

    def normalize_findings(self, raw_output: Any) -> List[Dict[str, Any]]:
        if not raw_output:
            return []

        findings = []
        for item in raw_output:
            if not isinstance(item, dict):
                continue

            original = item.get("original_finding", {})
            confirmed = item.get("confirmed")
            vuln_type = item.get("vuln_type", "unknown")

            if confirmed is True:
                severity = "Critical"
                confidence = 0.95
                status_str = "CONFIRMED"
            elif confirmed is False:
                severity = original.get("severity", "Medium")
                confidence = max(0.30, original.get("confidence", 0.50) * 0.5)
                status_str = "NOT_CONFIRMED"
            else:
                severity = original.get("severity", "Medium")
                confidence = original.get("confidence", 0.70)
                status_str = "POC_GENERATED"

            finding_id = hashlib.md5(
                f"exploit:{vuln_type}:{item.get('file', '')}:{item.get('line', 0)}".encode()
            ).hexdigest()[:12]

            findings.append(
                {
                    "id": f"EXP-{finding_id}",
                    "type": f"exploit_{status_str.lower()}",
                    "severity": severity,
                    "confidence": confidence,
                    "location": {
                        "file": item.get("file", original.get("location", {}).get("file", "")),
                        "line": item.get("line", original.get("location", {}).get("line", 0)),
                        "function": original.get("location", {}).get("function", ""),
                    },
                    "message": f"[{status_str}] {original.get('message', original.get('description', vuln_type))}",
                    "description": f"PoC synthesis result: {status_str}. {original.get('description', '')}",
                    "recommendation": original.get("recommendation", "Review generated PoC"),
                    "swc_id": original.get("swc_id"),
                    "cwe_id": original.get("cwe_id"),
                    "tool": "exploit_synthesizer",
                }
            )

        return findings
