"""
Main MCP Server for Pentest MCP Server.
Implements MCP tools for SSH + Tmux session management.
"""

import asyncio
import logging
import json
import re
import time
from typing import Any, Dict, List, Optional, Sequence
from mcp.server import Server
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server
from mcp.types import (
    Resource,
    Tool,
    TextContent,
    ImageContent,
    EmbeddedResource,
)

from .config import Config
from .ssh_manager import SSHManager, SSHConnectionError
from .tmux_manager import TmuxManager, TmuxSessionError

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class PentestMCPServer:
    """Main MCP Server for autonomous pentesting operations."""

    def __init__(self):
        self.config = Config()
        self.ssh_manager: Optional[SSHManager] = None
        self.tmux_manager: Optional[TmuxManager] = None
        self.server = Server("pentest-mcp")
        self._setup_resources()
        self._setup_tools()

    def _setup_resources(self):
        """Setup MCP resources."""
        @self.server.list_resources()
        async def handle_list_resources() -> List[Resource]:
            return [
                Resource(
                    uri="pentest://sessions",
                    name="Active Sessions",
                    description="List of all active tmux sessions",
                    mimeType="application/json"
                ),
                Resource(
                    uri="pentest://system-status", 
                    name="System Status",
                    description="Current system resource usage and connection status",
                    mimeType="application/json"
                )
            ]

        @self.server.read_resource()
        async def handle_read_resource(uri: str) -> List[TextContent]:
            if uri == "pentest://sessions":
                sessions = await self.tmux_manager.list_sessions()
                return [TextContent(type="text", text=json.dumps(sessions))]
            elif uri == "pentest://system-status":
                status = await self._handle_get_system_status({})
                return [TextContent(type="text", text=json.dumps(status))]
            else:
                raise ValueError(f"Unknown resource: {uri}")

    def _setup_tools(self):
        """Setup MCP tools."""

        @self.server.list_tools()
        async def handle_list_tools() -> List[Tool]:
            return [
                Tool(
                    name="create_session",
                    description="Create a new persistent tmux session",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "Unique identifier for the session"
                            },
                            "shell": {
                                "type": "string", 
                                "description": "Shell to use (default: bash)",
                                "default": "bash"
                            }
                        },
                        "required": ["session_id"]
                    }
                ),
                Tool(
                    name="list_sessions",
                    description="List all active tmux sessions",
                    inputSchema={
                        "type": "object",
                        "properties": {}
                    }
                ),
                Tool(
                    name="kill_session", 
                    description="Terminate a tmux session",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "Session to terminate"
                            }
                        },
                        "required": ["session_id"]
                    }
                ),
                Tool(
                    name="execute",
                    description="Execute a command in a persistent tmux session with trigger-based monitoring",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "Target session ID"
                            },
                            "command": {
                                "type": "string",
                                "description": "Command to execute"
                            },
                            "background": {
                                "type": "boolean",
                                "description": "Run in background (default: false)",
                                "default": False
                            },
                            "triggers": {
                                "type": "array",
                                "description": "Triggers for AI suspension/resumption",
                                "items": {
                                    "type": "object",
                                    "properties": {
                                        "type": {
                                            "type": "string",
                                            "description": "prompt | regex | timeout | file_exists",
                                            "enum": ["prompt", "regex", "timeout", "file_exists"]
                                        },
                                        "pattern": {
                                            "type": "string",
                                            "description": "Regex pattern (for regex type)"
                                        },
                                        "timeout_seconds": {
                                            "type": "integer",
                                            "description": "Timeout in seconds (for timeout type)"
                                        },
                                        "path": {
                                            "type": "string",
                                            "description": "File path (for file_exists type)"
                                        },
                                        "name": {
                                            "type": "string",
                                            "description": "Trigger identifier"
                                        }
                                    },
                                    "required": ["type"]
                                }
                            },
                            "max_timeout": {
                                "type": "integer",
                                "description": "Maximum wait time in seconds (default: 300)",
                                "default": 300
                            }
                        },
                        "required": ["session_id", "command"]
                    }
                ),
                Tool(
                    name="read_output",
                    description="Get current output from a session",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "Target session ID"
                            },
                            "lines": {
                                "type": "integer", 
                                "description": "Number of lines to return (default: all)"
                            }
                        },
                        "required": ["session_id"]
                    }
                ),
                Tool(
                    name="send_input",
                    description="Send input to an interactive session",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "Target session ID"
                            },
                            "input": {
                                "type": "string",
                                "description": "Text to send"
                            },
                            "press_enter": {
                                "type": "boolean",
                                "description": "Press Enter after input (default: true)",
                                "default": True
                            }
                        },
                        "required": ["session_id", "input"]
                    }
                ),
                Tool(
                    name="get_system_status",
                    description="Get your OS system resource usage and status",
                    inputSchema={
                        "type": "object",
                        "properties": {}
                    }
                ),
                Tool(
                    name="recover_sessions",
                    description="Find and reconnect to orphaned tmux sessions",
                    inputSchema={
                        "type": "object",
                        "properties": {}
                    }
                ),
                Tool(
                    name="upload_file",
                    description="Upload a file to the your OS system",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "local_path": {
                                "type": "string",
                                "description": "Local file path to upload"
                            },
                            "remote_path": {
                                "type": "string", 
                                "description": "Remote destination path"
                            }
                        },
                        "required": ["local_path", "remote_path"]
                    }
                ),
                Tool(
                    name="download_file", 
                    description="Download a file from the your OS system",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "remote_path": {
                                "type": "string",
                                "description": "Remote file path to download"
                            },
                            "local_path": {
                                "type": "string",
                                "description": "Local destination path"
                            }
                        },
                        "required": ["remote_path", "local_path"]
                    }
                ),
                Tool(
                    name="get_session_history",
                    description="Get command history for a specific tmux session",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "Session to get command history for"
                            }
                        },
                        "required": ["session_id"]
                    }
                ),
                Tool(
                    name="parse_tool_output",
                    description="Parse common pentest tool outputs (nmap XML, JSON, etc.)",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "tool": {
                                "type": "string",
                                "description": "nmap | masscan | sqlmap | custom",
                                "enum": ["nmap", "masscan", "sqlmap", "custom"]
                            },
                            "file_path": {
                                "type": "string",
                                "description": "Path to output file on your OS system"
                            },
                            "format": {
                                "type": "string",
                                "description": "xml | json | txt",
                                "enum": ["xml", "json", "txt"]
                            }
                        },
                        "required": ["tool", "file_path"]
                    }
                )
            ]

        @self.server.call_tool()
        async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
            """Handle tool calls."""

            try:
                # Ensure connection is established
                if not await self._ensure_initialized():
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "status": "error",
                            "message": "Failed to initialize SSH/Tmux connection"
                        })
                    )]

                # Route to appropriate handler
                if name == "create_session":
                    result = await self._handle_create_session(arguments)
                elif name == "list_sessions":
                    result = await self._handle_list_sessions(arguments)
                elif name == "kill_session":
                    result = await self._handle_kill_session(arguments)
                elif name == "execute":
                    result = await self._handle_execute(arguments)
                elif name == "read_output":
                    result = await self._handle_read_output(arguments)
                elif name == "send_input":
                    result = await self._handle_send_input(arguments)
                elif name == "get_system_status":
                    result = await self._handle_get_system_status(arguments)
                elif name == "recover_sessions":
                    result = await self._handle_recover_sessions(arguments)
                elif name == "upload_file":
                    result = await self._handle_upload_file(arguments)
                elif name == "download_file":
                    result = await self._handle_download_file(arguments)
                elif name == "get_session_history":
                    result = await self._handle_get_session_history(arguments)
                elif name == "parse_tool_output":
                    result = await self._handle_parse_tool_output(arguments)
                elif name == "list_files":
                    result = await self._handle_list_files(arguments)
                else:
                    result = {
                        "status": "error",
                        "message": f"Unknown tool: {name}"
                    }

                return [TextContent(type="text", text=json.dumps(result))]

            except Exception as e:
                logger.error(f"Error handling tool {name}: {e}")
                return [TextContent(
                    type="text",
                    text=json.dumps({
                        "status": "error",
                        "message": f"Tool execution failed: {str(e)}"
                    })
                )]

    async def _ensure_initialized(self) -> bool:
        """Ensure SSH and Tmux managers are initialized."""
        max_attempts = 3
        for attempt in range(max_attempts):
            try:
                if not self.ssh_manager:
                    self.ssh_manager = SSHManager(self.config)
                    connected = await self.ssh_manager.connect()
                    if not connected:
                        continue

                if not self.tmux_manager:
                    self.tmux_manager = TmuxManager(self.ssh_manager)
                    initialized = await self.tmux_manager.initialize()
                    if not initialized:
                        continue

                # Ensure connection is alive
                if await self.ssh_manager.ensure_connected():
                    return True

            except Exception as e:
                logger.error(f"Initialization attempt {attempt + 1} failed: {e}")
                if attempt == max_attempts - 1:
                    return False
                await asyncio.sleep(2 ** attempt)  # Exponential backoff

        return False

    # Tool Handlers
    async def _handle_create_session(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle create_session tool call."""
        session_id = args.get("session_id")
        shell = args.get("shell", "bash")

        if not session_id:
            return {"status": "error", "message": "session_id is required"}

        return await self.tmux_manager.create_session(session_id, shell)

    async def _handle_list_sessions(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle list_sessions tool call."""
        sessions = await self.tmux_manager.list_sessions()
        return {"status": "success", "sessions": sessions}

    async def _handle_kill_session(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle kill_session tool call."""
        session_id = args.get("session_id")

        if not session_id:
            return {"status": "error", "message": "session_id is required"}

        return await self.tmux_manager.kill_session(session_id)

    async def _handle_execute(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Enhanced execute tool call with comprehensive trigger support."""
        session_id = args.get("session_id")
        command = args.get("command")
        background = args.get("background", False)
        triggers = args.get("triggers", [{"type": "prompt"}])  # Default to prompt trigger
        max_timeout = args.get("max_timeout", 300)
        poll_interval = args.get("poll_interval", 1.0)

        if not session_id or not command:
            return {"status": "error", "message": "session_id and command are required"}

        # Validate command (optional security check)
        if not await self._validate_command(command):
            return {"status": "error", "message": f"Command blocked for security: {command}"}

        # Ensure we're initialized
        if not await self._ensure_initialized():
            return {"status": "error", "message": "SSH/Tmux connection not available"}

        # Execute command
        execute_result = await self.tmux_manager.execute_command(session_id, command)
        if execute_result.get("status") != "sent":
            return {
                "status": "error",
                "message": f"Failed to execute command: {execute_result.get('message', 'Unknown error')}"
            }

        # Background execution: return immediately
        if background:
            return {
                "status": "background",
                "message": "Command sent, running in background",
                "session_id": session_id,
                "command": command
            }

        # Foreground execution: monitor with enhanced trigger system
        try:
            from .trigger_system import TriggerSystem
            trigger_system = TriggerSystem(self.ssh_manager, self.tmux_manager)

            monitor_result = await trigger_system.monitor_session_with_triggers(
                session_id=session_id,
                triggers=triggers,
                max_timeout=max_timeout,
                poll_interval=poll_interval
            )

            # Add command context to result
            monitor_result["command"] = command
            monitor_result["session_id"] = session_id

            return monitor_result

        except Exception as e:
            logger.error(f"Error during trigger monitoring: {e}")
            # Fallback: wait briefly and return output
            await asyncio.sleep(2)
            output_result = await self.tmux_manager.capture_pane(session_id)

            return {
                "status": "completed",
                "command": command,
                "output": output_result.get("output", ""),
                "execution_time": 2.0,
                "note": "Trigger monitoring failed, using fallback"
            }

    async def _handle_read_output(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle read_output tool call."""
        session_id = args.get("session_id")
        lines = args.get("lines")

        if not session_id:
            return {"status": "error", "message": "session_id is required"}

        result = await self.tmux_manager.capture_pane(session_id, lines)
        return result

    async def _handle_send_input(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle send_input tool call."""
        session_id = args.get("session_id")
        input_text = args.get("input")
        press_enter = args.get("press_enter", True)

        if not session_id or input_text is None:
            return {"status": "error", "message": "session_id and input are required"}

        return await self.tmux_manager.send_input(session_id, input_text, press_enter)

    async def _handle_get_system_status(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle get_system_status tool call."""
        try:
            # Get system info
            system_info = await self.ssh_manager.get_system_info()

            # Get session count
            sessions = await self.tmux_manager.list_sessions()

            # Get connection status
            connection_status = self.ssh_manager.get_connection_status()

            return {
                "status": "success",
                "system_info": system_info,
                "active_sessions": len(sessions),
                "connection_status": connection_status,
                "sessions": sessions
            }

        except Exception as e:
            return {"status": "error", "message": f"Failed to get system status: {e}"}

    async def _handle_recover_sessions(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle recover_sessions tool call."""
        try:
            recovered = await self.tmux_manager.recover_sessions()

            recovered_info = []
            for session_id in recovered:
                session_info = self.tmux_manager.get_session_info(session_id)
                if session_info:
                    # Get last output
                    output_result = await self.tmux_manager.capture_pane(session_id, 20)
                    last_output = output_result.get("output", "")[-500:]  # Last 500 chars

                    recovered_info.append({
                        "session_id": session_id,
                        "age_seconds": int(time.time() - session_info.created_at),
                        "last_output": last_output
                    })

            return {
                "status": "success",
                "recovered_sessions": recovered_info
            }

        except Exception as e:
            return {"status": "error", "message": f"Recovery failed: {e}"}

    async def _handle_get_session_history(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle get_session_history tool call."""
        session_id = args.get("session_id")

        if not session_id:
            return {"status": "error", "message": "session_id is required"}

        try:
            result = await self.tmux_manager.get_session_history(session_id)
            return result
        except Exception as e:
            return {"status": "error", "message": f"Failed to get session history: {e}"}

    async def _handle_upload_file(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle upload_file tool call."""
        try:
            import os

            local_path = args.get("local_path")
            remote_path = args.get("remote_path")

            if not await self._ensure_initialized():
                return {"status": "error", "message": "SSH connection not available"}

            # Validate local path exists on the MCP host (likely Windows)
            if not local_path or not os.path.isfile(local_path):
                return {"status": "error", "message": f"Local file not found: {local_path}"}

            # Ensure remote directory exists (use SSH command via tmux/exec)
            remote_dir = os.path.dirname(remote_path) or "/"
            try:
                await self.ssh_manager.run_command(f"mkdir -p '{remote_dir}'", timeout=10)
            except Exception:
                # Non-fatal: continue and let SFTP report precise errors
                pass

            # Use async context manager for SFTP lifecycle
            sftp = await self.ssh_manager.start_sftp_client()
            try:
                await sftp.put(local_path, remote_path)
            finally:
                try:
                    await sftp.exit()
                except Exception:
                    # Some asyncssh versions use .close(); attempt fallback
                    try:
                        await sftp.close()  # type: ignore[attr-defined]
                    except Exception:
                        pass

            return {
                "status": "success",
                "message": f"File uploaded to {remote_path}"
            }

        except Exception as e:
            return {"status": "error", "message": f"Upload failed: {e}"}

    async def _handle_download_file(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle download_file tool call."""
        try:
            import os

            remote_path = args.get("remote_path")
            local_path = args.get("local_path")

            if not await self._ensure_initialized():
                return {"status": "error", "message": "SSH connection not available"}

            # Ensure local destination directory exists on MCP host
            local_dir = os.path.dirname(local_path) or "."
            try:
                os.makedirs(local_dir, exist_ok=True)
            except Exception as mkerr:
                return {"status": "error", "message": f"Failed to create local directory '{local_dir}': {mkerr}"}

            sftp = await self.ssh_manager.start_sftp_client()
            try:
                await sftp.get(remote_path, local_path)
            finally:
                try:
                    await sftp.exit()
                except Exception:
                    try:
                        await sftp.close()  # type: ignore[attr-defined]
                    except Exception:
                        pass

            return {
                "status": "success",
                "message": f"File downloaded to {local_path}"
            }

        except Exception as e:
            return {"status": "error", "message": f"Download failed: {e}"}

    async def _handle_parse_tool_output(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle parse_tool_output tool call."""
        try:
            tool = args.get("tool")
            file_path = args.get("file_path")
            format_type = args.get("format", "xml")

            if not await self._ensure_initialized():
                return {"status": "error", "message": "SSH connection not available"}

            # Download the file first
            import tempfile
            import os

            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_path = temp_file.name

            try:
                # Download the file
                sftp = await self.ssh_manager.start_sftp_client()
                try:
                    await sftp.get(file_path, temp_path)
                finally:
                    try:
                        await sftp.exit()
                    except Exception:
                        try:
                            await sftp.close()  # type: ignore[attr-defined]
                        except Exception:
                            pass

                # Parse based on tool and format
                parsed_data = await self._parse_output_file(tool, format_type, temp_path)

                return {
                    "status": "success",
                    "parsed_data": parsed_data,
                    "summary": f"Successfully parsed {tool} output"
                }

            finally:
                # Cleanup temp file
                if os.path.exists(temp_path):
                    os.unlink(temp_path)

        except Exception as e:
            return {"status": "error", "message": f"Output parsing failed: {e}"}

    async def _handle_list_files(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Handle list_files tool call."""
        try:
            path = args.get("path")
            recursive = args.get("recursive", False)

            if not await self._ensure_initialized():
                return {"status": "error", "message": "SSH connection not available"}

            if recursive:
                cmd = f"find {path} -type f -printf '%p %s %T@\\n' 2>/dev/null || echo 'NOT_FOUND'"
            else:
                cmd = f"ls -la {path} 2>/dev/null || echo 'NOT_FOUND'"

            result = await self.ssh_manager.run_command(cmd)

            if "NOT_FOUND" in result.stdout:
                return {"status": "error", "message": f"Path not found: {path}"}

            return {
                "status": "success",
                "output": result.stdout,
                "path": path,
                "recursive": recursive
            }

        except Exception as e:
            return {"status": "error", "message": f"File listing failed: {e}"}

    async def _parse_output_file(self, tool: str, format_type: str, file_path: str) -> Dict[str, Any]:
        """Parse tool output files."""
        try:
            if tool == "nmap" and format_type == "xml":
                return await self._parse_nmap_xml(file_path)
            elif tool == "masscan" and format_type == "xml":
                return await self._parse_masscan_xml(file_path)
            else:
                # For unsupported combinations, return basic file info
                import os
                return {
                    "file_size": os.path.getsize(file_path),
                    "tool": tool,
                    "format": format_type,
                    "note": "Automatic parsing not implemented for this tool/format"
                }
        except Exception as e:
            return {"error": f"Parsing failed: {e}"}

    async def _parse_nmap_xml(self, file_path: str) -> Dict[str, Any]:
        """Parse nmap XML output."""
        try:
            import xml.etree.ElementTree as ET

            tree = ET.parse(file_path)
            root = tree.getroot()

            hosts = []
            for host in root.findall('host'):
                host_info = {}

                # Get address
                address_elem = host.find('address')
                if address_elem is not None:
                    host_info['address'] = address_elem.get('addr')
                    host_info['type'] = address_elem.get('addrtype')

                # Get ports
                ports_info = []
                ports = host.find('ports')
                if ports is not None:
                    for port in ports.findall('port'):
                        port_info = {
                            'port': port.get('portid'),
                            'protocol': port.get('protocol'),
                            'state': port.find('state').get('state') if port.find('state') is not None else 'unknown',
                            'service': port.find('service').get('name') if port.find('service') is not None else 'unknown'
                        }
                        ports_info.append(port_info)

                host_info['ports'] = ports_info
                hosts.append(host_info)

            return {
                "hosts": hosts,
                "scan_info": {
                    "scanner": "nmap",
                    "hosts_found": len(hosts)
                }
            }

        except Exception as e:
            return {"error": f"Nmap XML parsing failed: {e}"}

    async def _parse_masscan_xml(self, file_path: str) -> Dict[str, Any]:
        """Parse masscan XML output."""
        try:
            import xml.etree.ElementTree as ET

            tree = ET.parse(file_path)
            root = tree.getroot()

            hosts = []
            for host in root.findall('host'):
                host_info = {}

                # Get address
                address_elem = host.find('address')
                if address_elem is not None:
                    host_info['address'] = address_elem.get('addr')
                    host_info['type'] = address_elem.get('addrtype')

                # Get ports
                ports_info = []
                for port in host.findall('ports/port'):
                    port_info = {
                        'port': port.get('portid'),
                        'protocol': port.get('protocol'),
                        'state': 'open'  # masscan typically only shows open ports
                    }
                    ports_info.append(port_info)

                host_info['ports'] = ports_info
                hosts.append(host_info)

            return {
                "hosts": hosts,
                "scan_info": {
                    "scanner": "masscan",
                    "hosts_found": len(hosts)
                }
            }

        except Exception as e:
            return {"error": f"Masscan XML parsing failed: {e}"}

    async def _validate_command(self, command: str) -> bool:
        """
        Validate command for security (optional).

        Args:
            command: Command to validate

        Returns:
            bool: True if command is allowed
        """
        # Check against blocked patterns
        for pattern in self.config.BLOCKED_COMMANDS:
            if re.search(pattern, command, re.IGNORECASE):
                logger.warning(f"Blocked dangerous command: {command}")
                return False

        return True

    async def shutdown(self):
        """Cleanup resources on shutdown."""
        try:
            if self.ssh_manager:
                await self.ssh_manager.disconnect()
            logger.info("MCP Server shutdown completed")
        except Exception as e:
            logger.error(f"Error during shutdown: {e}")


async def main():
    """Main entry point for the MCP server."""
    try:
        # Validate configuration
        Config.validate_config()

        # Create and run server
        server_instance = PentestMCPServer()

        # Setup signal handlers for clean shutdown
        import signal

        def signal_handler(signum, frame):
            logger.info(f"Received signal {signum}, shutting down...")
            asyncio.create_task(server_instance.shutdown())

        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)

        # Run the server
        async with stdio_server() as (read_stream, write_stream):
            await server_instance.server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name="pentest-mcp",
                    server_version="1.0.0",
                    capabilities={
                        "tools": {},
                        "resources": {}
                    }
                )
            )

    except Exception as e:
        logger.error(f"Server failed to start: {e}")
        raise


if __name__ == "__main__":
    asyncio.run(main())
