"""
BREACH.AI - Autonomous Hacker Brain

AI-powered autonomous attack engine that:
1. Analyzes targets and decides what to attack
2. Generates custom attack scripts
3. Executes scripts and learns from results
4. Chains multiple attacks together
5. Adapts when attacks fail
"""

import asyncio
import json
import time
import os
import subprocess
import tempfile
from dataclasses import dataclass, field
from typing import Optional, Any, Callable
from enum import Enum
from datetime import datetime

from breach.utils.logger import logger


class ScriptLanguage(Enum):
    """Languages the agent can write."""
    PYTHON = "python"
    BASH = "bash"
    JAVASCRIPT = "javascript"


class AttackPhase(Enum):
    """Phases of the autonomous attack."""
    RECONNAISSANCE = "reconnaissance"
    VULNERABILITY_SCAN = "vulnerability_scan"
    EXPLOITATION = "exploitation"
    POST_EXPLOITATION = "post_exploitation"
    DATA_EXFILTRATION = "data_exfiltration"
    REPORTING = "reporting"


@dataclass
class GeneratedScript:
    """A script generated by the AI."""
    id: str
    language: ScriptLanguage
    code: str
    purpose: str
    target: str
    phase: AttackPhase
    created_at: datetime = field(default_factory=datetime.now)

    # Execution
    executed: bool = False
    execution_time: float = 0.0
    success: bool = False
    output: str = ""
    error: str = ""

    # Learning
    findings: list = field(default_factory=list)
    follow_up_actions: list = field(default_factory=list)


@dataclass
class AttackState:
    """Current state of the autonomous attack."""
    target: str
    start_time: datetime
    current_phase: AttackPhase

    # Progress
    scripts_generated: int = 0
    scripts_executed: int = 0
    scripts_successful: int = 0

    # Findings
    vulnerabilities: list = field(default_factory=list)
    credentials: list = field(default_factory=list)
    access_gained: list = field(default_factory=list)
    data_extracted: list = field(default_factory=list)

    # Attack surface
    subdomains: list = field(default_factory=list)
    endpoints: list = field(default_factory=list)
    parameters: list = field(default_factory=list)
    technologies: list = field(default_factory=list)

    # Memory
    failed_attacks: list = field(default_factory=list)
    successful_attacks: list = field(default_factory=list)
    scripts_history: list = field(default_factory=list)

    # Control
    should_continue: bool = True
    max_runtime_hours: float = 24.0


# System prompts for the AI
HACKER_BRAIN_SYSTEM_PROMPT = """You are the autonomous security testing brain of BREACH.AI.

Your mission is to thoroughly test the security of authorized targets.

CURRENT STATE:
Target: {target}
Phase: {phase}
Time Elapsed: {elapsed_hours:.2f} hours
Scripts Generated: {scripts_generated}
Scripts Executed: {scripts_executed}
Vulnerabilities Found: {vuln_count}
Credentials Found: {cred_count}
Access Gained: {access_gained}

WHAT YOU KNOW:
Subdomains: {subdomains}
Endpoints: {endpoints}
Technologies: {technologies}

RECENT FINDINGS:
{recent_findings}

FAILED ATTACKS (don't repeat these):
{failed_attacks}

YOUR CAPABILITIES:
1. Write Python scripts for web attacks, API testing, data extraction
2. Write Bash scripts for system enumeration, networking
3. Write JavaScript for browser-based testing

RULES:
1. Generate executable scripts with error handling
2. Output findings in JSON format
3. Chain attacks - use findings from one in the next
4. If one approach fails, try another
5. Extract evidence to prove impact"""


SCRIPT_GENERATION_PROMPT = """Generate an attack script for:

TARGET: {target}
GOAL: {goal}
CONTEXT: {context}

Write a complete, executable {language} script that:
1. Attempts to achieve the goal
2. Handles errors gracefully
3. Outputs results in JSON format
4. Can be run standalone

Return ONLY the script code, no explanations."""


ANALYSIS_PROMPT = """Analyze the results:

SCRIPT PURPOSE: {purpose}
TARGET: {target}
OUTPUT:
{output}

ERROR (if any):
{error}

Return analysis as JSON:
{{
    "success": true/false,
    "vulnerabilities": ["vuln1", "vuln2"],
    "data_extracted": ["data1", "data2"],
    "access_gained": "none|user|admin|root",
    "next_actions": ["action1", "action2", "action3"],
    "reasoning": "why you think this"
}}"""


