import unittest
import time
from sia.docker_module import DockerModule

class DockerModuleTest(unittest.TestCase):
    def test_initialization(self):
        """Test DockerModule initialization."""
        docker_module = DockerModule()
        self.assertIsInstance(docker_module, DockerModule)

    def test_start_container_short_lived(self):
        """Test starting a short-lived container."""
        docker_module = DockerModule()
        
        # Test parameters
        image = "busybox:latest"
        timeout = 1000
        command = "echo"
        arguments = ["Hello World"]
        
        output = docker_module.start_container(
            image=image,
            timeout=timeout,
            command=command,
            arguments=arguments
        )
        
        # Verify output was returned
        self.assertEqual(output, "Hello World\n")

    def test_start_container_long_running(self):
        """Test starting a long-running container."""
        docker_module = DockerModule()
        
        name = "test_start_container_long_running"
        image = "busybox:latest"
        
        container_name = docker_module.start_container(
            image=image,
            name=name
        )
        
        # Verify container was started and stored
        self.assertEqual(container_name, name)
        self.assertIn(name, docker_module.containers)

    def test_start_container_with_volumes(self):
        """Test starting a container with volume mappings."""
        docker_module = DockerModule()

        arguments = ["sh", "-c", "echo 'Hello World' > /write_here/test_start_container_with_volumes"]
        volumes = {"/tmp": "/write_here"}

        output = docker_module.start_container(
            image="busybox:latest",
            timeout=1000,
            arguments=arguments,
            volumes=volumes
        )
        print(f"output: {output}", flush=True)
        # Verify output was written to volume
        with open("/tmp/test_start_container_with_volumes", "r") as f:
            self.assertEqual(f.read(), "Hello World\n")

    def test_invalid_container_start(self):
        """Test error handling for invalid container start."""
        docker_module = DockerModule()
        
        # Test missing required parameters
        with self.assertRaises(ValueError):
            docker_module.start_container(image="busybox:latest")

    def test_stop_container(self):
        """Test stopping a container."""
        docker_module = DockerModule()
        
        # Start a container first
        name = "test_stop_container"
        docker_module.start_container(
            image="busybox:latest",
            name=name
        )
        
        # Stop the container
        docker_module.stop_container(name)
        
        # Verify container was stopped and removed
        self.assertNotIn(name, docker_module.containers)

    def test_stop_nonexistent_container(self):
        """Test stopping a container that doesn't exist."""
        docker_module = DockerModule()
        
        with self.assertRaises(KeyError):
            docker_module.stop_container("nonexistent")

    def test_container_io_operations(self):
        """Test container I/O operations."""
        docker_module = DockerModule()

        arguments = [
            "sh",
            "-c",
            "read input; echo \"$input\" | tr '[:lower:]' '[:upper:]'; echo \"$input\" | tr '[:upper:]' '[:lower:]' >&2"
        ]
        
        # Start a container
        name = "test_container_io_operations"
        docker_module.start_container(
            image="busybox:latest",
            name=name,
            arguments=arguments
        )
        
        # Test stdin write
        test_input = "Test Input\n"
        docker_module.write_container_stdin(name, test_input)
        time.sleep(1) # for processing
        
        # Test stdout read
        stdout = docker_module.read_container_stdout(name)
        print(f"stdout: {stdout}", flush=True)
        self.assertEqual(stdout, "test input\n")
        
        # Test stderr read
        stderr = docker_module.read_container_stderr(name)
        print(f"stderr: {stderr}", flush=True)
        self.assertEqual(stderr, "TEST INPUT\n")

#    def test_wait_container(self):
#        """Test waiting for container completion."""
#        docker_module = DockerModule()
#        
#        # Start a container
#        name = "test-container"
#        docker_module.start_container(
#            image="test:latest",
#            name=name
#        )
#        
#        # Test wait operation
#        exit_code, output = docker_module.wait_container(name, timeout=1000)
#        
#        # Verify wait was called with correct timeout
#        self.mock_container.wait.assert_called_once_with(timeout=1.0)
#        self.assertEqual(exit_code, 0)
#        self.assertEqual(output, "Test output")
#
#    def test_get_container_status(self):
#        """Test getting container status information."""
#        docker_module = DockerModule()
#        
#        # Start a container
#        name = "test-container"
#        docker_module.start_container(
#            image="test:latest",
#            name=name
#        )
#        
#        # Configure mock logs sizes
#        self.mock_container.logs.side_effect = [b"stdout", b"stderr"]
#        
#        # Get status
#        status = docker_module.get_container_status(name)
#        
#        # Verify status information
#        self.assertEqual(status['name'], name)
#        self.assertEqual(status['status'], "running")
#        self.assertEqual(status['started_at'], "2024-10-23T10:00:00Z")
#        self.assertEqual(status['exit_code'], 0)
#        self.assertEqual(status['error'], "")
#        self.assertEqual(status['stdout_size'], 6)  # len(b"stdout")
#        self.assertEqual(status['stderr_size'], 6)  # len(b"stderr")
#        
#        # Verify container was reloaded
#        self.mock_container.reload.assert_called_once()

if __name__ == '__main__':
    unittest.main()