import errno
import gc
import os
import shutil
import sys
import tempfile
import textwrap
import time

import anyio
import pytest
from anyio.abc import Process

from mcp.client.session import ClientSession
from mcp.client.stdio import (
    StdioServerParameters,
    _create_platform_compatible_process,
    _terminate_process_tree,
    stdio_client,
)
from mcp.os.win32.utilities import FallbackProcess
from mcp.shared.exceptions import McpError
from mcp.shared.message import SessionMessage
from mcp.types import CONNECTION_CLOSED, JSONRPCMessage, JSONRPCRequest, JSONRPCResponse

from ..shared.test_win32_utils import escape_path_for_python

# Timeout for cleanup of processes that ignore SIGTERM
# This timeout ensures the test fails quickly if the cleanup logic doesn't have
# proper fallback mechanisms (SIGINT/SIGKILL) for processes that ignore SIGTERM
SIGTERM_IGNORING_PROCESS_TIMEOUT = 5.0

tee = shutil.which("tee")


@pytest.mark.anyio
@pytest.mark.skipif(tee is None, reason="could not find tee command")
async def test_stdio_context_manager_exiting():
    assert tee is not None
    async with stdio_client(StdioServerParameters(command=tee)) as (_, _):
        pass


@pytest.mark.anyio
@pytest.mark.skipif(tee is None, reason="could not find tee command")
async def test_stdio_client():
    assert tee is not None
    server_parameters = StdioServerParameters(command=tee)

    async with stdio_client(server_parameters) as (read_stream, write_stream):
        # Test sending and receiving messages
        messages = [
            JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
            JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
        ]

        async with write_stream:
            for message in messages:
                session_message = SessionMessage(message)
                await write_stream.send(session_message)

        read_messages: list[JSONRPCMessage] = []
        async with read_stream:
            async for message in read_stream:
                if isinstance(message, Exception):  # pragma: no cover
                    raise message

                read_messages.append(message.message)
                if len(read_messages) == 2:
                    break

        assert len(read_messages) == 2
        assert read_messages[0] == JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
        assert read_messages[1] == JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={}))


@pytest.mark.anyio
async def test_stdio_client_bad_path():
    """Check that the connection doesn't hang if process errors."""
    server_params = StdioServerParameters(command=sys.executable, args=["-c", "non-existent-file.py"])
    async with stdio_client(server_params) as (read_stream, write_stream):
        async with ClientSession(read_stream, write_stream) as session:
            # The session should raise an error when the connection closes
            with pytest.raises(McpError) as exc_info:
                await session.initialize()

            # Check that we got a connection closed error
            assert exc_info.value.error.code == CONNECTION_CLOSED
            assert "Connection closed" in exc_info.value.error.message


@pytest.mark.anyio
async def test_stdio_client_nonexistent_command():
    """Test that stdio_client raises an error for non-existent commands."""
    # Create a server with a non-existent command
    server_params = StdioServerParameters(
        command="/path/to/nonexistent/command",
        args=["--help"],
    )

    # Should raise an error when trying to start the process
    with pytest.raises(OSError) as exc_info:
        async with stdio_client(server_params) as (_, _):
            pass  # pragma: no cover

    # The error should indicate the command was not found (ENOENT: No such file or directory)
    assert exc_info.value.errno == errno.ENOENT


@pytest.mark.anyio
async def test_stdio_client_universal_cleanup():
    """
    Test that stdio_client completes cleanup within reasonable time
    even when connected to processes that exit slowly.
    """

    # Use a Python script that simulates a long-running process
    # This ensures consistent behavior across platforms
    long_running_script = textwrap.dedent(
        """
        import time
        import sys

        # Simulate a long-running process
        for i in range(100):
            time.sleep(0.1)
            # Flush to ensure output is visible
            sys.stdout.flush()
            sys.stderr.flush()
        """
    )

    server_params = StdioServerParameters(
        command=sys.executable,
        args=["-c", long_running_script],
    )

    start_time = time.time()

    with anyio.move_on_after(8.0) as cancel_scope:
        async with stdio_client(server_params) as (_, _):
            # Immediately exit - this triggers cleanup while process is still running
            pass

        end_time = time.time()
        elapsed = end_time - start_time

        # On Windows: 2s (stdin wait) + 2s (terminate wait) + overhead = ~5s expected
        assert elapsed < 6.0, (
            f"stdio_client cleanup took {elapsed:.1f} seconds, expected < 6.0 seconds. "
            f"This suggests the timeout mechanism may not be working properly."
        )

    # Check if we timed out
    if cancel_scope.cancelled_caught:  # pragma: no cover
        pytest.fail(
            "stdio_client cleanup timed out after 8.0 seconds. "
            "This indicates the cleanup mechanism is hanging and needs fixing."
        )


