import sys
import contextlib
import typing
import threading
import collections


class _MultiStream:
    def __init__(
        self,
        map: typing.Mapping[threading.Thread, list[typing.TextIO]],
        default: typing.TextIO,
    ) -> None:
        self._map = map
        self._default = [default]

    def __getattr__(self, attr):
        sseq = self._map.get(threading.current_thread(), []) or self._default
        stream = sseq[0]
        return getattr(stream, attr)


class _IoReplacer:
    def __init__(self, stream_name: str) -> None:
        self.stream_name = stream_name
        self.map: collections.defaultdict[threading.Thread, list[typing.TextIO]] = (
            collections.defaultdict(list)
        )
        self.old_stream: typing.TextIO = getattr(sys, self.stream_name)
        self.lock = threading.Lock()
        self.stream = _MultiStream(self.map, self.old_stream)

    def push(self, target: typing.TextIO) -> None:
        thread = threading.current_thread()
        with self.lock:
            if not self.map:
                setattr(sys, self.stream_name, self.stream)
            self.map[thread].append(target)

    def pop(self) -> None:
        thread = threading.current_thread()
        with self.lock:
            self.map[thread].pop()
            if not self.map[thread]:
                del self.map[thread]
            if not self.map:
                setattr(sys, self.stream_name, self.old_stream)

    @contextlib.contextmanager
    def __call__(self, new_target: typing.TextIO) -> typing.Iterator[None]:
        try:
            self.push(new_target)
            yield
        finally:
            assert self.map[threading.current_thread()][-1] is new_target
            self.pop()
            self.push


redirect_stdout = _IoReplacer("stdout")
redirect_stderr = _IoReplacer("stderr")