"""
Container and Process Isolation for Agent Workspaces.

Provides multiple isolation strategies:
- TMUX: Process isolation via tmux sessions (default)
- DOCKER: Container isolation via Docker
- PODMAN: Container isolation via Podman

Features:
- Auto-credential injection
- Configurable resource limits
- Network isolation options
- Workspace volume mounting
"""

import json
import logging
import os
import shutil
import subprocess
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)


class IsolationType(Enum):
    """Types of process isolation."""

    TMUX = "tmux"  # tmux session (default, lightweight)
    DOCKER = "docker"  # Docker container
    PODMAN = "podman"  # Podman container (rootless)


class NetworkMode(Enum):
    """Container network modes."""

    HOST = "host"  # Share host network
    BRIDGE = "bridge"  # Default bridge network
    NONE = "none"  # No network access


@dataclass
class ResourceLimits:
    """Resource limits for containerized agents."""

    memory_mb: int = 4096  # Memory limit in MB
    cpu_cores: float = 2.0  # CPU cores limit
    disk_gb: int = 10  # Disk space limit in GB
    pids_limit: int = 100  # Max number of processes

    def to_docker_args(self) -> list[str]:
        """Convert to Docker run arguments."""
        args = [
            f"--memory={self.memory_mb}m",
            f"--cpus={self.cpu_cores}",
            f"--pids-limit={self.pids_limit}",
        ]
        return args


@dataclass
class CredentialConfig:
    """Configuration for credential injection."""

    inject_git_config: bool = True  # Inject git user config
    inject_ssh_keys: bool = False  # Mount SSH keys
    inject_aws_credentials: bool = False  # Mount AWS credentials
    inject_env_vars: list[str] = field(default_factory=list)  # Env vars to pass
    custom_env: dict[str, str] = field(default_factory=dict)  # Custom env vars

    def get_env_dict(self) -> dict[str, str]:
        """Get environment variables to inject."""
        env = dict(self.custom_env)

        # Include specified env vars from current environment
        for var in self.inject_env_vars:
            if var in os.environ:
                env[var] = os.environ[var]

        return env


@dataclass
class IsolationConfig:
    """Configuration for agent isolation."""

    isolation_type: IsolationType = IsolationType.TMUX
    container_image: str = "ubuntu:22.04"  # Base image for containers
    network_mode: NetworkMode = NetworkMode.BRIDGE
    resource_limits: ResourceLimits = field(default_factory=ResourceLimits)
    credentials: CredentialConfig = field(default_factory=CredentialConfig)

    # Additional options
    privileged: bool = False  # Run container privileged (not recommended)
    read_only_root: bool = False  # Read-only root filesystem
    mount_docker_socket: bool = False  # Mount Docker socket (for DinD)
    extra_mounts: list[tuple[str, str]] = field(default_factory=list)  # (host, container)
    extra_args: list[str] = field(default_factory=list)  # Additional container args

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "isolation_type": self.isolation_type.value,
            "container_image": self.container_image,
            "network_mode": self.network_mode.value,
            "resource_limits": {
                "memory_mb": self.resource_limits.memory_mb,
                "cpu_cores": self.resource_limits.cpu_cores,
                "disk_gb": self.resource_limits.disk_gb,
                "pids_limit": self.resource_limits.pids_limit,
            },
            "credentials": {
                "inject_git_config": self.credentials.inject_git_config,
                "inject_ssh_keys": self.credentials.inject_ssh_keys,
                "inject_aws_credentials": self.credentials.inject_aws_credentials,
                "inject_env_vars": self.credentials.inject_env_vars,
            },
        }


@dataclass
class IsolatedEnvironment:
    """Represents an isolated environment for an agent."""

    agent_id: str
    isolation_type: IsolationType
    workspace_path: Path
    container_id: Optional[str] = None
    tmux_session: Optional[str] = None
    created_at: datetime = field(default_factory=datetime.now)
    config: Optional[IsolationConfig] = None

    @property
    def is_container(self) -> bool:
        """Check if this is a container-based isolation."""
        return self.isolation_type in (IsolationType.DOCKER, IsolationType.PODMAN)

    @property
    def runtime_cmd(self) -> str:
        """Get the container runtime command."""
        if self.isolation_type == IsolationType.DOCKER:
            return "docker"
        elif self.isolation_type == IsolationType.PODMAN:
            return "podman"
        return ""


