import os
import sys
from io import StringIO
from tempfile import TemporaryFile
from typing import Self, Any, IO

from .comparison import StringComparison, compare


class OutputCapture:
    """
    A context manager for capturing output to the
    :any:`sys.stdout` and :any:`sys.stderr` streams.

    :param separate: If ``True``, ``stdout`` and ``stderr`` will be captured
                     separately and their expected values must be passed to
                     :meth:`~OutputCapture.compare`.

    :param fd: If ``True``, the underlying file descriptors will be captured,
               rather than just the attributes on :mod:`sys`. This allows
               you to capture things like subprocesses that write directly
               to the file descriptors, but is more invasive, so only use it
               when you need it.

    :param strip_whitespace:
        When ``True``, which is the default, leading and training whitespace
        is trimmed from both the expected and actual values when comparing.

    .. note:: If ``separate`` is passed as ``True``,
              :attr:`OutputCapture.captured` will be an empty string.
    """

    output: IO
    stdout: IO
    stderr: IO

    original_stdout: IO[str] | int | None = None
    original_stderr: IO[str] | int | None = None

    def __init__(self, separate: bool = False, fd: bool = False, strip_whitespace: bool = True):
        self.separate = separate
        self.fd = fd
        self.strip_whitespace = strip_whitespace

    def __enter__(self) -> Self:
        if self.fd:
            self.output = TemporaryFile()
            self.stdout = TemporaryFile()
            self.stderr = TemporaryFile()
        else:
            self.output = StringIO()
            self.stdout = StringIO()
            self.stderr = StringIO()
        self.enable()
        return self

    def __exit__(self, *args: Any) -> None:
        self.disable()

    def disable(self) -> None:
        "Disable the output capture if it is enabled."
        if self.fd:
            for original, current in (
                (self.original_stdout, sys.stdout),
                (self.original_stderr, sys.stderr),
            ):
                os.dup2(original, current.fileno())  # type: ignore[arg-type]
                os.close(original)  # type: ignore[arg-type]

        else:
            sys.stdout = self.original_stdout
            sys.stderr = self.original_stderr

    def enable(self) -> None:
        "Enable the output capture if it is disabled."
        if self.original_stdout is None:
            if self.fd:
                self.original_stdout = os.dup(sys.stdout.fileno())
                self.original_stderr = os.dup(sys.stderr.fileno())
            else:
                self.original_stdout = sys.stdout
                self.original_stderr = sys.stderr
        if self.separate:
            if self.fd:
                os.dup2(self.stdout.fileno(), sys.stdout.fileno())
                os.dup2(self.stderr.fileno(), sys.stderr.fileno())
            else:
                sys.stdout = self.stdout
                sys.stderr = self.stderr
        else:
            if self.fd:
                os.dup2(self.output.fileno(), sys.stdout.fileno())
                os.dup2(self.output.fileno(), sys.stderr.fileno())
            else:
                sys.stdout = sys.stderr = self.output

    def _read(self, stream: IO | StringIO) -> str:
        if self.fd:
            stream.seek(0)
            return stream.read().decode()  # type: ignore[union-attr]
        else:
            return stream.getvalue()  # type: ignore[union-attr]

    @property
    def captured(self) -> str:
        "A property containing any output that has been captured so far."
        return self._read(self.output)

    def compare(
            self,
            expected: str | StringComparison = '',
            stdout: str | StringComparison = '',
            stderr: str | StringComparison = '',
            raises: bool = True,
    ) -> str | None:
        """
        Compare the captured output to that expected. If the output is
        not the same, an :class:`AssertionError` will be raised.

        :param expected: A string containing the expected combined output
                         of ``stdout`` and ``stderr``.

        :param stdout: A string containing the expected output to ``stdout``.

        :param stderr: A string containing the expected output to ``stderr``.

        :param raises: If ``False``, the message that would be raised in the
                       :class:`AssertionError` will be returned instead of the
                       exception being raised.
        """
        __tracebackhide__ = True
        expected_mapping = {}
        actual_mapping = {}
        for prefix, _expected, captured in (
                ('captured', expected, self.captured),
                ('stdout', stdout, self._read(self.stdout)),
                ('stderr', stderr, self._read(self.stderr)),
        ):
            if self.strip_whitespace:
                if isinstance(_expected, str):
                    _expected = _expected.strip()
                captured = captured.strip()
            if _expected != captured:
                expected_mapping[prefix] = _expected
                actual_mapping[prefix] = captured
        if len(expected_mapping) == 1:
            return compare(expected=tuple(expected_mapping.values())[0],
                           actual=tuple(actual_mapping.values())[0],
                           raises=raises)
        return compare(expected=expected_mapping, actual=actual_mapping, raises=raises)