class AutonomousHackerBrain:
    """
    The autonomous security testing engine.

    This brain:
    - Decides what to test next
    - Writes custom scripts
    - Executes them safely
    - Learns from results
    - Adapts and persists
    """

    def __init__(
        self,
        api_key: str = None,
        model: str = "claude-sonnet-4-20250514",
        max_runtime_hours: float = 24.0,
    ):
        self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
        self.model = model
        self.max_runtime_hours = max_runtime_hours

        # Lazy load client
        self._client = None

        # State
        self.state: Optional[AttackState] = None

        # Script execution sandbox
        self.sandbox_dir = tempfile.mkdtemp(prefix="breach_autonomous_")

        # Callbacks for UI updates
        self.on_script_generated: Optional[Callable] = None
        self.on_script_executed: Optional[Callable] = None
        self.on_finding: Optional[Callable] = None
        self.on_phase_change: Optional[Callable] = None
        self.on_log: Optional[Callable] = None

    @property
    def client(self):
        """Lazy load Anthropic client."""
        if self._client is None:
            try:
                from anthropic import Anthropic
                self._client = Anthropic(api_key=self.api_key)
            except ImportError:
                logger.error("Anthropic SDK not installed. Run: pip install anthropic")
                raise
        return self._client

    def _log(self, message: str, level: str = "info"):
        """Log message and call callback."""
        if level == "info":
            logger.info(message)
        elif level == "warning":
            logger.warning(message)
        elif level == "error":
            logger.error(message)

        if self.on_log:
            self.on_log(message, level)

    async def run_autonomous_attack(
        self,
        target: str,
        aggressive: bool = True,
    ) -> AttackState:
        """
        Run a fully autonomous security assessment.

        This will run for up to max_runtime_hours, continuously
        testing the target's security.
        """
        self._log(f"Starting autonomous assessment against {target}")
        self._log(f"Maximum runtime: {self.max_runtime_hours} hours")

        # Initialize state
        self.state = AttackState(
            target=target,
            start_time=datetime.now(),
            current_phase=AttackPhase.RECONNAISSANCE,
            max_runtime_hours=self.max_runtime_hours,
        )

        # Main attack loop
        iteration = 0
        while self._should_continue():
            iteration += 1
            self._log(f"\n{'='*60}")
            self._log(f"Iteration {iteration} | Phase: {self.state.current_phase.value}")
            self._log(f"Elapsed: {self._elapsed_hours():.2f}h / {self.max_runtime_hours}h")
            self._log(f"{'='*60}")

            try:
                # 1. Decide what to do next
                action = await self._decide_next_action()

                if not action:
                    self._log("No action decided, advancing phase", "warning")
                    self._advance_phase()
                    continue

                # 2. Generate a script for the action
                script = await self._generate_script(action)

                if not script:
                    self._log("Failed to generate script", "warning")
                    continue

                # 3. Execute the script
                result = await self._execute_script(script)

                # 4. Analyze results
                analysis = await self._analyze_results(script, result)

                # 5. Update state based on analysis
                await self._update_state(script, analysis)

                # 6. Check if we should advance phase
                if self._should_advance_phase():
                    self._advance_phase()

                # Small delay between iterations
                await asyncio.sleep(1)

            except Exception as e:
                self._log(f"Error in attack loop: {e}", "error")
                await asyncio.sleep(5)

        # Final report
        self._log("\n" + "="*60)
        self._log("AUTONOMOUS ASSESSMENT COMPLETE")
        self._log(f"Total runtime: {self._elapsed_hours():.2f} hours")
        self._log(f"Scripts generated: {self.state.scripts_generated}")
        self._log(f"Scripts executed: {self.state.scripts_executed}")
        self._log(f"Vulnerabilities found: {len(self.state.vulnerabilities)}")
        self._log(f"Credentials found: {len(self.state.credentials)}")
        self._log(f"Access gained: {self.state.access_gained}")
        self._log("="*60)

        return self.state

    async def _decide_next_action(self) -> Optional[dict]:
        """Use AI to decide what to do next."""
        prompt = HACKER_BRAIN_SYSTEM_PROMPT.format(
            target=self.state.target,
            phase=self.state.current_phase.value,
            elapsed_hours=self._elapsed_hours(),
            scripts_generated=self.state.scripts_generated,
            scripts_executed=self.state.scripts_executed,
            vuln_count=len(self.state.vulnerabilities),
            cred_count=len(self.state.credentials),
            access_gained=", ".join(self.state.access_gained) or "none",
            subdomains=", ".join(self.state.subdomains[:10]) or "none discovered",
            endpoints=", ".join(self.state.endpoints[:10]) or "none discovered",
            technologies=", ".join(self.state.technologies[:10]) or "unknown",
            recent_findings=self._format_recent_findings(),
            failed_attacks=self._format_failed_attacks(),
        )

        decision_prompt = """Based on the current state, decide what to test next.

Return your decision as JSON:
{
    "action": "short description of what to do",
    "goal": "what we're trying to achieve",
    "language": "python|bash|javascript",
    "priority": 1-10,
    "reasoning": "why this action now"
}"""

        try:
            response = self.client.messages.create(
                model=self.model,
                max_tokens=1000,
                system=prompt,
                messages=[{"role": "user", "content": decision_prompt}]
            )

            content = response.content[0].text

            # Parse JSON from response
            json_start = content.find("{")
            json_end = content.rfind("}") + 1
            if json_start != -1 and json_end > json_start:
                json_str = content[json_start:json_end]
                action = json.loads(json_str)

                self._log(f"Next action: {action.get('action', 'unknown')}")
                self._log(f"Goal: {action.get('goal', 'unknown')}")

                return action

            return None

        except Exception as e:
            self._log(f"Decision failed: {e}", "error")
            return None

    async def _generate_script(self, action: dict) -> Optional[GeneratedScript]:
        """Generate a security testing script based on the decided action."""
        language_str = action.get("language", "python")
        try:
            language = ScriptLanguage(language_str)
        except ValueError:
            language = ScriptLanguage.PYTHON

        prompt = SCRIPT_GENERATION_PROMPT.format(
            target=self.state.target,
            goal=action.get("goal", "test security"),
            context=self._get_context_for_script(),
            language=language.value,
        )

        try:
            response = self.client.messages.create(
                model=self.model,
                max_tokens=4000,
                messages=[{"role": "user", "content": prompt}]
            )

            code = response.content[0].text

            # Clean up code (remove markdown if present)
            if "```" in code:
                lines = code.split("\n")
                in_code_block = False
                code_lines = []
                for line in lines:
                    if line.startswith("```"):
                        in_code_block = not in_code_block
                        continue
                    if in_code_block:
                        code_lines.append(line)
                code = "\n".join(code_lines)

            script = GeneratedScript(
                id=f"script_{self.state.scripts_generated + 1}",
                language=language,
                code=code,
                purpose=action.get("action", "unknown"),
                target=self.state.target,
                phase=self.state.current_phase,
            )

            self.state.scripts_generated += 1
            self.state.scripts_history.append(script)

            if self.on_script_generated:
                self.on_script_generated(script)

            self._log(f"Generated {language.value} script ({len(code)} chars)")

            return script

        except Exception as e:
            self._log(f"Script generation failed: {e}", "error")
            return None

    async def _execute_script(self, script: GeneratedScript) -> dict:
        """Execute a generated script in a sandbox."""
        self._log(f"Executing {script.language.value} script...")

        start_time = time.time()
        result = {"output": "", "error": "", "success": False}

        try:
            if script.language == ScriptLanguage.PYTHON:
                result = await self._execute_python(script.code)
            elif script.language == ScriptLanguage.BASH:
                result = await self._execute_bash(script.code)
            elif script.language == ScriptLanguage.JAVASCRIPT:
                result = await self._execute_javascript(script.code)

            script.execution_time = time.time() - start_time
            script.executed = True
            script.output = result.get("output", "")
            script.error = result.get("error", "")
            script.success = result.get("success", False)

            self.state.scripts_executed += 1
            if script.success:
                self.state.scripts_successful += 1

            if self.on_script_executed:
                self.on_script_executed(script, result)

            self._log(f"Execution complete ({script.execution_time:.2f}s)")

        except Exception as e:
            script.error = str(e)
            script.executed = True
            self._log(f"Execution failed: {e}", "error")

        return result

    async def _execute_python(self, code: str) -> dict:
        """Execute Python code."""
        script_path = os.path.join(self.sandbox_dir, f"script_{time.time()}.py")

        with open(script_path, "w") as f:
            f.write(code)

        try:
            result = subprocess.run(
                ["python3", script_path],
                capture_output=True,
                text=True,
                timeout=300,
                cwd=self.sandbox_dir,
            )

            return {
                "output": result.stdout,
                "error": result.stderr,
                "success": result.returncode == 0,
            }

        except subprocess.TimeoutExpired:
            return {"output": "", "error": "Script timed out after 5 minutes", "success": False}
        except Exception as e:
            return {"output": "", "error": str(e), "success": False}

    async def _execute_bash(self, code: str) -> dict:
        """Execute Bash code."""
        script_path = os.path.join(self.sandbox_dir, f"script_{time.time()}.sh")

        with open(script_path, "w") as f:
            f.write(code)

        os.chmod(script_path, 0o755)

        try:
            result = subprocess.run(
                ["bash", script_path],
                capture_output=True,
                text=True,
                timeout=300,
                cwd=self.sandbox_dir,
            )

            return {
                "output": result.stdout,
                "error": result.stderr,
                "success": result.returncode == 0,
            }

        except subprocess.TimeoutExpired:
            return {"output": "", "error": "Script timed out", "success": False}
        except Exception as e:
            return {"output": "", "error": str(e), "success": False}

    async def _execute_javascript(self, code: str) -> dict:
        """Execute JavaScript code (Node.js)."""
        script_path = os.path.join(self.sandbox_dir, f"script_{time.time()}.js")

        with open(script_path, "w") as f:
            f.write(code)

        try:
            result = subprocess.run(
                ["node", script_path],
                capture_output=True,
                text=True,
                timeout=300,
                cwd=self.sandbox_dir,
            )

            return {
                "output": result.stdout,
                "error": result.stderr,
                "success": result.returncode == 0,
            }

        except subprocess.TimeoutExpired:
            return {"output": "", "error": "Script timed out", "success": False}
        except FileNotFoundError:
            return {"output": "", "error": "Node.js not installed", "success": False}
        except Exception as e:
            return {"output": "", "error": str(e), "success": False}

    async def _analyze_results(self, script: GeneratedScript, result: dict) -> dict:
        """Use AI to analyze execution results."""
        prompt = ANALYSIS_PROMPT.format(
            purpose=script.purpose,
            target=script.target,
            output=result.get("output", "")[:5000],
            error=result.get("error", "")[:1000],
        )

        try:
            response = self.client.messages.create(
                model=self.model,
                max_tokens=1000,
                messages=[{"role": "user", "content": prompt}]
            )

            content = response.content[0].text

            # Parse JSON from response
            json_start = content.find("{")
            json_end = content.rfind("}") + 1
            if json_start != -1 and json_end > json_start:
                json_str = content[json_start:json_end]
                analysis = json.loads(json_str)
                return analysis

            return self._default_analysis()

        except Exception as e:
            self._log(f"Analysis failed: {e}", "error")
            return self._default_analysis()

    def _default_analysis(self) -> dict:
        """Return default analysis when AI fails."""
        return {
            "success": False,
            "vulnerabilities": [],
            "data_extracted": [],
            "access_gained": "none",
            "next_actions": ["Try a different approach"],
        }

    async def _update_state(self, script: GeneratedScript, analysis: dict):
        """Update state based on analysis results."""
        # Add vulnerabilities
        for vuln in analysis.get("vulnerabilities", []):
            if vuln and vuln not in self.state.vulnerabilities:
                self.state.vulnerabilities.append(vuln)
                if self.on_finding:
                    self.on_finding("vulnerability", vuln)
                self._log(f"VULNERABILITY FOUND: {vuln}")

        # Add extracted data
        for data in analysis.get("data_extracted", []):
            if data and data not in self.state.data_extracted:
                self.state.data_extracted.append(data)
                if self.on_finding:
                    self.on_finding("data", data)
                self._log(f"DATA EXTRACTED: {str(data)[:100]}")

        # Update access level
        access = analysis.get("access_gained", "none")
        if access != "none" and access not in self.state.access_gained:
            self.state.access_gained.append(access)
            if self.on_finding:
                self.on_finding("access", access)
            self._log(f"ACCESS GAINED: {access}")

        # Track success/failure
        if analysis.get("success"):
            self.state.successful_attacks.append({
                "script_id": script.id,
                "purpose": script.purpose,
                "findings": analysis,
            })
        else:
            self.state.failed_attacks.append({
                "script_id": script.id,
                "purpose": script.purpose,
                "error": script.error or "Unknown failure",
            })

        script.follow_up_actions = analysis.get("next_actions", [])

    def _should_advance_phase(self) -> bool:
        """Determine if we should advance to the next phase."""
        phase = self.state.current_phase

        if phase == AttackPhase.RECONNAISSANCE:
            return len(self.state.endpoints) >= 5 or len(self.state.subdomains) >= 3

        if phase == AttackPhase.VULNERABILITY_SCAN:
            return len(self.state.vulnerabilities) >= 1

        if phase == AttackPhase.EXPLOITATION:
            return len(self.state.access_gained) >= 1

        if phase == AttackPhase.POST_EXPLOITATION:
            return len(self.state.data_extracted) >= 1

        if phase == AttackPhase.DATA_EXFILTRATION:
            return True

        return False

    def _advance_phase(self):
        """Advance to the next attack phase."""
        phases = list(AttackPhase)
        current_index = phases.index(self.state.current_phase)

        if current_index < len(phases) - 1:
            self.state.current_phase = phases[current_index + 1]
            self._log(f"Advancing to phase: {self.state.current_phase.value}")

            if self.on_phase_change:
                self.on_phase_change(self.state.current_phase)

    def _should_continue(self) -> bool:
        """Determine if the attack should continue."""
        if not self.state.should_continue:
            return False

        if self._elapsed_hours() >= self.max_runtime_hours:
            self._log("Maximum runtime reached")
            return False

        if self.state.current_phase == AttackPhase.REPORTING:
            self._log("Reached reporting phase - assessment complete")
            return False

        return True

    def _elapsed_hours(self) -> float:
        """Get elapsed time in hours."""
        elapsed = datetime.now() - self.state.start_time
        return elapsed.total_seconds() / 3600

    def stop(self):
        """Stop the autonomous attack."""
        self._log("Stopping autonomous assessment...")
        if self.state:
            self.state.should_continue = False

    def _format_recent_findings(self) -> str:
        """Format recent findings for the prompt."""
        findings = []

        for vuln in self.state.vulnerabilities[-5:]:
            findings.append(f"- Vulnerability: {vuln}")

        for data in self.state.data_extracted[-5:]:
            findings.append(f"- Data: {str(data)[:100]}")

        for access in self.state.access_gained[-5:]:
            findings.append(f"- Access: {access}")

        return "\n".join(findings) if findings else "No findings yet"

    def _format_failed_attacks(self) -> str:
        """Format failed attacks for the prompt."""
        failed = []
        for attack in self.state.failed_attacks[-10:]:
            failed.append(f"- {attack.get('purpose', 'unknown')}: {attack.get('error', 'failed')[:50]}")
        return "\n".join(failed) if failed else "No failed attacks yet"

    def _get_context_for_script(self) -> str:
        """Get context for script generation."""
        context = {
            "phase": self.state.current_phase.value,
            "known_endpoints": self.state.endpoints[:20],
            "known_subdomains": self.state.subdomains[:10],
            "known_technologies": self.state.technologies,
            "vulnerabilities_found": self.state.vulnerabilities[:10],
            "access_gained": self.state.access_gained,
        }
        return json.dumps(context, indent=2)


async def run_autonomous_attack(
    target: str,
    max_hours: float = 24.0,
    aggressive: bool = True,
    callbacks: dict = None,
    api_key: str = None,
) -> AttackState:
    """
    Run a fully autonomous security assessment.

    Args:
        target: The target URL/domain
        max_hours: Maximum runtime in hours (default 24)
        aggressive: Whether to use aggressive techniques
        callbacks: Optional callbacks for events
        api_key: Anthropic API key

    Returns:
        AttackState with all findings
    """
    brain = AutonomousHackerBrain(
        api_key=api_key,
        max_runtime_hours=max_hours
    )

    if callbacks:
        brain.on_script_generated = callbacks.get("on_script_generated")
        brain.on_script_executed = callbacks.get("on_script_executed")
        brain.on_finding = callbacks.get("on_finding")
        brain.on_phase_change = callbacks.get("on_phase_change")
        brain.on_log = callbacks.get("on_log")

    return await brain.run_autonomous_attack(target, aggressive)
