# -*- coding: utf-8 -*-
"""Location: ./tests/unit/mcpgateway/test_wrapper.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti + contributors

Tests for the MCP *wrapper* module (single file, full coverage).
This suite fakes the "mcp" dependency tree so that no real network or
pydantic models are required and exercises almost every branch inside
*mcpgateway.wrapper*.
"""

# Standard
import asyncio
import contextlib
import errno
import sys
import types
from unittest.mock import AsyncMock

# Third-Party
import pytest

# First-Party
import mcpgateway.wrapper as wrapper


# Ensure shutdown flag is clear before each test run
def setup_function():
    wrapper._shutdown.clear()


# -------------------
# Utilities
# -------------------
def test_convert_url_variants():
    assert wrapper.convert_url("http://x/servers/uuid") == "http://x/servers/uuid/mcp/"
    assert wrapper.convert_url("http://x/servers/uuid/") == "http://x/servers/uuid//mcp/"
    assert wrapper.convert_url("http://x/servers/uuid/mcp") == "http://x/servers/uuid/mcp/"
    assert wrapper.convert_url("http://x/servers/uuid/mcp/") == "http://x/servers/uuid/mcp/"
    assert wrapper.convert_url("http://x/servers/uuid/sse") == "http://x/servers/uuid/mcp/"


def test_make_error_defaults_and_data():
    err = wrapper.make_error("oops")
    assert err["error"]["message"] == "oops"
    assert err["error"]["code"] == wrapper.JSONRPC_INTERNAL_ERROR
    err2 = wrapper.make_error("bad", code=-32099, data={"x": 1})
    assert err2["error"]["data"] == {"x": 1}
    assert err2["error"]["code"] == -32099


def test_setup_logging_on_and_off():
    wrapper.setup_logging("DEBUG")
    assert wrapper.logger.disabled is False
    wrapper.logger.debug("hello debug")
    wrapper.setup_logging("OFF")
    assert wrapper.logger.disabled is True


def test_setup_logging_none_disables():
    wrapper.setup_logging("DEBUG")
    assert wrapper.logger.disabled is False
    wrapper.setup_logging(None)
    assert wrapper.logger.disabled is True


def test_shutting_down_and_mark_shutdown():
    wrapper._shutdown.clear()
    assert not wrapper.shutting_down()
    wrapper._mark_shutdown()
    assert wrapper.shutting_down()
    # Calling it again should hit the "already set" branch.
    wrapper._mark_shutdown()
    assert wrapper.shutting_down()
    wrapper._shutdown.clear()
    assert not wrapper.shutting_down()


def test_send_to_stdout_json_and_str(monkeypatch):
    captured = []

    def fake_write(s):
        captured.append(s)
        return len(s)

    # Mock both buffer and text write to catch whichever is used
    monkeypatch.setattr(sys.stdout, "write", fake_write)
    if hasattr(sys.stdout, "buffer"):
        monkeypatch.setattr(sys.stdout.buffer, "write", fake_write)

    monkeypatch.setattr(sys.stdout, "flush", lambda: None)
    if hasattr(sys.stdout, "buffer"):
        monkeypatch.setattr(sys.stdout.buffer, "flush", lambda: None)

    wrapper.send_to_stdout({"a": 1})
    wrapper.send_to_stdout("plain text")

    # decode captured bytes to str for assertion simplicity
    decoded_captured = []
    for s in captured:
        if isinstance(s, bytes):
            decoded_captured.append(s.decode("utf-8"))
        else:
            decoded_captured.append(s)

    assert any('"a":1' in s or '"a": 1' in s for s in decoded_captured)
    assert any("plain text" in s for s in decoded_captured)


def test_send_to_stdout_oserror(monkeypatch):
    wrapper._shutdown.clear()

    def bad_write(_):
        raise OSError(errno.EPIPE, "broken pipe")

    monkeypatch.setattr(sys.stdout, "write", bad_write)
    if hasattr(sys.stdout, "buffer"):
        monkeypatch.setattr(sys.stdout.buffer, "write", bad_write)

    monkeypatch.setattr(sys.stdout, "flush", lambda: None)
    if hasattr(sys.stdout, "buffer"):
        monkeypatch.setattr(sys.stdout.buffer, "flush", lambda: None)

    wrapper.send_to_stdout({"x": 1})
    assert wrapper.shutting_down()


