"""
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 json
import logging
import os
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()
            visibility = 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