class IsolationManager:
    """
    Manages isolated environments for agents.

    Supports tmux sessions, Docker containers, and Podman containers.
    """

    def __init__(
        self,
        default_config: Optional[IsolationConfig] = None,
    ) -> None:
        """
        Initialize the isolation manager.

        Args:
            default_config: Default isolation configuration
        """
        self.default_config = default_config or IsolationConfig()
        self.environments: dict[str, IsolatedEnvironment] = {}

        # Check available runtimes
        self._tmux_available = shutil.which("tmux") is not None
        self._docker_available = shutil.which("docker") is not None
        self._podman_available = shutil.which("podman") is not None

    def get_available_runtimes(self) -> list[IsolationType]:
        """Get list of available isolation runtimes."""
        available = []
        if self._tmux_available:
            available.append(IsolationType.TMUX)
        if self._docker_available:
            available.append(IsolationType.DOCKER)
        if self._podman_available:
            available.append(IsolationType.PODMAN)
        return available

    def create_environment(
        self,
        agent_id: str,
        workspace_path: Path,
        config: Optional[IsolationConfig] = None,
    ) -> IsolatedEnvironment:
        """
        Create an isolated environment for an agent.

        Args:
            agent_id: Unique identifier for the agent
            workspace_path: Path to the agent's workspace
            config: Isolation configuration (uses default if not provided)

        Returns:
            IsolatedEnvironment object

        Raises:
            RuntimeError: If the required runtime is not available
        """
        config = config or self.default_config

        # Validate runtime availability
        if config.isolation_type == IsolationType.TMUX and not self._tmux_available:
            raise RuntimeError("tmux is not available")
        if config.isolation_type == IsolationType.DOCKER and not self._docker_available:
            raise RuntimeError("Docker is not available")
        if config.isolation_type == IsolationType.PODMAN and not self._podman_available:
            raise RuntimeError("Podman is not available")

        # Create the environment
        if config.isolation_type == IsolationType.TMUX:
            env = self._create_tmux_environment(agent_id, workspace_path, config)
        else:
            env = self._create_container_environment(agent_id, workspace_path, config)

        self.environments[agent_id] = env
        logger.info(
            f"Created {config.isolation_type.value} environment for agent {agent_id}"
        )

        return env

    def _create_tmux_environment(
        self,
        agent_id: str,
        workspace_path: Path,
        config: IsolationConfig,
    ) -> IsolatedEnvironment:
        """Create a tmux-based environment."""
        session_name = f"agent-{agent_id}"

        # Check if session already exists
        result = subprocess.run(
            ["tmux", "has-session", "-t", session_name],
            capture_output=True,
        )

        if result.returncode != 0:
            # Create new session
            cmd = [
                "tmux", "new-session",
                "-d",  # Detached
                "-s", session_name,
                "-c", str(workspace_path),
            ]

            # Inject environment variables
            env_dict = config.credentials.get_env_dict()
            for key, value in env_dict.items():
                cmd.extend(["-e", f"{key}={value}"])

            subprocess.run(cmd, check=True, capture_output=True)

        return IsolatedEnvironment(
            agent_id=agent_id,
            isolation_type=IsolationType.TMUX,
            workspace_path=workspace_path,
            tmux_session=session_name,
            config=config,
        )

    def _create_container_environment(
        self,
        agent_id: str,
        workspace_path: Path,
        config: IsolationConfig,
    ) -> IsolatedEnvironment:
        """Create a container-based environment."""
        runtime = "docker" if config.isolation_type == IsolationType.DOCKER else "podman"
        container_name = f"agent-{agent_id}"

        # Build container run command
        cmd = [
            runtime, "run",
            "-d",  # Detached
            "--name", container_name,
            "-w", "/workspace",
            "-v", f"{workspace_path}:/workspace:rw",
        ]

        # Add resource limits
        cmd.extend(config.resource_limits.to_docker_args())

        # Network mode
        cmd.extend(["--network", config.network_mode.value])

        # Environment variables
        env_dict = config.credentials.get_env_dict()
        for key, value in env_dict.items():
            cmd.extend(["-e", f"{key}={value}"])

        # Git config injection
        if config.credentials.inject_git_config:
            git_config = Path.home() / ".gitconfig"
            if git_config.exists():
                cmd.extend(["-v", f"{git_config}:/root/.gitconfig:ro"])

        # SSH key injection
        if config.credentials.inject_ssh_keys:
            ssh_dir = Path.home() / ".ssh"
            if ssh_dir.exists():
                cmd.extend(["-v", f"{ssh_dir}:/root/.ssh:ro"])

        # AWS credentials injection
        if config.credentials.inject_aws_credentials:
            aws_dir = Path.home() / ".aws"
            if aws_dir.exists():
                cmd.extend(["-v", f"{aws_dir}:/root/.aws:ro"])

        # Extra mounts
        for host_path, container_path in config.extra_mounts:
            cmd.extend(["-v", f"{host_path}:{container_path}"])

        # Docker socket mount
        if config.mount_docker_socket:
            cmd.extend(["-v", "/var/run/docker.sock:/var/run/docker.sock"])

        # Privileged mode
        if config.privileged:
            cmd.append("--privileged")

        # Read-only root
        if config.read_only_root:
            cmd.append("--read-only")

        # Extra args
        cmd.extend(config.extra_args)

        # Image and command
        cmd.extend([config.container_image, "tail", "-f", "/dev/null"])

        # Remove existing container if present
        subprocess.run(
            [runtime, "rm", "-f", container_name],
            capture_output=True,
        )

        # Create container
        result = subprocess.run(cmd, check=True, capture_output=True)
        container_id = result.stdout.decode().strip()

        return IsolatedEnvironment(
            agent_id=agent_id,
            isolation_type=config.isolation_type,
            workspace_path=workspace_path,
            container_id=container_id,
            config=config,
        )

    def execute_command(
        self,
        agent_id: str,
        command: str,
        interactive: bool = False,
    ) -> Optional[str]:
        """
        Execute a command in an agent's environment.

        Args:
            agent_id: The agent's identifier
            command: Command to execute
            interactive: Whether to run interactively

        Returns:
            Command output (if not interactive)
        """
        if agent_id not in self.environments:
            raise ValueError(f"No environment found for agent: {agent_id}")

        env = self.environments[agent_id]

        if env.isolation_type == IsolationType.TMUX:
            return self._execute_tmux_command(env, command, interactive)
        else:
            return self._execute_container_command(env, command, interactive)

    def _execute_tmux_command(
        self,
        env: IsolatedEnvironment,
        command: str,
        interactive: bool,
    ) -> Optional[str]:
        """Execute command in tmux session."""
        if interactive:
            subprocess.run(
                ["tmux", "send-keys", "-t", env.tmux_session, command, "Enter"],
                check=True,
                capture_output=True,
            )
            return None
        else:
            # For non-interactive, send and capture
            subprocess.run(
                ["tmux", "send-keys", "-t", env.tmux_session, command, "Enter"],
                check=True,
                capture_output=True,
            )
            # Give it a moment and capture output
            import time
            time.sleep(0.5)
            return self.capture_output(env.agent_id, lines=50)

    def _execute_container_command(
        self,
        env: IsolatedEnvironment,
        command: str,
        interactive: bool,
    ) -> Optional[str]:
        """Execute command in container."""
        runtime = env.runtime_cmd
        container_name = f"agent-{env.agent_id}"

        cmd = [runtime, "exec"]
        if interactive:
            cmd.extend(["-it"])

        cmd.extend([container_name, "sh", "-c", command])

        if interactive:
            subprocess.run(cmd)
            return None
        else:
            result = subprocess.run(cmd, capture_output=True)
            return result.stdout.decode()

    def capture_output(self, agent_id: str, lines: int = 100) -> str:
        """
        Capture recent output from an agent's environment.

        Args:
            agent_id: The agent's identifier
            lines: Number of lines to capture

        Returns:
            Captured output
        """
        if agent_id not in self.environments:
            raise ValueError(f"No environment found for agent: {agent_id}")

        env = self.environments[agent_id]

        if env.isolation_type == IsolationType.TMUX:
            result = subprocess.run(
                ["tmux", "capture-pane", "-t", env.tmux_session, "-p", "-S", f"-{lines}"],
                capture_output=True,
            )
            return result.stdout.decode()
        else:
            # For containers, get logs
            runtime = env.runtime_cmd
            result = subprocess.run(
                [runtime, "logs", "--tail", str(lines), f"agent-{agent_id}"],
                capture_output=True,
            )
            return result.stdout.decode()

    def is_running(self, agent_id: str) -> bool:
        """Check if an agent's environment is running."""
        if agent_id not in self.environments:
            return False

        env = self.environments[agent_id]

        if env.isolation_type == IsolationType.TMUX:
            result = subprocess.run(
                ["tmux", "has-session", "-t", env.tmux_session],
                capture_output=True,
            )
            return result.returncode == 0
        else:
            runtime = env.runtime_cmd
            result = subprocess.run(
                [runtime, "inspect", "-f", "{{.State.Running}}", f"agent-{agent_id}"],
                capture_output=True,
            )
            return result.stdout.decode().strip() == "true"

    def stop(self, agent_id: str, force: bool = False) -> None:
        """
        Stop an agent's environment.

        Args:
            agent_id: The agent's identifier
            force: Force stop (kill)
        """
        if agent_id not in self.environments:
            return

        env = self.environments[agent_id]

        if env.isolation_type == IsolationType.TMUX:
            subprocess.run(
                ["tmux", "kill-session", "-t", env.tmux_session],
                capture_output=True,
            )
        else:
            runtime = env.runtime_cmd
            cmd = [runtime, "kill" if force else "stop", f"agent-{agent_id}"]
            subprocess.run(cmd, capture_output=True)

        logger.info(f"Stopped environment for agent {agent_id}")

    def cleanup(self, agent_id: str) -> None:
        """
        Fully clean up an agent's environment.

        Args:
            agent_id: The agent's identifier
        """
        if agent_id not in self.environments:
            return

        env = self.environments[agent_id]

        # Stop first
        self.stop(agent_id, force=True)

        # Remove container if applicable
        if env.is_container:
            runtime = env.runtime_cmd
            subprocess.run(
                [runtime, "rm", "-f", f"agent-{agent_id}"],
                capture_output=True,
            )

        del self.environments[agent_id]
        logger.info(f"Cleaned up environment for agent {agent_id}")

    def cleanup_all(self) -> None:
        """Clean up all environments."""
        agent_ids = list(self.environments.keys())
        for agent_id in agent_ids:
            self.cleanup(agent_id)

    def list_environments(self) -> list[IsolatedEnvironment]:
        """List all active environments."""
        return list(self.environments.values())

    def get_environment(self, agent_id: str) -> Optional[IsolatedEnvironment]:
        """Get a specific environment."""
        return self.environments.get(agent_id)

    def get_stats(self, agent_id: str) -> Optional[dict[str, Any]]:
        """
        Get resource usage stats for an agent's environment.

        Only available for container environments.
        """
        if agent_id not in self.environments:
            return None

        env = self.environments[agent_id]

        if not env.is_container:
            return {"type": "tmux", "stats": "not available for tmux sessions"}

        runtime = env.runtime_cmd
        result = subprocess.run(
            [runtime, "stats", "--no-stream", "--format", "json", f"agent-{agent_id}"],
            capture_output=True,
        )

        if result.returncode == 0:
            try:
                return json.loads(result.stdout.decode())
            except json.JSONDecodeError:
                return None

        return None