def test_send_to_stdout_oserror_other_errno(monkeypatch):
    wrapper._shutdown.clear()

    def bad_write(_):
        raise OSError(errno.EIO, "io error")

    monkeypatch.setattr(sys.stdout, "write", bad_write)
    if hasattr(sys.stdout, "buffer"):
        monkeypatch.setattr(sys.stdout.buffer, "write", bad_write)

    monkeypatch.setattr(sys.stdout, "flush", lambda: None)
    if hasattr(sys.stdout, "buffer"):
        monkeypatch.setattr(sys.stdout.buffer, "flush", lambda: None)

    wrapper.send_to_stdout({"x": 2})
    assert wrapper.shutting_down()


def test_send_to_stdout_no_buffer_fallback(monkeypatch):
    wrapper._shutdown.clear()
    captured = []

    class DummyStdout:
        def write(self, text):
            captured.append(text)
            return len(text)

        def flush(self):
            return None

    def raises(_):
        raise TypeError("boom")

    monkeypatch.setattr(wrapper, "orjson", types.SimpleNamespace(dumps=raises))
    monkeypatch.setattr(sys, "stdout", DummyStdout())

    wrapper.send_to_stdout("plain text")
    wrapper.send_to_stdout(b"bytes")

    assert any("plain text" in s for s in captured)
    assert any("bytes" in s for s in captured)


# -------------------
# Async stream parsers
# -------------------
@pytest.mark.asyncio
async def test_ndjson_lines_basic_and_tail():
    wrapper._shutdown.clear()

    async def fake_iter_bytes():
        # basic multi-line + a final line without newline to test tail handling
        yield b'{"a":1}\n{"b":2}\n'
        yield b'{"c":3}'

    resp = types.SimpleNamespace(aiter_bytes=fake_iter_bytes, aiter_lines=None)  # aiter_lines not used in optimized version
    lines = [l async for l in wrapper.ndjson_lines(resp)]
    # lines are bytes now
    assert b'{"a":1}' in lines
    assert b'{"b":2}' in lines
    assert b'{"c":3}' in lines


@pytest.mark.asyncio
async def test_ndjson_lines_skips_empty_chunk_and_blank_lines():
    wrapper._shutdown.clear()

    async def fake_iter_bytes():
        yield b""
        yield b'{"a":1}\n\n{"b":2}\n'

    resp = types.SimpleNamespace(aiter_bytes=fake_iter_bytes)
    lines = [l async for l in wrapper.ndjson_lines(resp)]
    assert lines == [b'{"a":1}', b'{"b":2}']


@pytest.mark.asyncio
async def test_ndjson_lines_breaks_on_shutdown():
    wrapper._shutdown.clear()

    async def fake_iter_bytes():
        yield b'{"a":1}\n'
        wrapper._mark_shutdown()
        yield b'{"b":2}\n'

    resp = types.SimpleNamespace(aiter_bytes=fake_iter_bytes)
    lines = [l async for l in wrapper.ndjson_lines(resp)]
    assert lines == [b'{"a":1}']


@pytest.mark.asyncio
async def test_sse_events_basic_and_tail():
    wrapper._shutdown.clear()

    async def fake_iter_bytes():
        # two events with proper separators, plus a tail-only chunk
        yield b"data: first\n\n"
        yield b"data: second\n\n"
        yield b"data: tailonly\n\n"

    resp = types.SimpleNamespace(aiter_bytes=fake_iter_bytes)
    events = [e async for e in wrapper.sse_events(resp)]
    # events are bytes
    assert b"first" in events
    assert b"second" in events
    assert b"tailonly" in events


@pytest.mark.asyncio
async def test_sse_events_handles_comments_blank_events_and_partial_line():
    wrapper._shutdown.clear()

    async def fake_iter_bytes():
        yield b""
        yield b": comment\n\n"
        yield b"data\n\n"
        yield b"data: tail"

    resp = types.SimpleNamespace(aiter_bytes=fake_iter_bytes)
    events = [e async for e in wrapper.sse_events(resp)]
    assert b"" in events  # from `data` line without value
    assert b"tail" in events


@pytest.mark.asyncio
async def test_sse_events_breaks_on_shutdown():
    wrapper._shutdown.clear()

    async def fake_iter_bytes():
        yield b"data: one\n\n"
        wrapper._mark_shutdown()
        yield b"data: two\n\n"

    resp = types.SimpleNamespace(aiter_bytes=fake_iter_bytes)
    events = [e async for e in wrapper.sse_events(resp)]
    assert events == [b"one"]


