"""Step definitions for features/tdd_mcp_adapter_rlock_concurrency.feature.

Tests that MCPToolAdapter releases its RLock before making transport calls,
allowing concurrent operations to proceed without being blocked.

Bug #10512: MCPToolAdapter holds RLock during entire transport call,
blocking concurrent operations.
"""

from __future__ import annotations

import threading
import time
from typing import Any

from behave import given, then, when
from behave.runner import Context

from cleveragents.mcp.adapter import (
    MCPServerConfig,
    MCPToolAdapter,
    MCPToolResult,
)
from features.mocks.mock_mcp_transport import MockMCPTransport


def _mock_tool(
    name: str, desc: str = "", schema: dict[str, Any] | None = None
) -> dict[str, Any]:
    return {
        "name": name,
        "description": desc or f"Mock tool {name}",
        "inputSchema": schema or {},
    }


class SlowMCPTransport(MockMCPTransport):
    """Transport that introduces a configurable delay during call()."""

    def __init__(
        self,
        delay: float,
        tools: list[dict[str, Any]] | None = None,
    ) -> None:
        super().__init__(tools=tools)
        self._delay = delay

    def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        # Use the real time.sleep (not the capped test version) to ensure
        # the delay is meaningful for concurrency testing.
        _real_sleep = getattr(time, "_original_sleep", time.sleep)
        _real_sleep(self._delay)
        return super().call(method, params)


class ConcurrencyTrackingTransport(MockMCPTransport):
    """Transport that tracks the maximum number of concurrent in-flight calls.

    Uses a threading.Lock-protected counter to record how many calls are
    executing simultaneously.  The ``max_concurrent`` attribute reflects the
    peak concurrency observed across all calls.

    If the adapter holds its RLock during transport calls, only one call can
    be in-flight at a time (max_concurrent == 1).  If the lock is released
    before the transport call, multiple calls can overlap (max_concurrent > 1).
    """

    def __init__(
        self,
        delay: float,
        tools: list[dict[str, Any]] | None = None,
    ) -> None:
        super().__init__(tools=tools)
        self._delay = delay
        self._counter_lock = threading.Lock()
        self._current_concurrent: int = 0
        self.max_concurrent: int = 0

    def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        with self._counter_lock:
            self._current_concurrent += 1
            if self._current_concurrent > self.max_concurrent:
                self.max_concurrent = self._current_concurrent

        # Use the real time.sleep (not the capped test version) to ensure
        # the delay is meaningful for concurrency testing.
        _real_sleep = getattr(time, "_original_sleep", time.sleep)
        _real_sleep(self._delay)

        with self._counter_lock:
            self._current_concurrent -= 1

        return super().call(method, params)


class LockCheckingInvokeTransport(MockMCPTransport):
    """Transport that checks whether the adapter's RLock is held during call()."""

    def __init__(
        self,
        adapter_lock: threading.RLock,
        tools: list[dict[str, Any]] | None = None,
    ) -> None:
        super().__init__(tools=tools)
        self._adapter_lock = adapter_lock
        self.lock_was_held_during_call = False

    def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        if method == "tools/call":
            # Try to acquire the lock from a different thread.
            # If the lock is held by the calling thread, a different thread
            # will fail to acquire it (non-blocking).
            acquired_event = threading.Event()
            lock_held = [True]

            def _try_acquire() -> None:
                got_it = self._adapter_lock.acquire(blocking=False)
                if got_it:
                    lock_held[0] = False
                    self._adapter_lock.release()
                acquired_event.set()

            t = threading.Thread(target=_try_acquire, daemon=True)
            t.start()
            t.join(timeout=1.0)
            acquired_event.wait(timeout=1.0)
            self.lock_was_held_during_call = lock_held[0]

        return super().call(method, params)


