"""Thread-safe utilities for concurrent execution with isolated logging."""

import contextlib
import io
import sys
import threading
from typing import Generator, TextIO, Type

from typing_extensions import override

from fused._udf.state import StdStreamProxy


class MultipleTextIO:
    _outputs: list[TextIO]

    def __init__(self, *args: list[TextIO]) -> None:
        self._outputs = args

    def write(self, *args, **kwargs) -> int:
        res: int = 0
        for out in self._outputs:
            res = out.write(*args, **kwargs)
        return res

    def flush(self, *args, **kwargs):
        res = None
        for out in self._outputs:
            res = out.flush(*args, **kwargs)
        return res

    def getvalue(self) -> str:
        for out in self._outputs:
            if isinstance(out, io.StringIO):
                return out.getvalue()
        return ""

    def isatty(self) -> bool:
        return False

    @property
    def ranges(self) -> list[tuple[int, int, int]]:
        for out in self._outputs:
            if hasattr(out, "ranges"):
                return out.ranges
        raise AttributeError("Ranges not found")

    def fileno(self) -> int:
        """Return the file descriptor from the first output that supports it."""
        for out in self._outputs:
            try:
                if hasattr(out, "fileno"):
                    return out.fileno()
            except (AttributeError, io.UnsupportedOperation):
                continue
        # If none of the outputs support fileno, raise proper exception
        raise io.UnsupportedOperation("fileno() not supported on any underlying stream")


class ThreadLocalStreamRouter(TextIO):
    """
    A file-like object that routes writes to thread-local buffers.

    This is used for concurrent UDF execution where each thread
    needs isolated stdout/stderr. It's installed globally as sys.stdout/sys.stderr
    and routes writes based on the calling thread's ID.

    Architecture:
    - Parent thread installs this router globally
    - Each worker thread registers its own buffer (stored in thread-local storage)
    - Writes from threads with registered buffers are routed to their buffer
    - Writes from threads without registered buffers fall back to the original stream
    """

    def __init__(self, original_stream: TextIO):
        self.original_stream = original_stream
        self._thread_local = threading.local()

    def _get_buffer(self) -> TextIO:
        buffer = getattr(self._thread_local, "buffer", None)
        if buffer is None:
            # Fall back to original stream for threads without registered buffers
            # (e.g., main thread or threads that don't need isolation)
            return self.original_stream
        return buffer

    def register_buffer_for_current_thread(self, buffer: TextIO):
        """Register a buffer for the current thread."""
        self._thread_local.buffer = buffer

    def unregister_buffer_for_current_thread(self):
        """Unregister the current thread's buffer."""
        if hasattr(self._thread_local, "buffer"):
            delattr(self._thread_local, "buffer")

    @override
    def write(self, text: str) -> int:
        """Write to the appropriate buffer based on current thread."""
        buffer = self._get_buffer()
        return buffer.write(text)

    @override
    def flush(self) -> None:
        """Flush current thread's buffer."""
        buffer = self._get_buffer()
        buffer.flush()

    @override
    def isatty(self) -> bool:
        return False

    @override
    def close(self) -> None:
        # prevent closing the underlying stream
        self.flush()

    @override
    def writable(self) -> bool:
        return True

    @override
    def readable(self) -> bool:
        return False

    @override
    def fileno(self) -> int:
        """Return the file descriptor of the original stream."""
        try:
            return self.original_stream.fileno()
        except (AttributeError, io.UnsupportedOperation):
            # Raise proper exception instead of returning None
            # This allows multiprocessing to catch and handle it correctly
            raise io.UnsupportedOperation("fileno() not supported on underlying stream")


@contextlib.contextmanager
def isolate_streams_per_thread() -> Generator[
    tuple[ThreadLocalStreamRouter, ThreadLocalStreamRouter], None, None
]:
    """
    Context manager that sets up thread-local stream routing globally.

    If routers are already installed (e.g., from a parent thread pool), this will
    reuse them instead of creating new ones. This allows nested usage without conflicts.
    """
    # Check if routers are already installed
    if isinstance(sys.stdout, ThreadLocalStreamRouter) and isinstance(
        sys.stderr, ThreadLocalStreamRouter
    ):
        # Reuse existing routers - no need to install/restore
        yield sys.stdout, sys.stderr
    else:
        # Install new routers
        original_stdout = sys.stdout
        original_stderr = sys.stderr

        stdout_router = ThreadLocalStreamRouter(original_stdout)
        stderr_router = ThreadLocalStreamRouter(original_stderr)

        # Install globally
        sys.stdout = stdout_router
        sys.stderr = stderr_router

        try:
            yield stdout_router, stderr_router
        finally:
            # Restore original streams
            sys.stdout = original_stdout
            sys.stderr = original_stderr


@contextlib.contextmanager
def capture_current_thread_output(
    output_to_original_streams: bool = False, impl_class: Type[TextIO] = io.StringIO
):
    """
    Context manager for worker threads to capture their own stdout/stderr.

    This should be used within a worker thread when isolate_streams_per_thread()
    is active in the parent.

    Usage:
        with capture_current_thread_output() as (out_buf, err_buf):
            # Do work
            print("This goes to out_buf")

            # Logs are automatically captured
            logs = out_buf.getvalue()

    Returns:
        Tuple of (stdout_buffer, stderr_buffer) as StringIO objects
    """

    # Ensure routers are installed
    stdout_router = sys.stdout
    stderr_router = sys.stderr
    if not isinstance(stdout_router, ThreadLocalStreamRouter) or not isinstance(
        stderr_router, ThreadLocalStreamRouter
    ):
        raise ValueError(
            "capture_current_thread_output() must be used in the worker thread after isolate_streams_per_thread()"
        )
    # Save previous buffers to support nested execution
    prev_stdout_buf = getattr(stdout_router._thread_local, "buffer", None)
    prev_stderr_buf = getattr(stderr_router._thread_local, "buffer", None)

    with impl_class() as out_buf, impl_class() as err_buf:
        # use MultipleTextIO to capture the output to the original streams
        if output_to_original_streams:
            out_buf = MultipleTextIO(out_buf, stdout_router.original_stream)
            err_buf = MultipleTextIO(err_buf, stderr_router.original_stream)
        # Register buffers for current thread (no thread_id needed!)
        stdout_router.register_buffer_for_current_thread(out_buf)
        stderr_router.register_buffer_for_current_thread(err_buf)

        try:
            with (
                # ensure that the standard streams are not closed by the UDF code
                StdStreamProxy("stdin"),
                StdStreamProxy("__stdin__"),
                StdStreamProxy("__stdout__"),
                StdStreamProxy("__stderr__"),
            ):
                yield out_buf, err_buf
        finally:
            # Restore previous buffers (or unregister if there were none)
            if prev_stdout_buf is not None:
                stdout_router.register_buffer_for_current_thread(prev_stdout_buf)
            else:
                stdout_router.unregister_buffer_for_current_thread()

            if prev_stderr_buf is not None:
                stderr_router.register_buffer_for_current_thread(prev_stderr_buf)
            else:
                stderr_router.unregister_buffer_for_current_thread()