@pytest.mark.asyncio
async def test_sse_events_partial_line_without_colon():
    wrapper._shutdown.clear()

    async def fake_iter_bytes():
        # A normal event, then a trailing partial line with no ":" to hit the
        # `field, value = line, b""` branch in tail processing.
        yield b"event: message\ndata: one\n\n"
        yield b"data"

    resp = types.SimpleNamespace(aiter_bytes=fake_iter_bytes)
    events = [e async for e in wrapper.sse_events(resp)]
    assert b"one" in events
    assert b"" in events


@pytest.mark.asyncio
async def test_sse_events_tail_ignores_comment_and_non_data_fields():
    wrapper._shutdown.clear()

    async def fake_iter_bytes_comment_tail():
        yield b"data: one\n\n"
        yield b": this is a comment"

    resp1 = types.SimpleNamespace(aiter_bytes=fake_iter_bytes_comment_tail)
    events1 = [e async for e in wrapper.sse_events(resp1)]
    assert events1 == [b"one"]

    async def fake_iter_bytes_event_tail():
        yield b"data: one\n\n"
        yield b"event: message"

    resp2 = types.SimpleNamespace(aiter_bytes=fake_iter_bytes_event_tail)
    events2 = [e async for e in wrapper.sse_events(resp2)]
    assert events2 == [b"one"]


# -------------------
# Settings dataclass
# -------------------
def test_settings_defaults():
    s = wrapper.Settings("http://x/mcp", "Bearer token", 5, 10, 2, "DEBUG")
    assert s.server_url == "http://x/mcp"
    assert s.auth_header == "Bearer token"
    assert s.concurrency == 2


# -------------------
# parse_args
# -------------------
def test_parse_args_with_env(monkeypatch):
    monkeypatch.setenv("MCP_SERVER_URL", "http://localhost:4444/servers/uuid")
    monkeypatch.setenv("MCP_AUTH", "Bearer 123")
    sys_argv = sys.argv
    sys.argv = ["prog"]
    try:
        s = wrapper.parse_args()
        assert s.server_url.endswith("/mcp/")
        assert s.auth_header == "Bearer 123"
    finally:
        sys.argv = sys_argv


def test_parse_args_missing_url(monkeypatch):
    # no env and no arg should exit with SystemExit
    monkeypatch.delenv("MCP_SERVER_URL", raising=False)
    sys_argv = sys.argv
    sys.argv = ["prog"]
    try:
        with pytest.raises(SystemExit):
            wrapper.parse_args()
    finally:
        sys.argv = sys_argv


# -------------------
# stdin_reader
# -------------------
@pytest.mark.asyncio
async def test_stdin_reader_valid_and_invalid(monkeypatch):
    wrapper._shutdown.clear()
    q = asyncio.Queue()

    # synchronous readline callable used by asyncio.to_thread
    lines = iter([b'{"ok":1}\n', b"{bad json}\n", b"   \n", b""])

    def fake_readline():
        try:
            return next(lines)
        except StopIteration:
            return b""

    # Mock buffer.readline if available, else stdin.readline (but existing test ran on host python which likely has buffer)
    # The wrapper uses sys.stdin.buffer.readline if available.
    if hasattr(sys.stdin, "buffer"):
        monkeypatch.setattr(sys.stdin.buffer, "readline", fake_readline)
    else:
        monkeypatch.setattr(sys.stdin, "readline", fake_readline)

    task = asyncio.create_task(wrapper.stdin_reader(q))

    # Collect three items: valid dict, error dict for invalid json, and None for EOF
    got1 = await asyncio.wait_for(q.get(), timeout=1)
    got2 = await asyncio.wait_for(q.get(), timeout=1)
    got3 = await asyncio.wait_for(q.get(), timeout=1)

    # first should be parsed dict
    assert isinstance(got1, dict) and got1.get("ok") == 1
    # second should be an error object (from make_error)
    assert isinstance(got2, dict) and "error" in got2
    # third should be None (EOF sentinel)
    assert got3 is None

    task.cancel()
    with contextlib.suppress(asyncio.CancelledError):
        await task


