import asyncio
import os

try:
    import pty
except Exception as ex:
    print("You are not able to run a terminal. See error:")
    print(ex)
import subprocess

commands = {
    "alpine": "docker run -it alpine /bin/sh",
    "r": "docker run -v /usr/local/bin:/usr/local/bin -it ubuntu:latest run.sh",
}


class TerminalSession:
    def __init__(self, command):
        self.master, self.slave = None, None
        self.process = None
        self.sockets = []
        self.history = b""
        self.history_size = 1024 * 20
        self.command = command
        self.start_process(self.command)

    def start_process(self, command):
        if not self.is_running():
            if self.master:
                os.close(self.master)
                os.close(self.slave)
                self.master = None
                self.slave = None

            self.master, self.slave = pty.openpty()
            self.process = subprocess.Popen(
                command.split(" "),
                stdin=self.slave,
                stdout=self.slave,
                stderr=self.slave,
                bufsize=0,
                universal_newlines=True,
            )

    def is_running(self):
        if not self.process:
            return False
        asyncio.get_event_loop()
        return self.process.poll() is None

    async def add_websocket(self, ws):
        self.start_process(self.command)
        asyncio.create_task(self.read_output(ws))

    async def read_output(self, ws):
        self.sockets.append(ws)
        if len(self.sockets) > 1 and self.history:
            start = 0
            try:
                start = self.history.index(b"\n")
            except ValueError:
                pass
            await ws.send_bytes(self.history[start:])
            return
        loop = asyncio.get_event_loop()
        while True:
            try:
                data = await loop.run_in_executor(None, os.read, self.master, 1024)
                if not data:
                    break
                self.history += data
                if len(self.history) > self.history_size:
                    self.history = self.history[: 0 - self.history_size]
                try:
                    for ws in self.sockets:
                        await ws.send_bytes(data)  # Send raw bytes for ANSI support
                except:
                    self.sockets.remove(ws)
            except Exception:
                await self.close()
                break

    async def close(self):
        print("Terminating process")
        if self.process:
            self.process.terminate()
            self.process = None
        if self.master:
            os.close(self.master)
            os.close(self.slave)
            self.master = None
            self.slave = None

        print("Terminated process")
        for ws in self.sockets:
            try:
                await ws.close()
            except Exception:
                pass
        self.sockets = []

    async def write_input(self, data):
        try:
            data = data.encode()
        except AttributeError:
            pass
        try:
            await asyncio.get_event_loop().run_in_executor(
                None, os.write, self.master, data
            )
        except Exception as ex:
            print(ex)
            await self.close()