@pytest.mark.anyio
@pytest.mark.skipif(sys.platform == "win32", reason="Windows signal handling is different")
async def test_stdio_client_sigint_only_process():  # pragma: no cover
    """
    Test cleanup with a process that ignores SIGTERM but responds to SIGINT.
    """
    # Create a Python script that ignores SIGTERM but handles SIGINT
    script_content = textwrap.dedent(
        """
        import signal
        import sys
        import time

        # Ignore SIGTERM (what process.terminate() sends)
        signal.signal(signal.SIGTERM, signal.SIG_IGN)

        # Handle SIGINT (Ctrl+C signal) by exiting cleanly
        def sigint_handler(signum, frame):
            sys.exit(0)

        signal.signal(signal.SIGINT, sigint_handler)

        # Keep running until SIGINT received
        while True:
            time.sleep(0.1)
        """
    )

    server_params = StdioServerParameters(
        command=sys.executable,
        args=["-c", script_content],
    )

    start_time = time.time()

    try:
        # Use anyio timeout to prevent test from hanging forever
        with anyio.move_on_after(5.0) as cancel_scope:
            async with stdio_client(server_params) as (_, _):
                # Let the process start and begin ignoring SIGTERM
                await anyio.sleep(0.5)
                # Exit context triggers cleanup - this should not hang
                pass

        if cancel_scope.cancelled_caught:  # pragma: no cover
            raise TimeoutError("Test timed out")

        end_time = time.time()
        elapsed = end_time - start_time

        # Should complete quickly even with SIGTERM-ignoring process
        # This will fail if cleanup only uses process.terminate() without fallback
        assert elapsed < SIGTERM_IGNORING_PROCESS_TIMEOUT, (
            f"stdio_client cleanup took {elapsed:.1f} seconds with SIGTERM-ignoring process. "
            f"Expected < {SIGTERM_IGNORING_PROCESS_TIMEOUT} seconds. "
            "This suggests the cleanup needs SIGINT/SIGKILL fallback."
        )
    except (TimeoutError, Exception) as e:  # pragma: no cover
        if isinstance(e, TimeoutError) or "timed out" in str(e):
            pytest.fail(
                f"stdio_client cleanup timed out after {SIGTERM_IGNORING_PROCESS_TIMEOUT} seconds "
                "with SIGTERM-ignoring process. "
                "This confirms the cleanup needs SIGINT/SIGKILL fallback for processes that ignore SIGTERM."
            )
        else:
            raise


async def _wait_for_first_write(path: str) -> None:
    """Poll until the file at *path* exists and has grown beyond its initial empty state.

    The marker files below are created empty before the writer is spawned, so any
    growth proves the writing process booted and reached its write loop. Polling
    replaces fixed startup sleeps, which flake on loaded machines where interpreter
    startup can exceed any fixed window. Bounded so a writer that never starts
    fails the test instead of hanging it.
    """
    with anyio.fail_after(15):
        while not os.path.exists(path) or os.path.getsize(path) == 0:
            await anyio.sleep(0.05)


async def _wait_for_writes_to_stop(path: str) -> None:
    """Poll until the file at *path* stops growing.

    Returns once the size is unchanged across three successive 0.3 second gaps
    (each three times the writers' 0.1 second write interval), so a writer that
    is merely starved of CPU for a single gap is not mistaken for a terminated
    one. Any observed growth resets the consecutive-stable counter. The sentinel
    forces at least one non-stable iteration before counting starts. If the file
    never stops growing, the timeout fails the test: a writer that survives
    _terminate_process_tree is a genuine cleanup failure that must not be masked.
    """
    last_size = -1
    stable_pairs = 0
    with anyio.fail_after(15):
        while True:
            current_size = os.path.getsize(path)
            if current_size == last_size:
                stable_pairs += 1
            else:
                stable_pairs = 0
                last_size = current_size
            if stable_pairs == 3:
                return
            await anyio.sleep(0.3)