@pytest.mark.asyncio
async def test_stdin_reader_no_buffer_path(monkeypatch):
    wrapper._shutdown.clear()
    q = asyncio.Queue()

    lines = iter(["{bad json}\n", ""])

    class DummyStdin:
        def readline(self):
            return next(lines)

    async def fake_to_thread(func, *args, **kwargs):
        return func(*args, **kwargs)

    monkeypatch.setattr(sys, "stdin", DummyStdin())
    monkeypatch.setattr(wrapper.asyncio, "to_thread", fake_to_thread)

    task = asyncio.create_task(wrapper.stdin_reader(q))
    got1 = await asyncio.wait_for(q.get(), timeout=1)
    got2 = await asyncio.wait_for(q.get(), timeout=1)

    assert isinstance(got1, dict) and "error" in got1
    assert got2 is None

    task.cancel()
    with contextlib.suppress(asyncio.CancelledError):
        await task


@pytest.mark.asyncio
async def test_stdin_reader_decode_failure_falls_back_to_str(monkeypatch):
    """Force line.decode() to raise so stdin_reader uses str(line) fallback."""
    wrapper._shutdown.clear()
    q: asyncio.Queue = asyncio.Queue()

    class FakeLine:
        def strip(self):
            return self

        def decode(self, *_a, **_k):
            raise UnicodeError("boom")

        def __len__(self):
            return 1

        def __str__(self):
            return "<FakeLine>"

    lines = iter([FakeLine(), b""])

    def fake_readline():
        return next(lines)

    async def fake_to_thread(func, *args, **kwargs):
        return func(*args, **kwargs)

    if hasattr(sys.stdin, "buffer"):
        monkeypatch.setattr(sys.stdin.buffer, "readline", fake_readline)
    else:
        monkeypatch.setattr(sys.stdin, "readline", fake_readline)

    monkeypatch.setattr(wrapper.asyncio, "to_thread", fake_to_thread)

    task = asyncio.create_task(wrapper.stdin_reader(q))
    err_obj = await asyncio.wait_for(q.get(), timeout=1)
    eof = await asyncio.wait_for(q.get(), timeout=1)

    assert isinstance(err_obj, dict) and err_obj.get("error", {}).get("data") == "<FakeLine>"
    assert eof is None

    task.cancel()
    with contextlib.suppress(asyncio.CancelledError):
        await task


# -------------------
# forward_once (JSON / NDJSON / SSE / HTTP error)
# -------------------
class DummyResp:
    def __init__(self, status=200, ctype="application/json", body=b'{"ok":1}'):
        self.status_code = status
        # use dict-like headers as wrapper expects resp.headers.get(...)
        self._headers = {"Content-Type": ctype}
        self._body = body

    @property
    def headers(self):
        return self._headers

    async def aread(self):
        # return full body for application/json path
        return self._body

    # context manager to be used with `async with client.stream(...) as resp:`
    async def __aenter__(self):
        return self

    async def __aexit__(self, *a):
        return False

    async def aiter_bytes(self):
        # yield the body as a single chunk (ndjson and sse parsers will process it)
        yield self._body


class DummyClient:
    def __init__(self, resp):
        self._resp = resp

    def stream(self, *a, **k):
        # returning the response instance which implements async context manager
        return self._resp


class CapturingClient:
    def __init__(self, resp):
        self._resp = resp
        self.calls = []

    def stream(self, *a, **k):
        self.calls.append(k)
        return self._resp


@pytest.mark.asyncio
async def test_forward_once_json_and_invalid(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    # valid JSON response
    client = DummyClient(DummyResp(200, "application/json", b'{"ok":123}'))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"a": 1})
    assert any(isinstance(o, dict) and o.get("ok") == 123 for o in captured)

    # invalid JSON body (application/json but not JSON)
    client = DummyClient(DummyResp(200, "application/json", b"notjson"))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"b": 2})
    # should have produced an error object for invalid JSON response
    assert any(isinstance(o, dict) and "error" in o for o in captured)


@pytest.mark.asyncio
async def test_forward_once_application_json_empty_body_returns_no_output(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    client = DummyClient(DummyResp(202, "application/json", b""))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"jsonrpc": "2.0", "method": "notifications/initialized"})

    assert captured == []


@pytest.mark.asyncio
async def test_forward_once_ndjson_and_sse_and_http_error(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    # ndjson: multiple JSON lines
    ndj = b'{"x":1}\n{"y":2}\n'
    client = DummyClient(DummyResp(200, "application/x-ndjson", ndj))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"z": 3})
    assert any(isinstance(d, dict) and ("x" in d or "y" in d) for d in captured)

    # sse: streaming events with data: lines containing JSON
    sse_chunk = b'data: {"foo": 42}\n\n'
    client = DummyClient(DummyResp(200, "text/event-stream", sse_chunk))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"w": 4})
    assert any(isinstance(d, dict) and d.get("foo") == 42 for d in captured)

    # http error (non-2xx)
    client = DummyClient(DummyResp(500, "application/json", b""))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"e": 1})
    assert any(isinstance(d, dict) and "error" in d for d in captured)


