"""
Enhanced Trigger System for Pentest MCP Server.
Implements intelligent AI suspension/resumption based on output patterns.
"""

import asyncio
import re
import time
import logging
from typing import Dict, List, Any, Optional, Tuple
from .ssh_manager import SSHManager
from .tmux_manager import TmuxManager

logger = logging.getLogger(__name__)


class TriggerSystem:
    """Manages trigger-based monitoring for AI suspension with enhanced features."""

    # Comprehensive shell prompt patterns
    PROMPT_PATTERNS = [
        r'[\$#>]\s*$',              # Standard bash/sh prompts
        r'root@.*[#\$]\s*$',        # root@hostname#
        r'kali@.*[$]\s*$',          # kali@hostname$
        r'parrot@.*[$]\s*$',        # parrot@hostname$
        r'backbox@.*[$]\s*$',       # backbox@hostname$
        r'msf6?\s*>',               # Metasploit prompt
        r'\(meterpreter\)\s*>',     # Meterpreter prompt
        r'mysql>\s*$',              # MySQL prompt
        r'ftp>\s*$',                # FTP prompt
        r'>>>\s*$',                 # Python prompt
        r'In \[\d+\]:\s*$',         # IPython prompt
        r'\[sudo\]\s*password',     # Sudo password prompt
    ]

    # Error detection patterns
    ERROR_PATTERNS = [
        r'error:',
        r'failed',
        r'exception',
        r'traceback',
        r'command not found',
        r'permission denied',
        r'connection refused',
        r'no route to host',
        r'timeout',
        r'killed',
        r'segmentation fault'
    ]

    def __init__(self, ssh_manager: SSHManager, tmux_manager: TmuxManager):
        self.ssh = ssh_manager
        self.tmux = tmux_manager
        self.active_monitors = {}
        self.output_buffers = {}  # Track output per session

    async def monitor_session_with_triggers(
        self,
        session_id: str,
        triggers: List[Dict[str, Any]],
        max_timeout: int = 300,
        poll_interval: float = 1.0
    ) -> Dict[str, Any]:
        """
        Enhanced monitoring with comprehensive trigger support.

        Args:
            session_id: Tmux session identifier
            triggers: List of trigger configurations
            max_timeout: Maximum wait time in seconds
            poll_interval: How often to check output (seconds)

        Returns:
            Monitoring result with detailed trigger information
        """
        start_time = time.time()

        # Initialize output buffer for this session
        if session_id not in self.output_buffers:
            self.output_buffers[session_id] = OutputBuffer()

        output_buffer = self.output_buffers[session_id]

        # Add implicit timeout trigger if not present
        has_timeout = any(t.get("type") == "timeout" for t in triggers)
        if not has_timeout:
            triggers.append({"type": "timeout", "timeout_seconds": max_timeout})

        # Validate triggers before starting
        is_valid, error_msg = self.validate_triggers(triggers)
        if not is_valid:
            return {
                "status": "error",
                "error": f"Invalid trigger configuration: {error_msg}",
                "execution_time": 0
            }

        logger.info(f"Starting trigger monitoring for session '{session_id}' with {len(triggers)} triggers")

        while True:
            elapsed_time = time.time() - start_time

            # Check overall timeout
            if elapsed_time >= max_timeout:
                return {
                    "status": "timeout",
                    "message": f"Maximum timeout of {max_timeout} seconds reached",
                    "execution_time": elapsed_time
                }

            try:
                # Capture current output using tmux manager
                capture_result = await self.tmux.capture_pane(session_id)
                if capture_result["status"] != "success":
                    return {
                        "status": "error",
                        "error": f"Failed to capture output: {capture_result.get('message', 'Unknown error')}",
                        "execution_time": elapsed_time
                    }

                current_output = capture_result["output"]
                new_output, full_output = output_buffer.get_new_output(current_output)

                # Check all triggers
                for trigger in triggers:
                    trigger_result = await self._check_trigger(
                        trigger, new_output, full_output, start_time, session_id
                    )

                    if trigger_result["matched"]:
                        logger.info(f"Trigger matched: {trigger_result['trigger_name']} for session '{session_id}'")

                        # Check for errors in output
                        error_info = self._detect_error_in_output(full_output)
                        if error_info["error_detected"]:
                            trigger_result["error_detected"] = error_info

                        return {
                            "status": "trigger_matched",
                            "trigger": trigger_result,
                            "output": full_output,
                            "execution_time": elapsed_time,
                            "session_id": session_id
                        }

                # Sleep before next check
                await asyncio.sleep(poll_interval)

            except Exception as e:
                logger.error(f"Monitoring error for session '{session_id}': {e}")
                return {
                    "status": "error",
                    "error": f"Monitoring failed: {str(e)}",
                    "execution_time": time.time() - start_time
                }

    async def _check_trigger(
        self,
        trigger: Dict[str, Any],
        new_output: str,
        full_output: str,
        start_time: float,
        session_id: str
    ) -> Dict[str, Any]:
        """Check if a trigger condition is met with enhanced logic."""

        trigger_type = trigger.get("type")

        if trigger_type == "prompt":
            return self._check_prompt_trigger(full_output)

        elif trigger_type == "regex":
            return self._check_regex_trigger(trigger, new_output, full_output)

        elif trigger_type == "timeout":
            return self._check_timeout_trigger(start_time, trigger)

        elif trigger_type == "file_exists":
            return await self._check_file_exists_trigger(trigger, session_id)

        elif trigger_type == "multi":
            return await self._check_multi_trigger(trigger, new_output, full_output, start_time, session_id)

        else:
            return {"matched": False, "trigger_name": "unknown"}

    def _check_prompt_trigger(self, full_output: str) -> Dict[str, Any]:
        """Check if output ends with a shell prompt."""
        # Look at last 200 characters only (prompts are at end)
        tail = full_output[-200:] if len(full_output) > 200 else full_output

        for pattern in self.PROMPT_PATTERNS:
            match = re.search(pattern, tail, re.MULTILINE)
            if match:
                return {
                    "matched": True,
                    "trigger_name": "prompt_detected",
                    "trigger_type": "prompt",
                    "matched_pattern": pattern
                }

        return {"matched": False}

    def _check_regex_trigger(self, trigger: Dict[str, Any], new_output: str, full_output: str) -> Dict[str, Any]:
        """Check if regex pattern matches in output."""
        pattern = trigger["pattern"]
        search_in = trigger.get("search_in", "new")  # "new" or "full"

        text = new_output if search_in == "new" else full_output

        match = re.search(pattern, text, re.IGNORECASE | re.MULTILINE)

        if match:
            return {
                "matched": True,
                "trigger_name": trigger.get("name", "regex_match"),
                "trigger_type": "regex",
                "matched_text": match.group(0),
                "match_position": match.start()
            }

        return {"matched": False}

    def _check_timeout_trigger(self, start_time: float, trigger: Dict[str, Any]) -> Dict[str, Any]:
        """Check if timeout condition is met."""
        timeout_seconds = trigger.get("timeout_seconds", 300)
        elapsed = time.time() - start_time

        if elapsed >= timeout_seconds:
            return {
                "matched": True,
                "trigger_name": "timeout",
                "trigger_type": "timeout",
                "elapsed_seconds": elapsed,
                "timeout_seconds": timeout_seconds
            }

        return {"matched": False}

    async def _check_file_exists_trigger(self, trigger: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """Check if file exists on remote system."""
        path = trigger["path"]

        try:
            # Use the session to check file existence
            cmd = f"test -f '{path}' && echo EXISTS || echo NOTFOUND"
            result = await self.tmux.execute_command(session_id, cmd)

            if result.get("status") == "sent":
                # Wait a moment for command to execute
                await asyncio.sleep(0.5)
                output_result = await self.tmux.capture_pane(session_id)

                if "EXISTS" in output_result.get("output", ""):
                    # Get file size for validation
                    stat_cmd = f"stat -c '%s' '{path}' 2>/dev/null || echo 0"
                    await self.tmux.execute_command(session_id, stat_cmd)
                    await asyncio.sleep(0.5)
                    stat_output = await self.tmux.capture_pane(session_id)

                    file_size = 0
                    for line in stat_output.get("output", "").split('\n'):
                        if line.strip().isdigit():
                            file_size = int(line.strip())
                            break

                    return {
                        "matched": True,
                        "trigger_name": trigger.get("name", "file_exists"),
                        "trigger_type": "file_exists",
                        "file_path": path,
                        "file_size": file_size
                    }

        except Exception as e:
            logger.warning(f"File existence check failed for {path}: {e}")

        return {"matched": False}

    async def _check_multi_trigger(
        self,
        trigger: Dict[str, Any],
        new_output: str,
        full_output: str,
        start_time: float,
        session_id: str
    ) -> Dict[str, Any]:
        """Check multiple conditions with AND/OR logic."""
        operator = trigger.get("operator", "AND")
        conditions = trigger["conditions"]

        results = []
        for condition in conditions:
            result = await self._check_trigger(condition, new_output, full_output, start_time, session_id)
            results.append(result["matched"])

        if operator == "AND":
            matched = all(results)
        else:  # OR
            matched = any(results)

        return {
            "matched": matched,
            "trigger_name": trigger.get("name", "multi_condition"),
            "trigger_type": "multi",
            "conditions_met": results,
            "operator": operator
        }

    def _detect_error_in_output(self, output: str) -> Dict[str, Any]:
        """Automatically detect errors in output."""
        for pattern in self.ERROR_PATTERNS:
            match = re.search(pattern, output, re.IGNORECASE)
            if match:
                context = self._extract_error_context(output, pattern)
                return {
                    "error_detected": True,
                    "error_pattern": pattern,
                    "context": context,
                    "suggestion": self._get_error_suggestion(pattern)
                }

        return {"error_detected": False}

    def _extract_error_context(self, output: str, pattern: str) -> str:
        """Extract context around error for better debugging."""
        lines = output.split('\n')

        for i, line in enumerate(lines):
            if re.search(pattern, line, re.IGNORECASE):
                start = max(0, i - 3)
                end = min(len(lines), i + 4)
                return '\n'.join(lines[start:end])

        return ""

    def _get_error_suggestion(self, pattern: str) -> str:
        """Get helpful suggestions for common errors."""
        suggestions = {
            "command not found": "Tool may not be installed. Try: apt install <tool>",
            "permission denied": "Need elevated privileges. Try: sudo <command>",
            "connection refused": "Port may be closed or firewalled",
            "no route to host": "Target unreachable. Check network/VPN",
            "error:": "Check command syntax and parameters",
            "failed": "Operation failed, check dependencies and permissions"
        }

        for error_pattern, suggestion in suggestions.items():
            if error_pattern in pattern.lower():
                return suggestion

        return "Check command and system status"

    def validate_triggers(self, triggers: List[Dict[str, Any]]) -> Tuple[bool, str]:
        """Validate trigger configuration before execution."""
        if not triggers:
            return True, ""

        for trigger in triggers:
            if "type" not in trigger:
                return False, "Trigger missing 'type' field"

            trigger_type = trigger["type"]

            if trigger_type == "regex":
                if "pattern" not in trigger:
                    return False, "Regex trigger missing 'pattern' field"

            elif trigger_type == "timeout":
                if "timeout_seconds" not in trigger:
                    return False, "Timeout trigger missing 'timeout_seconds' field"
                if trigger["timeout_seconds"] <= 0:
                    return False, "Timeout must be positive"

            elif trigger_type == "file_exists":
                if "path" not in trigger:
                    return False, "File exists trigger missing 'path' field"

            elif trigger_type == "multi":
                if "conditions" not in trigger:
                    return False, "Multi trigger missing 'conditions' field"
                if not isinstance(trigger["conditions"], list):
                    return False, "Multi trigger conditions must be a list"

                # Validate sub-conditions
                for condition in trigger["conditions"]:
                    is_valid, error = self.validate_triggers([condition])
                    if not is_valid:
                        return False, f"Invalid condition in multi trigger: {error}"

            elif trigger_type not in ["prompt"]:
                return False, f"Unknown trigger type: {trigger_type}"

        return True, ""


class OutputBuffer:
    """Efficient output buffer management for trigger monitoring."""

    def __init__(self):
        self.last_line_count = 0
        self.full_output = ""

    def get_new_output(self, current_output: str) -> Tuple[str, str]:
        """Get only new lines since last check."""
        current_lines = current_output.split('\n')

        new_lines = current_lines[self.last_line_count:]
        self.last_line_count = len(current_lines)

        new_output = '\n'.join(new_lines)
        self.full_output = current_output

        return new_output, self.full_output

    def reset(self):
        """Reset buffer state."""
        self.last_line_count = 0
        self.full_output = ""