async def _dispose_process(proc: Process | FallbackProcess) -> None:
    """Reap a dead process and close its pipe streams inside the test that spawned it.

    Without this, the subprocess transports stay referenced by the per-test event
    loop, become garbage only after that loop closes, and their GC-time
    ResourceWarnings fire during a later test on the same worker (on Windows
    proactor the warning can itself die in __repr__ on a closed pipe). An in-test
    gc.collect() cannot catch that, so the process is reaped and closed
    deterministically here. Draining stdout to EOF guarantees the event loop has
    observed the pipe closure (anyio's reader aclose alone does not close the
    underlying transport), which lets asyncio close the subprocess transport
    before the test returns.

    Precondition: the WHOLE process tree must already be confirmed dead. wait()
    tolerates an already-exited process and returns promptly, but on the Windows
    fallback path it runs popen.wait in a thread and is effectively uncancellable,
    and stdout only reaches EOF once every tree member that inherited the pipe
    handle is gone. The timeout fails the test rather than hanging it if that
    precondition is ever violated.
    """
    with anyio.fail_after(15):
        await proc.wait()
        assert proc.stdin is not None
        await proc.stdin.aclose()
        assert proc.stdout is not None
        while True:
            try:
                await proc.stdout.receive()
            except anyio.EndOfStream:
                break
        await proc.stdout.aclose()