@pytest.mark.asyncio
async def test_forward_once_form_encoding_and_auth_header(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    client = CapturingClient(DummyResp(200, "application/json", b'{"ok":true}'))
    settings = wrapper.Settings("http://x/mcp", "Bearer token")
    settings.content_type = "application/x-www-form-urlencoded"

    await wrapper.forward_once(client, settings, {"a": "b", "c": 1})
    call = client.calls[0]
    assert call["headers"]["Authorization"] == "Bearer token"
    assert call["headers"]["Content-Type"] == "application/x-www-form-urlencoded"
    assert b"a=b" in call["data"]
    assert b"c=1" in call["data"]
    assert any(isinstance(o, dict) and o.get("ok") is True for o in captured)


@pytest.mark.asyncio
async def test_forward_once_returns_immediately_if_shutting_down(monkeypatch):
    wrapper._shutdown.clear()
    wrapper._mark_shutdown()

    class BoomClient:
        def stream(self, *a, **k):  # pragma: no cover - should not be called
            raise AssertionError("should not call stream")

    await wrapper.forward_once(BoomClient(), wrapper.Settings("http://x/mcp", None), {"a": 1})


@pytest.mark.asyncio
async def test_forward_once_form_encoding_non_dict_payload(monkeypatch):
    wrapper._shutdown.clear()
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda _obj: None)

    client = CapturingClient(DummyResp(200, "application/json", b'{"ok":true}'))
    settings = wrapper.Settings("http://x/mcp", None)
    settings.content_type = "application/x-www-form-urlencoded"

    await wrapper.forward_once(client, settings, "a=b")
    call = client.calls[0]
    assert call["headers"]["Content-Type"] == "application/x-www-form-urlencoded"
    assert call["data"] == b"a=b"


@pytest.mark.asyncio
async def test_forward_once_auto_detect_urlencode_and_json(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    client = CapturingClient(DummyResp(200, "application/json", b'{"ok":true}'))
    settings = wrapper.Settings("http://x/mcp", None)
    settings.content_type = "AUTO"

    await wrapper.forward_once(client, settings, {"a": "b", "c": 1})
    call1 = client.calls[-1]
    assert call1["headers"]["Content-Type"] == "application/x-www-form-urlencoded"
    assert b"a=b" in call1["data"]

    await wrapper.forward_once(client, settings, {"a": {"nested": 1}})
    call2 = client.calls[-1]
    assert call2["headers"]["Content-Type"].startswith("application/json")
    assert b'{"a":{"nested":1}}' in call2["data"]
    assert any(isinstance(o, dict) and o.get("ok") is True for o in captured)


@pytest.mark.asyncio
async def test_forward_once_returns_when_shutdown_triggered_inside_response(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    class ShutdownResp(DummyResp):
        async def __aenter__(self):
            wrapper._mark_shutdown()
            return self

    client = DummyClient(ShutdownResp(200, "application/json", b'{"ok":1}'))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"a": 1})
    assert captured == []


@pytest.mark.asyncio
async def test_forward_once_application_json_skips_processing_if_shutdown_after_read(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    class ShutdownAfterReadResp(DummyResp):
        async def aread(self):
            wrapper._mark_shutdown()
            return await super().aread()

    client = DummyClient(ShutdownAfterReadResp(200, "application/json", b'{"ok":1}'))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"a": 1})
    assert captured == []


@pytest.mark.asyncio
async def test_forward_once_process_line_returns_early_when_shutdown(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    async def fake_ndjson_lines(_resp):
        yield b'{"x": 1}'

    monkeypatch.setattr(wrapper, "ndjson_lines", fake_ndjson_lines)

    calls = {"n": 0}

    def seq_shutting_down():
        calls["n"] += 1
        # 1: start check, 2: after response, 3: outer loop check, 4: inside _process_line
        return calls["n"] == 4

    monkeypatch.setattr(wrapper, "shutting_down", seq_shutting_down)

    client = DummyClient(DummyResp(200, "application/x-ndjson", b""))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"z": 3})
    assert captured == []


