from typing import Dict, List, Optional, Tuple
import docker
from docker.models.containers import Container
import logging

class DockerModule:
    """
    Module for handling Docker container operations in SIA.
    Manages container lifecycle, I/O operations, and status monitoring.
    """

    def __init__(self):
        self.client = docker.from_env()
        self.containers: Dict[str, Container] = {}
        self.logger = logging.getLogger(__name__)
        
    def start_container(
        self,
        image: str,
        name: Optional[str] = None,
        timeout: int = -1,
        command: Optional[str] = None,
        arguments: Optional[List[str]] = None,
        volumes: Optional[Dict[str, str]] = None,
        ports: Optional[Dict[str, str]] = None,
        environment: Optional[Dict[str, str]] = None,
    ) -> str:
        """Start a new Docker container with the specified configuration.

        Args:
            image: Docker image to use
            name: Unique container name for long running containers
            timeout: Timeout in milliseconds for short running containers
            command: Main command to run in container
            arguments: List of command line arguments
            volumes: Dictionary mapping host paths to container paths
            ports: Dictionary mapping host ports to container ports
            environment: Dictionary of environment variables and values

        Returns:
            For short-lived containers (with timeout): Container output
            For long-running containers: Container name

        Raises:
            docker.errors.ImageNotFound: If specified image doesn't exist
            docker.errors.APIError: If container creation/start fails
        """
        # Validate inputs
        if name is None and timeout < 0:
            raise ValueError("Either name or timeout must be provided")

        # Prepare command with arguments
        cmd = []
        if command:
            cmd.append(command)
        if arguments:
            cmd.extend(arguments)

        # Prepare volume bindings
        volume_binds = {}
        if volumes:
            for host_path, container_path in volumes.items():
                volume_binds[host_path] = {'bind': container_path, 'mode': 'rw'}

        # Prepare port bindings
        port_bindings = {}
        if ports:
            for host_port, container_port in ports.items():
                port_bindings[container_port] = host_port

        try:
            container = self.client.containers.run(
                image,
                command=cmd if cmd else None,
                name=name,
                detach=True,
                volumes=volume_binds,
                ports=port_bindings,
                environment=environment,
            )

            if timeout >= 0:
                # For short-lived containers, wait for completion
                try:
                    container.wait(timeout=timeout/1000)  # Convert ms to seconds
                    output = container.logs().decode('utf-8')
                    container.remove(force=True)
                    return output
                except docker.errors.NotFound:
                    self.logger.warning(f"Container was removed before timeout")
                    return ""
            else:
                # For long-running containers, store reference
                self.containers[name] = container
                return name

        except Exception as e:
            self.logger.error(f"Error starting container: {str(e)}")
            raise

    def stop_container(self, name: str) -> None:
        """Stop and remove a Docker container by name.

        Args:
            name: Name of the container to stop

        Raises:
            KeyError: If container name not found
            docker.errors.APIError: If container stop/removal fails
        """
        if name not in self.containers:
            raise KeyError(f"Container {name} not found")

        try:
            container = self.containers[name]
            container.stop()
            container.remove()
            del self.containers[name]
        except Exception as e:
            self.logger.error(f"Error stopping container {name}: {str(e)}")
            raise

    def write_container_stdin(self, name: str, data: str) -> None:
        """Write data to a container's standard input.

        Args:
            name: Name of the target container
            data: Data to write to stdin

        Raises:
            KeyError: If container name not found
            docker.errors.APIError: If write operation fails
        """
        if name not in self.containers:
            raise KeyError(f"Container {name} not found")

        try:
            container = self.containers[name]
            socket = container.attach_socket(params={'stdin': 1, 'stream': 1})
            socket._sock.send(data.encode('utf-8'))
            socket.close()
        except Exception as e:
            self.logger.error(f"Error writing to container {name} stdin: {str(e)}")
            raise

    def read_container_stdout(self, name: str, n: int = -1) -> str:
        """Read from a container's standard output buffer.

        Args:
            name: Name of the container
            n: Number of bytes to read; -1 means read all available

        Returns:
            Data read from stdout

        Raises:
            KeyError: If container name not found
            docker.errors.APIError: If read operation fails
        """
        if name not in self.containers:
            raise KeyError(f"Container {name} not found")

        try:
            container = self.containers[name]
            logs = container.logs(stdout=True, stderr=False, tail=n if n > 0 else 'all')
            return logs.decode('utf-8')
        except Exception as e:
            self.logger.error(f"Error reading container {name} stdout: {str(e)}")
            raise

    def read_container_stderr(self, name: str, n: int = -1) -> str:
        """Read from a container's standard error buffer.

        Args:
            name: Name of the container
            n: Number of bytes to read; -1 means read all available

        Returns:
            Data read from stderr

        Raises:
            KeyError: If container name not found
            docker.errors.APIError: If read operation fails
        """
        if name not in self.containers:
            raise KeyError(f"Container {name} not found")

        try:
            container = self.containers[name]
            logs = container.logs(stdout=False, stderr=True, tail=n if n > 0 else 'all')
            return logs.decode('utf-8')
        except Exception as e:
            self.logger.error(f"Error reading container {name} stderr: {str(e)}")
            raise

    def wait_container(self, name: str, timeout: int) -> Tuple[int, str]:
        """Wait for a container to finish execution.

        Args:
            name: Name of the container to wait for
            timeout: Time to wait in milliseconds

        Returns:
            Tuple of (exit_code, output)

        Raises:
            KeyError: If container name not found
            docker.errors.APIError: If wait operation fails
            TimeoutError: If container doesn't finish within timeout
        """
        if name not in self.containers:
            raise KeyError(f"Container {name} not found")

        try:
            container = self.containers[name]
            result = container.wait(timeout=timeout/1000)  # Convert ms to seconds
            logs = container.logs().decode('utf-8')
            return (result['StatusCode'], logs)
        except Exception as e:
            self.logger.error(f"Error waiting for container {name}: {str(e)}")
            raise

    def get_container_status(self, name: str) -> Dict[str, any]:
        """Get current status information for a container.

        Args:
            name: Name of the container

        Returns:
            Dictionary with container status information

        Raises:
            KeyError: If container name not found
            docker.errors.APIError: If status check fails
        """
        if name not in self.containers:
            raise KeyError(f"Container {name} not found")

        try:
            container = self.containers[name]
            container.reload()  # Refresh container info
            
            return {
                'name': name,
                'status': container.status,
                'started_at': container.attrs['State']['StartedAt'],
                'exit_code': container.attrs['State']['ExitCode'],
                'error': container.attrs['State']['Error'],
                'stdout_size': len(container.logs(stdout=True, stderr=False)),
                'stderr_size': len(container.logs(stdout=False, stderr=True))
            }
        except Exception as e:
            self.logger.error(f"Error getting status for container {name}: {str(e)}")
            raise