class TestChildProcessCleanup:
    """
    Tests for child process cleanup functionality using _terminate_process_tree.

    These tests verify that child processes are properly terminated when the parent
    is killed, addressing the issue where processes like npx spawn child processes
    that need to be cleaned up. The tests cover various process tree scenarios:

    - Basic parent-child relationship (single child process)
    - Multi-level process trees (parent → child → grandchild)
    - Race conditions where parent exits during cleanup

    Note on Windows ResourceWarning:
    On Windows, we may see ResourceWarning about subprocess still running. This is
    expected behavior due to how Windows process termination works:
    - anyio's process.terminate() calls Windows TerminateProcess() API
    - TerminateProcess() immediately kills the process without allowing cleanup
    - subprocess.Popen objects in the killed process can't run their cleanup code
    - Python detects this during garbage collection and issues a ResourceWarning

    This warning does NOT indicate a process leak - the processes are properly
    terminated. It only means the Popen objects couldn't clean up gracefully.
    This is a fundamental difference between Windows and Unix process termination.
    """

    @pytest.mark.anyio
    @pytest.mark.filterwarnings("ignore::ResourceWarning" if sys.platform == "win32" else "default")
    async def test_basic_child_process_cleanup(self):
        """
        Test basic parent-child process cleanup.
        Parent spawns a single child process that writes continuously to a file.
        """
        # Create a marker file for the child process to write to
        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
            marker_file = f.name

        # Also create a file to verify parent started
        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
            parent_marker = f.name

        # Parent script that spawns a child process
        parent_script = textwrap.dedent(
            f"""
            import subprocess
            import sys
            import time
            import os

            # Mark that parent started
            with open({escape_path_for_python(parent_marker)}, 'w') as f:
                f.write('parent started\\n')

            # Child script that writes continuously
            child_script = f'''
            import time
            with open({escape_path_for_python(marker_file)}, 'a') as f:
                while True:
                    f.write(f"{time.time()}")
                    f.flush()
                    time.sleep(0.1)
            '''

            # Start the child process
            child = subprocess.Popen([sys.executable, '-c', child_script])

            # Parent just sleeps
            while True:
                time.sleep(0.1)
            """
        )

        # Start the parent process
        proc = await _create_platform_compatible_process(sys.executable, ["-c", parent_script])
        tree_killed = False

        try:
            # Wait for the parent to start and the child to reach its write loop
            await _wait_for_first_write(parent_marker)
            assert os.path.getsize(parent_marker) > 0, "Parent process didn't start"

            await _wait_for_first_write(marker_file)
            assert os.path.getsize(marker_file) > 0, "Child process should be writing"

            # Terminate using our function
            await _terminate_process_tree(proc)
            tree_killed = True

            # Verify the child stopped writing; a survivor times out and fails the test
            await _wait_for_writes_to_stop(marker_file)

            # Tree is dead: reap and close the process so nothing leaks into later tests
            await _dispose_process(proc)
        finally:
            if not tree_killed:  # pragma: no cover - cleanup only reached when the test failed mid-flight
                await _terminate_process_tree(proc)
                await _dispose_process(proc)
            # Clean up files
            for f in [marker_file, parent_marker]:
                try:
                    os.unlink(f)
                except OSError:  # pragma: no cover
                    pass
            # Collect subprocess transports now, while this test's warning filters
            # are active, so GC-time ResourceWarnings cannot hit a later test
            gc.collect()

    @pytest.mark.anyio
    @pytest.mark.filterwarnings("ignore::ResourceWarning" if sys.platform == "win32" else "default")
    async def test_nested_process_tree(self):
        """
        Test nested process tree cleanup (parent → child → grandchild).
        Each level writes to a different file to verify all processes are terminated.
        """
        # Create temporary files for each process level
        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f1:
            parent_file = f1.name
        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f2:
            child_file = f2.name
        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f3:
            grandchild_file = f3.name

        # Simple nested process tree test
        # We create parent -> child -> grandchild, each writing to a file
        parent_script = textwrap.dedent(
            f"""
            import subprocess
            import sys
            import time
            import os

            # Child will spawn grandchild and write to child file
            child_script = f'''import subprocess
            import sys
            import time

            # Grandchild just writes to file
            grandchild_script = \"\"\"import time
            with open({escape_path_for_python(grandchild_file)}, 'a') as f:
                while True:
                    f.write(f"gc {{time.time()}}")
                    f.flush()
                    time.sleep(0.1)\"\"\"

            # Spawn grandchild
            subprocess.Popen([sys.executable, '-c', grandchild_script])

            # Child writes to its file
            with open({escape_path_for_python(child_file)}, 'a') as f:
                while True:
                    f.write(f"c {time.time()}")
                    f.flush()
                    time.sleep(0.1)'''

            # Spawn child process
            subprocess.Popen([sys.executable, '-c', child_script])

            # Parent writes to its file
            with open({escape_path_for_python(parent_file)}, 'a') as f:
                while True:
                    f.write(f"p {time.time()}")
                    f.flush()
                    time.sleep(0.1)
            """
        )

        # Start the parent process
        proc = await _create_platform_compatible_process(sys.executable, ["-c", parent_script])
        tree_killed = False

        try:
            # Wait for every level of the tree to reach its write loop
            for file_path, name in [(parent_file, "parent"), (child_file, "child"), (grandchild_file, "grandchild")]:
                await _wait_for_first_write(file_path)
                assert os.path.getsize(file_path) > 0, f"{name} process should be writing"

            # Terminate the whole tree
            await _terminate_process_tree(proc)
            tree_killed = True

            # Verify every level stopped writing; a survivor times out and fails the test
            for file_path in (parent_file, child_file, grandchild_file):
                await _wait_for_writes_to_stop(file_path)

            # Tree is dead: reap and close the process so nothing leaks into later tests
            await _dispose_process(proc)
        finally:
            if not tree_killed:  # pragma: no cover - cleanup only reached when the test failed mid-flight
                await _terminate_process_tree(proc)
                await _dispose_process(proc)
            # Clean up all marker files
            for f in [parent_file, child_file, grandchild_file]:
                try:
                    os.unlink(f)
                except OSError:  # pragma: no cover
                    pass
            # Collect subprocess transports now, while this test's warning filters
            # are active, so GC-time ResourceWarnings cannot hit a later test
            gc.collect()

    @pytest.mark.anyio
    @pytest.mark.filterwarnings("ignore::ResourceWarning" if sys.platform == "win32" else "default")
    async def test_early_parent_exit(self):
        """
        Test cleanup when parent exits during termination sequence.
        Tests the race condition where parent might die during our termination
        sequence but we can still clean up the children via the process group.
        """
        # Create a temporary file for the child
        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
            marker_file = f.name

        # Parent that spawns child and waits briefly
        parent_script = textwrap.dedent(
            f"""
            import subprocess
            import sys
            import time
            import signal

            # Child that continues running
            child_script = f'''import time
            with open({escape_path_for_python(marker_file)}, 'a') as f:
                while True:
                    f.write(f"child {time.time()}")
                    f.flush()
                    time.sleep(0.1)'''

            # Start child in same process group
            subprocess.Popen([sys.executable, '-c', child_script])

            # Parent waits a bit then exits on SIGTERM
            def handle_term(sig, frame):
                sys.exit(0)

            signal.signal(signal.SIGTERM, handle_term)

            # Wait
            while True:
                time.sleep(0.1)
            """
        )

        # Start the parent process
        proc = await _create_platform_compatible_process(sys.executable, ["-c", parent_script])
        tree_killed = False

        try:
            # Wait for the child to reach its write loop
            await _wait_for_first_write(marker_file)
            assert os.path.getsize(marker_file) > 0, "Child should be writing"

            # Terminate - this will kill the process group even if parent exits first
            await _terminate_process_tree(proc)
            tree_killed = True

            # Verify the child stopped writing; a survivor times out and fails the test
            await _wait_for_writes_to_stop(marker_file)

            # Tree is dead: reap and close the process so nothing leaks into later tests
            await _dispose_process(proc)
        finally:
            if not tree_killed:  # pragma: no cover - cleanup only reached when the test failed mid-flight
                await _terminate_process_tree(proc)
                await _dispose_process(proc)
            # Clean up marker file
            try:
                os.unlink(marker_file)
            except OSError:  # pragma: no cover
                pass
            # Collect subprocess transports now, while this test's warning filters
            # are active, so GC-time ResourceWarnings cannot hit a later test
            gc.collect()


