import fcntl
import os
import signal
import struct
import subprocess
import sys
import termios
import tty
from dataclasses import dataclass, field
from typing import List, Optional, TextIO, Tuple, Union
from shared_utils import hard_exit

tees = []


@dataclass
class Tee:
    orig_stdout: TextIO = field(init=False)
    stdout_pipe_proc: Optional[Tuple[TextIO, subprocess.Popen]] = field(init=False, default=None)

    orig_stderr: TextIO = field(init=False)
    stderr_pipe_proc: Optional[Tuple[TextIO, subprocess.Popen]] = field(init=False, default=None)

    def __post_init__(self):
        tees.append(self)  # Do not let this get GC'd
        self.orig_stdout = _dup(sys.stdout)
        self.orig_stderr = _dup(sys.stderr)
        hard_exit.register_atexit(self.close)

    def to(self, stdout, stderr):
        if not isinstance(stdout, list):
            stdout = [stdout]
        if not isinstance(stderr, list):
            stderr = [stderr]

        old_stdout_pipe_proc = self.stdout_pipe_proc
        self.stdout_pipe_proc = _tee(self.orig_stdout, stdout, stdout=self.orig_stdout)
        self.stdout = stdout
        old_stderr_pipe_proc = self.stderr_pipe_proc
        self.stderr_pipe_proc = _tee(self.orig_stderr, stderr, stdout=self.orig_stdout)
        self.stderr = stderr

        self._drain(stdout_pipe_proc=old_stdout_pipe_proc, stderr_pipe_proc=old_stderr_pipe_proc)
        self.resume()

        return self

    def close(self):
        self.pause()
        self._drain(self.stdout_pipe_proc, self.stderr_pipe_proc)

    def _drain(self, stdout_pipe_proc, stderr_pipe_proc):
        # One sharp edge is that if you've spawned a subprocess with
        # the redirected stdout/stderr, the tee processes will not
        # die. In that case maybe we should set a timeout, or just
        # leak them? Not sure.
        if stdout_pipe_proc is not None:
            pipe, proc = stdout_pipe_proc
            pipe.close()
            os.kill(signal.SIGINT, proc.pid)
            proc.wait()
        if stderr_pipe_proc is not None:
            pipe, proc = stderr_pipe_proc
            pipe.close()
            os.kill(signal.SIGINT, proc.pid)
            proc.wait()

    def resume(self):
        os.dup2(self.stdout_pipe_proc[0].fileno(), sys.stdout.fileno())
        os.dup2(self.stderr_pipe_proc[0].fileno(), sys.stderr.fileno())

    def pause(self):
        os.dup2(self.orig_stdout.fileno(), sys.stdout.fileno())
        os.dup2(self.orig_stderr.fileno(), sys.stderr.fileno())


def _tee(src, to, stdout):
    if src.isatty():
        w, r = os.openpty()
        flags = termios.tcgetattr(src.fileno())
        termios.tcsetattr(w, termios.TCSANOW, flags)
        termios.tcsetattr(r, termios.TCSANOW, flags)
    else:
        r, w = os.pipe()

    r = os.fdopen(r, "rb", buffering=0)
    w = os.fdopen(w, "wb", buffering=0)

    if src.isatty():
        tty.setraw(r)

        # Copy window size
        packed = fcntl.ioctl(sys.stdin, termios.TIOCGWINSZ, struct.pack("HHHH", 0, 0, 0, 0))
        fcntl.ioctl(r, termios.TIOCSWINSZ, packed)

        def set_ctty():
            signal.signal(signal.SIGINT, signal.SIG_IGN)
            w.close()
            fcntl.ioctl(r, termios.TIOCSCTTY, 0)

    else:

        def set_ctty():
            signal.signal(signal.SIGINT, signal.SIG_IGN)
            w.close()

    for path in to:
        os.makedirs(os.path.dirname(path), exist_ok=True)

    proc = subprocess.Popen(
        ["parent-lifetime", "--term", "tee", "-a"] + list(to),
        stdin=r,
        start_new_session=True,
        stderr=subprocess.DEVNULL,
        stdout=stdout,
        preexec_fn=set_ctty,
    )
    r.close()
    return w, proc


def _dup(src):
    src = os.dup(src.fileno())
    src = os.fdopen(src, "rb", buffering=0)
    return src


def tee(stdout: Union[str, List[str]], stderr: Union[List[str], str]):
    return Tee().to(stdout=stdout, stderr=stderr)


if __name__ == "__main__":
    t = Tee()
    t.to("/tmp/out.orig", "/tmp/err.orig")
    t.to("/tmp/out.new", "/tmp/err.new")
    breakpoint()