@pytest.mark.asyncio
async def test_forward_once_sse_empty_payload_continues_and_breaks_on_shutdown(monkeypatch):
    wrapper._shutdown.clear()
    captured = []

    def send_and_shutdown(obj):
        captured.append(obj)
        wrapper._mark_shutdown()

    monkeypatch.setattr(wrapper, "send_to_stdout", send_and_shutdown)

    body = b"data:\n\n" b'data: {"foo": 1}\n\n' b'data: {"bar": 2}\n\n'
    client = DummyClient(DummyResp(200, "text/event-stream", body))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"w": 4})
    assert any(isinstance(d, dict) and d.get("foo") == 1 for d in captured)
    assert not any(isinstance(d, dict) and d.get("bar") == 2 for d in captured)


@pytest.mark.asyncio
async def test_forward_once_ndjson_breaks_on_shutdown(monkeypatch):
    wrapper._shutdown.clear()
    captured = []

    def send_and_shutdown(obj):
        captured.append(obj)
        wrapper._mark_shutdown()

    monkeypatch.setattr(wrapper, "send_to_stdout", send_and_shutdown)

    ndj = b'{"x":1}\n{"y":2}\n'
    client = DummyClient(DummyResp(200, "application/x-ndjson", ndj))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"z": 3})
    assert any(isinstance(d, dict) and d.get("x") == 1 for d in captured)
    assert not any(isinstance(d, dict) and d.get("y") == 2 for d in captured)


@pytest.mark.asyncio
async def test_forward_once_fallback_parses_as_ndjson(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    body = b'{"a":1}\n{"b":2}\n'
    client = DummyClient(DummyResp(200, "text/plain", body))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"z": 3})
    assert any(isinstance(d, dict) and d.get("a") == 1 for d in captured)
    assert any(isinstance(d, dict) and d.get("b") == 2 for d in captured)


@pytest.mark.asyncio
async def test_forward_once_fallback_breaks_on_shutdown(monkeypatch):
    wrapper._shutdown.clear()
    captured = []

    def send_and_shutdown(obj):
        captured.append(obj)
        wrapper._mark_shutdown()

    monkeypatch.setattr(wrapper, "send_to_stdout", send_and_shutdown)

    body = b'{"a":1}\n{"b":2}\n'
    client = DummyClient(DummyResp(200, "text/plain", body))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"z": 3})
    assert any(isinstance(d, dict) and d.get("a") == 1 for d in captured)
    assert not any(isinstance(d, dict) and d.get("b") == 2 for d in captured)


@pytest.mark.asyncio
async def test_forward_once_sse_invalid_json(monkeypatch):
    wrapper._shutdown.clear()
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    sse_chunk = b"data: notjson\n\n"
    client = DummyClient(DummyResp(200, "text/event-stream", sse_chunk))
    await wrapper.forward_once(client, wrapper.Settings("x", None), {"w": 4})
    assert any(isinstance(d, dict) and "error" in d for d in captured)


# -------------------
# make_request retry path
# -------------------
@pytest.mark.asyncio
async def test_make_request_retries(monkeypatch):
    wrapper._shutdown.clear()
    called = {"n": 0}

    async def bad_forward(*a, **k):
        called["n"] += 1
        raise RuntimeError("fail")

    monkeypatch.setattr(wrapper, "forward_once", bad_forward)
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    # small base_delay so test runs quickly
    await wrapper.make_request(None, wrapper.Settings("x", None), {"a": 1}, max_retries=2, base_delay=0.001)
    # forward_once should have been called multiple times
    assert called["n"] >= 2
    # on exhausting retries, make_request sends a max retries error
    assert any(isinstance(o, dict) and "error" in o for o in captured)


@pytest.mark.asyncio
async def test_make_request_success(monkeypatch):
    wrapper._shutdown.clear()

    called = {"n": 0}

    async def ok_forward(*_a, **_k):
        called["n"] += 1
        return None

    monkeypatch.setattr(wrapper, "forward_once", ok_forward)
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    await wrapper.make_request(None, wrapper.Settings("x", None), {"a": 1}, max_retries=2, base_delay=0.001)
    assert called["n"] == 1
    assert captured == []


@pytest.mark.asyncio
async def test_make_request_does_nothing_if_already_shutting_down(monkeypatch):
    wrapper._shutdown.clear()
    wrapper._mark_shutdown()

    async def boom(*_a, **_k):  # pragma: no cover - should not be called
        raise AssertionError("forward_once should not be called")

    monkeypatch.setattr(wrapper, "forward_once", boom)
    await wrapper.make_request(None, wrapper.Settings("x", None), {"a": 1}, max_retries=1, base_delay=0.001)