class LockCheckingDiscoveryTransport(MockMCPTransport):
    """Transport that checks whether the adapter's RLock is held during tools/list."""

    def __init__(
        self,
        adapter_lock: threading.RLock,
        tools: list[dict[str, Any]] | None = None,
    ) -> None:
        super().__init__(tools=tools)
        self._adapter_lock = adapter_lock
        self.lock_was_held_during_call = False

    def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        if method == "tools/list":
            acquired_event = threading.Event()
            lock_held = [True]

            def _try_acquire() -> None:
                got_it = self._adapter_lock.acquire(blocking=False)
                if got_it:
                    lock_held[0] = False
                    self._adapter_lock.release()
                acquired_event.set()

            t = threading.Thread(target=_try_acquire, daemon=True)
            t.start()
            t.join(timeout=1.0)
            acquired_event.wait(timeout=1.0)
            self.lock_was_held_during_call = lock_held[0]

        return super().call(method, params)


# ---------------------------------------------------------------
# Given Steps
# ---------------------------------------------------------------


@given(
    'a connected MCP adapter with a slow transport tool "{tool_name}" taking {delay:f} seconds'
)
def step_slow_transport_tool(context: Context, tool_name: str, delay: float) -> None:
    tools = [_mock_tool(tool_name)]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = SlowMCPTransport(delay=delay, tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given(
    'a connected MCP adapter with a concurrency-tracking transport tool "{tool_name}" taking {delay:f} seconds'
)
def step_concurrency_tracking_transport_tool(
    context: Context, tool_name: str, delay: float
) -> None:
    tools = [_mock_tool(tool_name)]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = ConcurrencyTrackingTransport(delay=delay, tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given(
    "a connected MCP adapter with a slow discovery transport taking {delay:f} seconds"
)
def step_slow_discovery_transport(context: Context, delay: float) -> None:
    tools = [_mock_tool("tool_0"), _mock_tool("tool_1")]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = SlowMCPTransport(delay=delay, tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()


@given(
    "a connected MCP adapter with a concurrency-tracking discovery transport taking {delay:f} seconds"
)
def step_concurrency_tracking_discovery_transport(
    context: Context, delay: float
) -> None:
    tools = [_mock_tool("tool_0"), _mock_tool("tool_1")]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = ConcurrencyTrackingTransport(delay=delay, tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()


@given('a connected MCP adapter with a lock-checking transport tool "{tool_name}"')
def step_lock_checking_invoke_transport(context: Context, tool_name: str) -> None:
    tools = [_mock_tool(tool_name)]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    # Create adapter first to get its lock, then set up transport
    transport_placeholder = MockMCPTransport(tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=transport_placeholder)
    # Now create the lock-checking transport with the adapter's actual lock
    lock_transport = LockCheckingInvokeTransport(
        adapter_lock=context.mcp_adapter._lock,
        tools=tools,
    )
    # Replace the transport
    context.mcp_adapter._transport = lock_transport
    context.mcp_transport = lock_transport
    context.mcp_adapter._connected = True
    context.mcp_adapter._capabilities = {"tools": True}
    context.mcp_adapter.discover_tools()


@given("a connected MCP adapter with a lock-checking discovery transport")
def step_lock_checking_discovery_transport(context: Context) -> None:
    tools = [_mock_tool("tool_0"), _mock_tool("tool_1")]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    transport_placeholder = MockMCPTransport(tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=transport_placeholder)
    lock_transport = LockCheckingDiscoveryTransport(
        adapter_lock=context.mcp_adapter._lock,
        tools=tools,
    )
    context.mcp_adapter._transport = lock_transport
    context.mcp_transport = lock_transport
    context.mcp_adapter._connected = True
    context.mcp_adapter._capabilities = {"tools": True}


# ---------------------------------------------------------------
# When Steps
# ---------------------------------------------------------------


@when('I invoke "{tool_name}" concurrently from {count:d} threads')
def step_concurrent_invoke(context: Context, tool_name: str, count: int) -> None:
    # Use threading.Thread directly instead of ThreadPoolExecutor to avoid
    # deadlocks in forked worker processes (behave-parallel uses fork).
    # ThreadPoolExecutor holds internal locks that can be in a locked state
    # after fork(), causing child processes to deadlock on acquisition.
    results: list[MCPToolResult | None] = [None] * count
    errors: list[BaseException | None] = [None] * count
    errors_lock = threading.Lock()

    def _invoke(idx: int) -> None:
        try:
            results[idx] = context.mcp_adapter.invoke(tool_name, {})
        except BaseException as exc:
            with errors_lock:
                errors[idx] = exc

    threads = [threading.Thread(target=_invoke, args=(i,)) for i in range(count)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    # Re-raise the first error encountered, if any
    for err in errors:
        if err is not None:
            raise err

    context.concurrent_invoke_results = [r for r in results if r is not None]


@when("I call discover_tools concurrently from {count:d} threads")
def step_concurrent_discover(context: Context, count: int) -> None:
    # Use threading.Thread directly instead of ThreadPoolExecutor to avoid
    # deadlocks in forked worker processes (behave-parallel uses fork).
    # ThreadPoolExecutor holds internal locks that can be in a locked state
    # after fork(), causing child processes to deadlock on acquisition.
    results: list[Any] = []
    errors: list[Exception] = []
    collect_lock = threading.Lock()

    def _discover() -> None:
        try:
            result = context.mcp_adapter.discover_tools()
            with collect_lock:
                results.append(result)
        except Exception as exc:
            with collect_lock:
                errors.append(exc)

    threads = [threading.Thread(target=_discover) for _ in range(count)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    context.concurrent_discover_results = results
    context.concurrent_discover_errors = errors


# ---------------------------------------------------------------
# Then Steps
# ---------------------------------------------------------------


@then("all {count:d} concurrent invocations should succeed")
def step_all_invocations_succeed(context: Context, count: int) -> None:
    results = context.concurrent_invoke_results
    assert len(results) == count, f"Expected {count} results, got {len(results)}"
    for i, result in enumerate(results):
        assert result.success, f"Invocation {i} failed: {result.error}"


@then(
    "at least {min_concurrent:d} invocations should have been in-flight simultaneously"
)
def step_at_least_n_invocations_concurrent(
    context: Context, min_concurrent: int
) -> None:
    transport = context.mcp_transport
    actual = transport.max_concurrent
    assert actual >= min_concurrent, (
        f"Expected at least {min_concurrent} concurrent in-flight invocations, "
        f"but observed max_concurrent={actual}. "
        f"The adapter's RLock appears to be held during transport calls, "
        f"serializing concurrent operations."
    )


@then("all {count:d} concurrent discoveries should succeed")
def step_all_discoveries_succeed(context: Context, count: int) -> None:
    results = context.concurrent_discover_results
    errors = context.concurrent_discover_errors
    assert not errors, f"Expected no errors, got {len(errors)}: {errors}"
    assert len(results) == count, f"Expected {count} results, got {len(results)}"


@then(
    "at least {min_concurrent:d} discoveries should have been in-flight simultaneously"
)
def step_at_least_n_discoveries_concurrent(
    context: Context, min_concurrent: int
) -> None:
    transport = context.mcp_transport
    actual = transport.max_concurrent
    assert actual >= min_concurrent, (
        f"Expected at least {min_concurrent} concurrent in-flight discoveries, "
        f"but observed max_concurrent={actual}. "
        f"The adapter's RLock appears to be held during transport calls, "
        f"serializing concurrent operations."
    )


@then("the lock should not have been held during the transport call")
def step_lock_not_held_during_invoke(context: Context) -> None:
    transport = context.mcp_transport
    assert not transport.lock_was_held_during_call, (
        "The adapter's RLock was held during the transport.call() invocation. "
        "The lock should be released before making the transport call."
    )


@then("the lock should not have been held during the discovery transport call")
def step_lock_not_held_during_discovery(context: Context) -> None:
    transport = context.mcp_transport
    assert not transport.lock_was_held_during_call, (
        "The adapter's RLock was held during the transport.call('tools/list') invocation. "
        "The lock should be released before making the transport call."
    )
