"""Docker runtime for PentestAgent."""

import asyncio
import io
import tarfile
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Optional

from .runtime import CommandResult, Runtime

if TYPE_CHECKING:
    from ..mcp import MCPManager


@dataclass
class DockerConfig:
    """Docker runtime configuration."""

    image: str = "pentestagent:latest"  # Built from Dockerfile (use pentestagent-kali:latest for Kali image)
    container_name: str = "pentestagent-sandbox"
    network_mode: str = "bridge"
    cap_add: list = None
    volumes: dict = None
    environment: dict = None

    def __post_init__(self):
        if self.cap_add is None:
            self.cap_add = ["NET_ADMIN"]  # For VPN
        if self.volumes is None:
            self.volumes = {}
        if self.environment is None:
            self.environment = {}


class DockerRuntime(Runtime):
    """Manages Docker sandbox for tool execution."""

    def __init__(
        self,
        config: Optional[DockerConfig] = None,
        vpn_config: Optional[Path] = None,
        mcp_manager: Optional["MCPManager"] = None,
    ):
        """
        Initialize the Docker runtime.

        Args:
            config: Docker configuration
            vpn_config: Path to OpenVPN config file
            mcp_manager: MCP manager for tool calls
        """
        super().__init__(mcp_manager)
        self.config = config or DockerConfig()
        self.vpn_config = vpn_config
        self.client = None
        self.container = None
        self._browser_context = None
        self._proxy_running = False
        self._proxy_port = 8080

    async def start(self):
        """Start the sandbox container."""
        try:
            import docker

            self.client = docker.from_env()
        except ImportError as e:
            raise ImportError(
                "docker is required for Docker runtime. "
                "Install with: pip install docker"
            ) from e

        # Check if container already exists
        try:
            self.container = self.client.containers.get(self.config.container_name)
            if self.container.status != "running":
                self.container.start()
                await asyncio.sleep(2)  # Wait for container to fully start
        except Exception:
            # Create new container
            volumes = {
                str(Path.home() / ".pentestagent"): {
                    "bind": "/root/.pentestagent",
                    "mode": "rw",
                },
                **self.config.volumes,
            }

            try:
                self.container = self.client.containers.run(
                    self.config.image,
                    name=self.config.container_name,
                    detach=True,
                    tty=True,
                    cap_add=self.config.cap_add,
                    volumes=volumes,
                    network_mode=self.config.network_mode,
                    environment=self.config.environment,
                )
                await asyncio.sleep(2)
            except Exception as pull_err:
                import docker as _docker
                if isinstance(pull_err, _docker.errors.ImageNotFound):
                    raise RuntimeError(
                        f"Docker image '{self.config.image}' not found. "
                        "Build it first with: docker compose build"
                    ) from pull_err
                raise

        # Setup VPN if configured
        if self.vpn_config:
            await self._setup_vpn()

    async def stop(self):
        """Stop and remove the sandbox container."""
        if self.container:
            try:
                self.container.stop(timeout=10)
                self.container.remove()
            except Exception as e:
                logging.getLogger(__name__).exception("Failed stopping/removing container: %s", e)
                try:
                    from ..interface.notifier import notify

                    notify("warning", f"DockerRuntime: failed stopping/removing container: {e}")
                except Exception as ne:
                    logging.getLogger(__name__).exception("Failed to notify operator about docker stop error: %s", ne)
            finally:
                self.container = None

    async def execute_command(self, command: str, timeout: int = 300) -> CommandResult:
        """
        Execute a command in the sandbox.

        Args:
            command: The shell command to execute
            timeout: Timeout in seconds

        Returns:
            CommandResult with output
        """
        if not self.container:
            raise RuntimeError("Sandbox not started")

        try:
            # Execute command
            exec_result = self.container.exec_run(
                cmd=["bash", "-c", command], demux=True
            )

            stdout = (
                exec_result.output[0].decode(errors="replace")
                if exec_result.output[0]
                else ""
            )
            stderr = (
                exec_result.output[1].decode(errors="replace")
                if exec_result.output[1]
                else ""
            )

            return CommandResult(
                exit_code=exec_result.exit_code, stdout=stdout, stderr=stderr
            )

        except Exception as e:
            return CommandResult(
                exit_code=-1, stdout="", stderr=f"Execution error: {str(e)}"
            )

    async def browser_action(self, action: str, **kwargs) -> dict:
        """
        Perform browser automation in the sandbox.

        Args:
            action: The browser action to perform
            **kwargs: Action-specific arguments

        Returns:
            Action result dictionary
        """
        # This would communicate with a browser automation service in the container
        # For now, we'll use a simple implementation via terminal commands

        if action == "navigate":
            url = kwargs.get("url", "")
            result = await self.execute_command(f"curl -s -L -o /tmp/page.html '{url}'")
            if result.success:
                content_result = await self.execute_command(
                    "head -c 5000 /tmp/page.html"
                )
                return {
                    "url": url,
                    "title": "Retrieved",
                    "content": content_result.stdout,
                }
            return {"error": result.stderr}

        elif action == "get_content":
            result = await self.execute_command(
                "cat /tmp/page.html 2>/dev/null || echo 'No page loaded'"
            )
            return {"content": result.stdout}

        elif action == "get_links":
            result = await self.execute_command(
                "grep -oP 'href=\"\\K[^\"]+' /tmp/page.html 2>/dev/null | head -50"
            )
            links = [
                {"href": link, "text": ""}
                for link in result.stdout.strip().split("\n")
                if link
            ]
            return {"links": links}

        elif action == "screenshot":
            return {
                "error": "Screenshot requires graphical browser - not available in sandbox"
            }

        return {"error": f"Unknown browser action: {action}"}

    async def proxy_action(self, action: str, **kwargs) -> dict:
        """
        Control the HTTP proxy in the sandbox.

        Args:
            action: The proxy action to perform
            **kwargs: Action-specific arguments

        Returns:
            Action result dictionary
        """
        port = kwargs.get("port", self._proxy_port)

        if action == "start":
            # Start mitmproxy in the background
            result = await self.execute_command(
                f"mitmdump -p {port} --set block_global=false -w /tmp/proxy.flow &"
            )
            if result.exit_code == 0:
                self._proxy_running = True
                self._proxy_port = port
                return {"status": "started", "port": port}
            return {"error": result.stderr}

        elif action == "stop":
            await self.execute_command("pkill -f mitmdump")
            self._proxy_running = False
            return {"status": "stopped"}

        elif action == "status":
            result = await self.execute_command("pgrep -f mitmdump")
            return {
                "status": "running" if result.exit_code == 0 else "stopped",
                "port": self._proxy_port,
                "request_count": 0,  # Would need to parse proxy logs
            }

        elif action == "get_history":
            # Would parse /tmp/proxy.flow
            return {"requests": []}

        elif action == "clear_history":
            await self.execute_command("rm -f /tmp/proxy.flow")
            return {"status": "cleared"}

        return {"error": f"Unknown proxy action: {action}"}

    async def is_running(self) -> bool:
        """Check if the container is running."""
        if not self.container:
            return False

        try:
            self.container.reload()
            return self.container.status == "running"
        except Exception as e:
            logging.getLogger(__name__).exception("Failed to determine container running state: %s", e)
            try:
                from ..interface.notifier import notify

                notify("warning", f"DockerRuntime: is_running check failed: {e}")
            except Exception as notify_error:
                logging.getLogger(__name__).warning(
                    "Failed to send notification for DockerRuntime.is_running error: %s",
                    notify_error,
                )
            return False

    async def get_status(self) -> dict:
        """Get runtime status information."""
        running = await self.is_running()

        status = {
            "type": "docker",
            "running": running,
            "container_name": self.config.container_name,
            "image": self.config.image,
            "proxy_running": self._proxy_running,
        }

        if running:
            # Get container info
            self.container.reload()
            status["container_id"] = self.container.short_id
            status["network_mode"] = self.config.network_mode

        return status

    async def _setup_vpn(self):
        """Configure VPN in the sandbox."""
        if not self.vpn_config or not self.vpn_config.exists():
            return

        # Copy VPN config to container
        config_content = self.vpn_config.read_bytes()
        tar_data = self._create_tar(config_content, "client.ovpn")

        self.container.put_archive("/etc/openvpn", tar_data)

        # Start OpenVPN
        await self.execute_command(
            "openvpn --config /etc/openvpn/client.ovpn --daemon --log /var/log/openvpn.log"
        )

        # Wait for connection
        await asyncio.sleep(5)

        # Verify connection
        result = await self.execute_command("curl -s --max-time 10 ifconfig.me")
        if result.success:
            print(f"[VPN] Connected. External IP: {result.stdout.strip()}")
        else:
            print("[VPN] Connection may have failed")

    def _create_tar(self, content: bytes, filename: str) -> bytes:
        """Create a tar archive for container upload."""
        tar_stream = io.BytesIO()
        tar = tarfile.open(fileobj=tar_stream, mode="w")

        file_data = io.BytesIO(content)
        info = tarfile.TarInfo(name=filename)
        info.size = len(content)
        tar.addfile(info, file_data)
        tar.close()

        tar_stream.seek(0)
        return tar_stream.read()

    async def copy_to_container(self, local_path: Path, container_path: str):
        """
        Copy a file to the container.

        Args:
            local_path: Local file path
            container_path: Destination path in container
        """
        if not self.container:
            raise RuntimeError("Container not running")

        content = local_path.read_bytes()
        filename = local_path.name
        tar_data = self._create_tar(content, filename)

        # Ensure directory exists
        dir_path = str(Path(container_path).parent)
        await self.execute_command(f"mkdir -p {dir_path}")

        self.container.put_archive(dir_path, tar_data)

    async def copy_from_container(self, container_path: str, local_path: Path):
        """
        Copy a file from the container.

        Args:
            container_path: Source path in container
            local_path: Local destination path
        """
        if not self.container:
            raise RuntimeError("Container not running")

        bits, stat = self.container.get_archive(container_path)

        # Extract from tar
        tar_stream = io.BytesIO()
        for chunk in bits:
            tar_stream.write(chunk)
        tar_stream.seek(0)

        tar = tarfile.open(fileobj=tar_stream)
        for member in tar.getmembers():
            f = tar.extractfile(member)
            if f:
                local_path.write_bytes(f.read())
                break
        tar.close()
