import docker
import socket
import select
import time

class DockerShell:
    def __init__(self, container_name, shell="/bin/bash"):
        self.container_name = container_name
        self.shell_cmd = shell
        self.client = docker.from_env()
        self.api_client = self.client.api
        self.sock = None
        self._connect_shell()

    def _connect_shell(self):
        container = self.client.containers.get(self.container_name)
        if container.status != "running":
            container.start()
            container.reload()

        exec_id = self.api_client.exec_create(
            container.id,
            cmd=self.shell_cmd,
            tty=True,
            stdin=True,
            stdout=True,
            stderr=True,
            privileged=True
        )["Id"]

        self.sock = self.api_client.exec_start(exec_id, tty=True, stream=True, socket=True)

        # Flush welcome message
        time.sleep(0.5)
        self._read_output()

    def send_command(self, command: str, timeout=2.0) -> str:
        if not self.sock:
            raise RuntimeError("Shell socket not initialized.")

        raw_sock = self.sock._sock  # unwrap the true socket
        raw_sock.send((command + "\n").encode("utf-8"))

        return self._read_output(timeout=timeout)

    def _read_output(self, timeout=2.0) -> str:
        raw_sock = self.sock._sock  # unwrap the real socket
        raw_sock.setblocking(0)

        output = b""
        start_time = time.time()

        while True:
            ready, _, _ = select.select([raw_sock], [], [], 0.1)
            if ready:
                try:
                    chunk = raw_sock.recv(1024)
                    if not chunk:
                        break
                    output += chunk
                except socket.error:
                    break

            if time.time() - start_time > timeout:
                break

        return output.decode("utf-8", errors="ignore")

    def close(self):
        if self.sock:
            self.sock.close()
            self.sock = None


if __name__ == "__main__":
    shell = DockerShell("repocod_astropy")

    print(shell.send_command("cd usr/src/app/"))
    print('================================')
    print(shell.send_command("ls"))
    

    shell.close()
    
    
    