@pytest.mark.asyncio
async def test_make_request_returns_if_shutdown_triggered_in_exception(monkeypatch):
    wrapper._shutdown.clear()

    async def fail_and_shutdown(*_a, **_k):
        wrapper._mark_shutdown()
        raise RuntimeError("boom")

    monkeypatch.setattr(wrapper, "forward_once", fail_and_shutdown)
    captured = []
    monkeypatch.setattr(wrapper, "send_to_stdout", lambda obj: captured.append(obj))

    await wrapper.make_request(None, wrapper.Settings("x", None), {"a": 1}, max_retries=5, base_delay=0.001)
    # Should not emit "max retries exceeded" since shutdown was triggered.
    assert captured == []


# -------------------
# main_async smoke test
# -------------------
@pytest.mark.asyncio
async def test_main_async_smoke(monkeypatch):
    wrapper._shutdown.clear()

    worker_done = asyncio.Event()

    async def fake_reader(queue):
        await queue.put({"foo": "bar"})
        # Let the worker run before sending EOF
        await worker_done.wait()
        await queue.put(None)

    # simple make_request that just records calls
    called = {"n": 0}

    async def fake_make_request(client, settings, payload):
        called["n"] += 1
        worker_done.set()

    class DummyResilient:
        def __init__(self, *a, **k):
            pass

        async def aclose(self):
            return None

    monkeypatch.setattr(wrapper, "stdin_reader", fake_reader)
    monkeypatch.setattr(wrapper, "make_request", fake_make_request)
    monkeypatch.setattr(wrapper, "ResilientHttpClient", DummyResilient)

    settings = wrapper.Settings("http://x/mcp", None)
    await wrapper.main_async(settings)
    assert called["n"] >= 1, "make_request should have been called at least once"


@pytest.mark.asyncio
async def test_main_async_cancels_inflight_tasks(monkeypatch):
    wrapper._shutdown.clear()

    async def fake_reader(queue):
        await queue.put({"foo": "bar"})
        # Give main loop a chance to schedule the worker before EOF arrives.
        await asyncio.sleep(0)
        await queue.put(None)

    gate = asyncio.Event()

    async def blocking_make_request(*_a, **_k):
        await gate.wait()

    class DummyResilient:
        def __init__(self, *a, **k):
            self.closed = False

        async def aclose(self):
            self.closed = True

    monkeypatch.setattr(wrapper, "stdin_reader", fake_reader)
    monkeypatch.setattr(wrapper, "make_request", blocking_make_request)
    monkeypatch.setattr(wrapper, "ResilientHttpClient", DummyResilient)

    settings = wrapper.Settings("http://x/mcp", None, concurrency=1)
    await wrapper.main_async(settings)
    assert wrapper.shutting_down()


@pytest.mark.asyncio
async def test_main_async_ssl_verify_importerror_fallback(monkeypatch):
    wrapper._shutdown.clear()

    async def fake_reader(queue):
        await queue.put(None)

    created = {}

    class DummyResilient:
        def __init__(self, *a, **k):
            created["client_args"] = k.get("client_args")

        async def aclose(self):
            return None

    # Force mcpgateway.config import to fail inside main_async.
    import builtins

    real_import = builtins.__import__

    def fake_import(name, globals=None, locals=None, fromlist=(), level=0):  # noqa: A002
        if name == "mcpgateway.config":
            raise ImportError("no config")
        return real_import(name, globals, locals, fromlist, level)

    monkeypatch.setattr(builtins, "__import__", fake_import)

    monkeypatch.setattr(wrapper, "stdin_reader", fake_reader)
    monkeypatch.setattr(wrapper, "ResilientHttpClient", DummyResilient)
    monkeypatch.setattr(wrapper, "make_request", AsyncMock())

    settings = wrapper.Settings("http://x/mcp", None)
    await wrapper.main_async(settings)
    assert created["client_args"]["verify"] is True


@pytest.mark.asyncio
async def test_main_async_exits_when_already_shutting_down(monkeypatch):
    """Cover the `while not shutting_down():` condition false branch."""
    wrapper._shutdown.clear()
    wrapper._mark_shutdown()

    async def fake_reader(_queue):
        return None

    def fake_create_task(coro):
        # main_async always cancels and awaits the reader task; return a done task
        # so cancellation is a no-op and no CancelledError leaks.
        coro.close()
        fut = asyncio.get_running_loop().create_future()
        fut.set_result(None)
        return fut

    class DummyResilient:
        def __init__(self, *a, **k):
            pass

        async def aclose(self):
            return None

    monkeypatch.setattr(wrapper, "stdin_reader", fake_reader)
    monkeypatch.setattr(wrapper.asyncio, "create_task", fake_create_task)
    monkeypatch.setattr(wrapper, "ResilientHttpClient", DummyResilient)
    monkeypatch.setattr(wrapper, "make_request", AsyncMock())

    await wrapper.main_async(wrapper.Settings("http://x/mcp", None))
    assert wrapper.shutting_down()


