# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
from __future__ import annotations

import subprocess
from abc import ABC, abstractmethod
from contextlib import contextmanager
from threading import Event, Thread
from typing import IO, TYPE_CHECKING, Any, cast

from airbyte import exceptions as exc
from airbyte._message_iterators import AirbyteMessageIterator


if TYPE_CHECKING:
    from collections.abc import Generator, Iterable, Iterator

    from airbyte.registry import ConnectorMetadata


_LATEST_VERSION = "latest"


class ExceptionHolder:
    def __init__(self) -> None:
        self.exception: Exception | None = None
        self.event = Event()

    def set_exception(
        self,
        ex: Exception,
    ) -> None:
        self.exception = ex
        self.event.set()  # Signal that an exception has occurred


def _pump_input(
    pipe: IO[str],
    messages: AirbyteMessageIterator,
    exception_holder: ExceptionHolder,
) -> None:
    """Pump lines into a pipe."""
    with pipe:
        try:
            pipe.writelines(message.model_dump_json() + "\n" for message in messages)
            pipe.flush()  # Ensure data is sent immediately
        except (BrokenPipeError, OSError) as ex:
            if isinstance(ex, BrokenPipeError):
                pass  # Expected during graceful shutdown
            else:
                exception_holder.set_exception(ex)
        except Exception as ex:
            exception_holder.set_exception(ex)


def _stream_from_file(file: IO[str]) -> Generator[str, Any, None]:
    """Stream lines from a file."""
    while True:
        line = file.readline()
        if not line:
            break
        yield line


@contextmanager
def _stream_from_subprocess(
    args: list[str],
    *,
    stdin: IO[str] | AirbyteMessageIterator | None = None,
    log_file: IO[str] | None = None,
    suppress_stderr: bool = False,
) -> Generator[Iterable[str], None, None]:
    """Stream lines from a subprocess.

    When stdin is an AirbyteMessageIterator, input is pumped to the subprocess
    in a separate thread while output is read concurrently. This avoids a
    potential deadlock where the subprocess blocks on stdout (buffer full)
    while we're waiting for input to finish before reading stdout.
    """
    input_thread: Thread | None = None
    exception_holder = ExceptionHolder()
    if isinstance(stdin, AirbyteMessageIterator):
        stderr_target = subprocess.DEVNULL if suppress_stderr else log_file
        process = subprocess.Popen(
            args,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=stderr_target,
            universal_newlines=True,
            encoding="utf-8",
        )
        input_thread = Thread(
            target=_pump_input,
            args=(
                process.stdin,
                stdin,
                exception_holder,
            ),
            daemon=True,  # Prevent blocking interpreter shutdown if thread gets stuck
        )
        input_thread.start()
        # Don't join here - let input and output happen concurrently to avoid deadlock.
        # The input thread will be joined in the finally block after reading is done.

    else:
        # stdin is None or a file-like object
        stderr_target = subprocess.DEVNULL if suppress_stderr else log_file
        process = subprocess.Popen(
            args,
            stdin=stdin,
            stdout=subprocess.PIPE,
            stderr=stderr_target,
            universal_newlines=True,
            encoding="utf-8",
        )

    if process.stdout is None:
        raise exc.AirbyteSubprocessError(
            message="Subprocess did not return a stdout stream.",
            context={
                "args": args,
                "returncode": process.returncode,
            },
        )

    try:
        yield _stream_from_file(process.stdout)
        process.wait()
    finally:
        try:
            # Terminate the process if it is still running
            if process.poll() is None:  # Check if the process is still running
                process.terminate()
                try:
                    # Wait for a short period to allow process to terminate gracefully
                    process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    # If the process does not terminate within the timeout, force kill it
                    process.kill()

            # Join the input thread if it exists (after process termination so stdin closes)
            if input_thread is not None:
                input_thread.join(timeout=10)

            # Now, the process is either terminated or killed. Check the exit code.
            exit_code = process.wait()

            # If the exit code is not 0 or -15 (SIGTERM), raise AirbyteSubprocessFailedError.
            # Include input thread exception as original_exception if present.
            if exit_code not in {0, -15}:
                raise exc.AirbyteSubprocessFailedError(
                    run_args=args,
                    exit_code=exit_code,
                    original_exception=(
                        exception_holder.exception
                        if exception_holder.exception
                        and not isinstance(exception_holder.exception, BrokenPipeError)
                        else None
                    ),
                )

            # Only raise input thread exception if exit code was OK (ignore BrokenPipeError)
            if exception_holder.exception and not isinstance(
                exception_holder.exception, BrokenPipeError
            ):
                raise exception_holder.exception
        finally:
            # Close the stdout stream
            if process.stdout:
                process.stdout.close()


class Executor(ABC):
    def __init__(
        self,
        *,
        name: str | None = None,
        metadata: ConnectorMetadata | None = None,
        target_version: str | None = None,
    ) -> None:
        """Initialize a connector executor.

        The 'name' param is required if 'metadata' is None.
        """
        if not name and not metadata:
            raise exc.PyAirbyteInternalError(message="Either name or metadata must be provided.")

        self.name: str = (
            name or cast("ConnectorMetadata", metadata).name
        )  # metadata is not None here
        self.metadata: ConnectorMetadata | None = metadata
        self.enforce_version: bool = target_version is not None

        self.reported_version: str | None = None
        self.target_version: str | None = None
        if target_version:
            if metadata and target_version == _LATEST_VERSION:
                self.target_version = metadata.latest_available_version
            else:
                self.target_version = target_version

    @property
    @abstractmethod
    def _cli(self) -> list[str]:
        """Get the base args of the CLI executable.

        Args will be appended to this list.
        """
        ...

    def map_cli_args(self, args: list[str]) -> list[str]:
        """Map CLI args if needed.

        By default, this is a no-op. Subclasses may override this method in order to
        map CLI args into the format expected by the connector.
        """
        return args

    def execute(
        self,
        args: list[str],
        *,
        stdin: IO[str] | AirbyteMessageIterator | None = None,
        suppress_stderr: bool = False,
    ) -> Iterator[str]:
        """Execute a command and return an iterator of STDOUT lines.

        If stdin is provided, it will be passed to the subprocess as STDIN.
        If suppress_stderr is True, stderr output will be suppressed to reduce noise.
        """
        mapped_args = self.map_cli_args(args)
        with _stream_from_subprocess(
            [*self._cli, *mapped_args],
            stdin=stdin,
            suppress_stderr=suppress_stderr,
        ) as stream_lines:
            yield from stream_lines

    @abstractmethod
    def ensure_installation(self, *, auto_fix: bool = True) -> None:
        _ = auto_fix
        pass

    @abstractmethod
    def install(self) -> None:
        pass

    @abstractmethod
    def uninstall(self) -> None:
        pass

    def get_installed_version(
        self,
        *,
        raise_on_error: bool = False,
        recheck: bool = False,
    ) -> str | None:
        """Detect the version of the connector installed."""
        _ = raise_on_error, recheck  # Unused
        raise NotImplementedError(
            f"'{type(self).__name__}' class cannot yet detect connector versions."
        )
