import sys
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from pathlib import Path


class Tee:
    def __init__(self, *streams):
        self.streams = streams

    def write(self, text):
        for stream in self.streams:
            stream.write(text)
            stream.flush()
        return len(text)

    def flush(self):
        for stream in self.streams:
            stream.flush()


@contextmanager
def redirect_output(logfile, mode="w", encoding="utf-8", include_stderr=True):
    path = Path(logfile)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open(mode, encoding=encoding) as fp:
        tee_out = Tee(sys.stdout, fp)
        if include_stderr:
            tee_err = Tee(sys.stderr, fp)
            with redirect_stdout(tee_out), redirect_stderr(tee_err):
                yield
        else:
            with redirect_stdout(tee_out):
                yield