@pytest.mark.asyncio
async def test_main_async_worker_skips_make_request_when_shutting_down(monkeypatch):
    """Cover worker branch when shutdown is observed before calling make_request."""
    wrapper._shutdown.clear()

    worker_checked = asyncio.Event()

    async def fake_reader(queue):
        await queue.put({"foo": "bar"})
        # Ensure the worker had a chance to observe shutdown before we send EOF.
        await asyncio.wait_for(worker_checked.wait(), timeout=1)
        await queue.put(None)

    # First two calls are the main loop's `while not shutting_down()` checks.
    # Third call is the worker's `if not shutting_down()` guard.
    calls = {"n": 0}

    def seq_shutting_down():
        calls["n"] += 1
        if calls["n"] == 3:
            worker_checked.set()
        return calls["n"] >= 3

    monkeypatch.setattr(wrapper, "shutting_down", seq_shutting_down)
    monkeypatch.setattr(wrapper, "stdin_reader", fake_reader)
    monkeypatch.setattr(wrapper, "make_request", AsyncMock())

    class DummyResilient:
        def __init__(self, *a, **k):
            pass

        async def aclose(self):
            return None

    monkeypatch.setattr(wrapper, "ResilientHttpClient", DummyResilient)

    await wrapper.main_async(wrapper.Settings("http://x/mcp", None, concurrency=1))
    assert wrapper.make_request.await_count == 0


# -------------------
# _install_signal_handlers runs (no-op on unsupported platforms)
# -------------------
def test_install_signal_handlers_runs():
    loop = asyncio.new_event_loop()
    try:
        wrapper._install_signal_handlers(loop)
    finally:
        loop.close()


def test_install_signal_handlers_sig_none_and_notimplemented(monkeypatch):
    # Cover both the `sig is None` branch and suppression of NotImplementedError.
    calls = {"n": 0}

    class DummyLoop:
        def add_signal_handler(self, *_a, **_k):
            calls["n"] += 1
            raise NotImplementedError

    monkeypatch.setattr(wrapper.signal, "SIGINT", None, raising=False)
    monkeypatch.setattr(wrapper.signal, "SIGTERM", None, raising=False)
    wrapper._install_signal_handlers(DummyLoop())
    assert calls["n"] == 0

    # Restore real signals for the NotImplementedError suppression path.
    import signal as _signal

    monkeypatch.setattr(wrapper.signal, "SIGINT", _signal.SIGINT, raising=False)
    monkeypatch.setattr(wrapper.signal, "SIGTERM", _signal.SIGTERM, raising=False)
    wrapper._install_signal_handlers(DummyLoop())


def test_main_runs_with_patched_deps(monkeypatch):
    wrapper._shutdown.clear()

    monkeypatch.setattr(wrapper, "parse_args", lambda: wrapper.Settings("http://x/mcp", None, log_level="DEBUG"))
    monkeypatch.setattr(wrapper, "_install_signal_handlers", lambda _loop: None)

    async def fake_main_async(_settings):
        await asyncio.sleep(0)

    monkeypatch.setattr(wrapper, "main_async", fake_main_async)
    # Avoid altering global event loop in the test process.
    monkeypatch.setattr(wrapper.asyncio, "set_event_loop", lambda _loop: None)

    wrapper.main()


def test_main_runs_with_logging_disabled(monkeypatch):
    """Cover main() branches when logging is disabled."""
    wrapper._shutdown.clear()

    monkeypatch.setattr(wrapper, "parse_args", lambda: wrapper.Settings("http://x/mcp", None, log_level=None))
    monkeypatch.setattr(wrapper, "_install_signal_handlers", lambda _loop: None)

    async def fake_main_async(_settings):
        await asyncio.sleep(0)

    monkeypatch.setattr(wrapper, "main_async", fake_main_async)
    monkeypatch.setattr(wrapper.asyncio, "set_event_loop", lambda _loop: None)

    wrapper.main()
