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

class DockerModule:
    """
    Module for handling Docker container operations in SIA.
    Manages container lifecycle, I/O operations, and status monitoring.
    """
    
    def _start_docker_engine():
        """
        Start the Docker engine
        """
        subprocess.run(['/etc/init.d/docker', 'start'])
        time.sleep(1)

    def __init__(self):
        DockerModule._start_docker_engine()
        self.client = docker.DockerClient(base_url='unix://var/run/docker.sock')
        self.containers: Dict[str, (Container, docker.APIClient.AttachSocket)] = {}
        self.logger = logging.getLogger(__name__)

    def __enter__(self):
        return self

    def __exit__(self, exec_type, exec_value, traceback):
        del self
    
    def __del__(self):
        """Cleanup all resources before shutdown."""
        try:
            # Stop and remove all containers
            for (container, socket) in self.containers.values():
                self._wait_and_stop(container, socket, 0)
            # Close the Docker client
            if hasattr(self, 'client'):
                self.client.close()
        except Exception as e:
            self.logger.error(f"Error during cleanup: {str(e)}")

    def _wait_and_stop(self, container: Container, socket, timeout: int) -> Tuple[int, str]:
        """Helper method for waiting on long- and short-lived containers."""
        status_code = None
        if timeout > 0:
            try:
                result = container.wait(timeout=timeout/1000)
                status_code = result.get('StatusCode')
            except requests.exceptions.ReadTimeout as e:
                container.stop()
                status_code = -1
            except requests.exceptions.ConnectionError as e:
                container.stop()
                status_code = -1
        else:
            container.stop()
            status_code = -1
        try:
            socket.close()
        except Exception as e:
            self.logger.error(f"Error closing socket: {str(e)}")
        output = container.logs(stdout=True, stderr=True).decode('utf-8')
        try:
            container.remove(force=True)
        except Exception as e:
            self.logger.error(f"Error removing container: {str(e)}")
        return status_code, output

    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,
    ) -> Optional[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: None
        """
        # Validate inputs
        if name is None and timeout < 0:
            raise ValueError("Either name or timeout must be provided")
        cmd = []
        if command:
            cmd.append(command)
        if arguments:
            cmd.extend(arguments)
        volume_binds = {}
        if volumes:
            for host_path, container_path in volumes.items():
                volume_binds[host_path] = {'bind': container_path, 'mode': 'rw'}
        port_bindings = {}
        if ports:
            for host_port, container_port in ports.items():
                port_bindings[container_port] = host_port
        self.client.images.pull(image)
        container = self.client.containers.create(
            image,
            command=cmd if cmd else None,
            name=name,
            detach=True,
            volumes=volume_binds,
            ports=port_bindings,
            environment=environment,
            tty=False,
            stdin_open=True
        )
        socket = container.attach_socket(params={
            'stdin': 1,
            'stdout': 1,
            'stderr': 1,
            'stream': 1,
            'demux': True
        })
        container.start()
        if timeout >= 0:
            # For short-lived containers
            return self._wait_and_stop(container, socket, timeout)[1]
        else:
            # For long-running containers
            self.containers[name] = (container, socket)

    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:
            _, socket = self.containers[name]
            socket._sock.send(data.encode('utf-8'))
        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, socket = 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, socket = 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)
        """
        if name not in self.containers:
            raise KeyError(f"Container {name} not found")

        container, socket = self.containers.pop(name)
        return self._wait_and_stop(container, socket, timeout)

    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, socket = 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
    
    def get_all_container_statuses(self) -> List[Dict[str, any]]:
        """Get current status information for all containers.
        Returns:
            List of dictionaries with container status information
        """
        statuses = []
        for name, (container, socket) in self.containers.items():
            statuses.append(self.get_container_status(name))
        return statuses