@pytest.mark.anyio
async def test_stdio_client_graceful_stdin_exit():
    """
    Test that a process exits gracefully when stdin is closed,
    without needing SIGTERM or SIGKILL.
    """
    # Create a Python script that exits when stdin is closed
    script_content = textwrap.dedent(
        """
        import sys

        # Read from stdin until it's closed
        try:
            while True:
                line = sys.stdin.readline()
                if not line:  # EOF/stdin closed
                    break
        except:
            pass

        # Exit gracefully
        sys.exit(0)
        """
    )

    server_params = StdioServerParameters(
        command=sys.executable,
        args=["-c", script_content],
    )

    start_time = time.time()

    # Use anyio timeout to prevent test from hanging forever
    with anyio.move_on_after(5.0) as cancel_scope:
        async with stdio_client(server_params) as (_, _):
            # Let the process start and begin reading stdin
            await anyio.sleep(0.2)
            # Exit context triggers cleanup - process should exit from stdin closure
            pass

    if cancel_scope.cancelled_caught:
        pytest.fail(
            "stdio_client cleanup timed out after 5.0 seconds. "
            "Process should have exited gracefully when stdin was closed."
        )  # pragma: no cover

    end_time = time.time()
    elapsed = end_time - start_time

    # Should complete quickly with just stdin closure (no signals needed)
    assert elapsed < 3.0, (
        f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-aware process. "
        f"Expected < 3.0 seconds since process should exit on stdin closure."
    )


@pytest.mark.anyio
async def test_stdio_client_stdin_close_ignored():
    """
    Test that when a process ignores stdin closure, the shutdown sequence
    properly escalates to SIGTERM.
    """
    # Create a Python script that ignores stdin closure but responds to SIGTERM
    script_content = textwrap.dedent(
        """
        import signal
        import sys
        import time

        # Set up SIGTERM handler to exit cleanly
        def sigterm_handler(signum, frame):
            sys.exit(0)

        signal.signal(signal.SIGTERM, sigterm_handler)

        # Close stdin immediately to simulate ignoring it
        sys.stdin.close()

        # Keep running until SIGTERM
        while True:
            time.sleep(0.1)
        """
    )

    server_params = StdioServerParameters(
        command=sys.executable,
        args=["-c", script_content],
    )

    start_time = time.time()

    # Use anyio timeout to prevent test from hanging forever
    with anyio.move_on_after(7.0) as cancel_scope:
        async with stdio_client(server_params) as (_, _):
            # Let the process start
            await anyio.sleep(0.2)
            # Exit context triggers cleanup
            pass

    if cancel_scope.cancelled_caught:
        pytest.fail(
            "stdio_client cleanup timed out after 7.0 seconds. "
            "Process should have been terminated via SIGTERM escalation."
        )  # pragma: no cover

    end_time = time.time()
    elapsed = end_time - start_time

    # Should take ~2 seconds (stdin close timeout) before SIGTERM is sent
    # Total time should be between 2-4 seconds
    assert 1.5 < elapsed < 4.5, (
        f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-ignoring process. "
        f"Expected between 2-4 seconds (2s stdin timeout + termination time)."
    )