"""Reproduction: talk-streaming.add-tests-for-streaming-path.

Expected: talk streaming handles chunk/done/error events, falls back cleanly when
streaming is unavailable, supports /stop during in-flight interact, and preserves
existing history replay + response contracts.
Actual:   One or more of these guarantees may fail.
"""

from __future__ import annotations

import asyncio
import io
import shutil
import sys
import time
import types
import uuid
from collections.abc import Awaitable, Callable, Iterator
from pathlib import Path
from typing import Any, cast

import pytest
from rich.console import Console
from returns.result import Success

from anima.shell.mcp_agent import AgentMCPServer


class _FakeTextContent:
    def __init__(self, text: str) -> None:
        self.text = text


class _FakeToolResult:
    def __init__(self, text: str) -> None:
        self.content = [_FakeTextContent(text)]


class _FakeStreamContext:
    async def __aenter__(self) -> tuple[object, object, Any]:
        return object(), object(), lambda: "session-id"

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: Any,
    ) -> None:
        return None


def _install_fake_mcp(
    monkeypatch: pytest.MonkeyPatch,
    call_tool: Callable[[str, dict[str, Any]], Awaitable[_FakeToolResult]],
) -> None:
    class FakeClientSession:
        async def __aenter__(self) -> "FakeClientSession":
            return self

        async def __aexit__(
            self,
            exc_type: type[BaseException] | None,
            exc: BaseException | None,
            tb: Any,
        ) -> None:
            return None

        def __init__(self, read_stream: object, write_stream: object) -> None:
            self.read_stream = read_stream
            self.write_stream = write_stream

        async def initialize(self) -> None:
            return None

        async def call_tool(self, tool: str, params: dict[str, Any]) -> _FakeToolResult:
            return await call_tool(tool, params)

    mcp_module = types.ModuleType("mcp")
    setattr(mcp_module, "ClientSession", FakeClientSession)

    client_module = types.ModuleType("mcp.client")

    stream_module = types.ModuleType("mcp.client.streamable_http")
    setattr(
        stream_module,
        "streamable_http_client",
        lambda _path, http_client=None: _FakeStreamContext(),
    )

    mcp_types_module = types.ModuleType("mcp.types")
    setattr(mcp_types_module, "TextContent", _FakeTextContent)

    monkeypatch.setitem(sys.modules, "mcp", mcp_module)
    monkeypatch.setitem(sys.modules, "mcp.client", client_module)
    monkeypatch.setitem(sys.modules, "mcp.client.streamable_http", stream_module)
    monkeypatch.setitem(sys.modules, "mcp.types", mcp_types_module)


def _make_input_handler(inputs_iter: Iterator[str]) -> Callable[[str], str]:
    def _input(_prompt: str) -> str:
        try:
            return next(inputs_iter)
        except StopIteration as exc:
            raise EOFError("end of test input") from exc

    return _input


def _run_talk(console: Console, stream: io.StringIO, *, json_output: bool) -> str:
    from anima.shell.cli_commands.talk_deps import default_deps
    from anima.shell.cli_commands.talk_runtime import run_talk_session

    asyncio.run(
        run_talk_session(
            socket_path=Path("/tmp/fake.sock"),
            agent_id="stream-test-agent",
            json_output=json_output,
            console=console,
            deps=default_deps(),
        )
    )
    return stream.getvalue()


def test_stream_events_chunk_done_error_are_handled() -> None:
    from anima.shell.cli_commands.talk_streaming import (
        _render_stream_chunk,
        _render_stream_done,
        _render_stream_error,
        parse_sse_events,
    )

    lines = [
        "event: chunk",
        'data: {"turn": 1, "seq": 1, "chunk": "hello"}',
        "",
        "event: done",
        (
            'data: {"turn": 1, "seq": 2, "done": true, '
            '"stop_reason": "complete", "final_response_chars": 5}'
        ),
        "",
        "event: error",
        'data: {"turn": 1, "seq": 3, "error": "boom", "code": "INTERNAL"}',
        "",
    ]
    events = parse_sse_events(lines)

    captured = io.StringIO()
    console = Console(file=captured, force_terminal=True, width=200)
    _render_stream_chunk(events[0], "stream-test-agent", True, console)
    _render_stream_done(events[1], "stream-test-agent", True, console)
    _render_stream_error(events[2], "stream-test-agent", True, console)

    output = captured.getvalue()
    assert "hello" in output, f"Expected chunk payload in output, got: {output}"
    assert "complete" in output, f"Expected done stop_reason in output, got: {output}"
    assert "boom" in output and "INTERNAL" in output, (
        f"Expected stream error payload in output, got: {output}"
    )


def test_fallback_when_stream_unavailable_uses_interact_path(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from anima.shell.cli_commands import talk_deps

    calls: list[str] = []
    probe_count = 0

    async def fake_call_tool(tool: str, _params: dict[str, Any]) -> _FakeToolResult:
        calls.append(tool)
        if tool == "anima.interact":
            return _FakeToolResult('{"response":"fallback reply","turn_number":3}')
        raise AssertionError(f"Unexpected tool call: {tool}")

    async def fake_probe(*args: Any, **kwargs: Any) -> bool:
        nonlocal probe_count
        probe_count += 1
        return False

    async def fail_if_stream_consumed(
        *args: Any, **kwargs: Any
    ) -> tuple[bool, str | None]:
        raise AssertionError("consume_sse_stream should not run in fallback mode")

    _install_fake_mcp(monkeypatch, fake_call_tool)
    monkeypatch.setattr(talk_deps, "probe_streaming_support", fake_probe)
    monkeypatch.setattr(talk_deps, "consume_sse_stream", fail_if_stream_consumed)

    stream = io.StringIO()
    console = Console(file=stream, force_terminal=True, width=200)
    setattr(console, "input", _make_input_handler(iter(["hello", "/quit"])))
    output = _run_talk(console, stream, json_output=False)

    assert "fallback reply" in output, (
        f"Expected blocking interact output when stream unavailable, got: {output}"
    )
    assert calls.count("anima.interact") == 1, f"Unexpected tool calls: {calls}"
    assert probe_count == 1, f"Expected one stream probe, got: {probe_count}"


def test_stop_mid_stream_calls_stop_turn_and_cancels_interact(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from anima.shell.cli_commands import talk_deps

    calls: list[str] = []
    interact_cancelled = asyncio.Event()
    stop_called = asyncio.Event()

    async def fake_call_tool(tool: str, _params: dict[str, Any]) -> _FakeToolResult:
        calls.append(tool)
        if tool == "anima.interact":
            try:
                await asyncio.sleep(5.0)
            except asyncio.CancelledError:
                interact_cancelled.set()
                raise
            return _FakeToolResult('{"response":"late reply","turn_number":1}')
        if tool == "anima.stop_turn":
            stop_called.set()
            return _FakeToolResult('{"status":"stopped"}')
        raise AssertionError(f"Unexpected tool call: {tool}")

    async def fake_probe(*args: Any, **kwargs: Any) -> bool:
        return True

    async def fake_consume(*args: Any, **kwargs: Any) -> tuple[bool, str | None]:
        await asyncio.wait_for(stop_called.wait(), timeout=2.0)
        return (False, "stopped")

    _install_fake_mcp(monkeypatch, fake_call_tool)
    monkeypatch.setattr(talk_deps, "probe_streaming_support", fake_probe)
    monkeypatch.setattr(talk_deps, "consume_sse_stream", fake_consume)

    stream = io.StringIO()
    console = Console(file=stream, force_terminal=True, width=200)
    setattr(console, "input", _make_input_handler(iter(["hello", "/stop", "/quit"])))
    _run_talk(console, stream, json_output=False)

    assert "anima.stop_turn" in calls, (
        f"Expected stop_turn call in tool history, got: {calls}"
    )
    assert interact_cancelled.is_set(), (
        "Expected in-flight interact to be cancelled on /stop"
    )


def test_explicit_stream_404_fallback_uses_interact_path(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Explicit /stream 404 fallback scenario.

    Scenario 4 from original step: When /stream endpoint returns 404,
    client correctly falls back to non-streaming interact path.

    This test differs from test_fallback_when_stream_unavailable_uses_interact_path
    by explicitly asserting the 404 response handling, not just generic unavailability.
    """
    from anima.shell.cli_commands import talk_deps
    import httpx

    calls: list[str] = []
    probe_get_call_count = 0

    class FakeAsyncClient:
        """Fake httpx.AsyncClient that returns 404 for /stream endpoint."""

        async def get(self, url: str) -> httpx.Response:
            nonlocal probe_get_call_count
            probe_get_call_count += 1
            # Explicitly return 404 for /stream endpoint probe
            if "/stream" in url:
                return httpx.Response(404, request=httpx.Request("GET", url))
            return httpx.Response(200, request=httpx.Request("GET", url))

        def stream(self, method: str, url: str, **kwargs: Any) -> "_FakeStreamContext":
            """Return async context manager for streaming requests.

            The probe uses client.stream("GET", url) which must be an async context
            manager that yields a Response. We return 404 for /stream endpoint.
            """
            nonlocal probe_get_call_count
            probe_get_call_count += 1

            class _FakeStreamResponse:
                """Async context manager yielding httpx.Response."""

                async def __aenter__(self) -> httpx.Response:
                    # Return 404 for /stream endpoint (matching get() behavior)
                    if "/stream" in url:
                        return httpx.Response(404, request=httpx.Request(method, url))
                    return httpx.Response(200, request=httpx.Request(method, url))

                async def __aexit__(self, *args: Any) -> None:
                    pass

            return _FakeStreamResponse()  # type: ignore[return-value]

        async def __aenter__(self) -> "FakeAsyncClient":
            return self

        async def __aexit__(self, *args: Any) -> None:
            pass

    async def fake_call_tool(tool: str, _params: dict[str, Any]) -> _FakeToolResult:
        calls.append(tool)
        if tool == "anima.interact":
            return _FakeToolResult('{"response":"404 fallback reply","turn_number":1}')
        raise AssertionError(f"Unexpected tool call: {tool}")

    async def fail_if_stream_consumed(
        *args: Any, **kwargs: Any
    ) -> tuple[bool, str | None]:
        raise AssertionError("consume_sse_stream MUST NOT run when /stream returns 404")

    def fake_create_client(*args: Any, **kwargs: Any) -> FakeAsyncClient:
        return FakeAsyncClient()

    _install_fake_mcp(monkeypatch, fake_call_tool)
    # Patch httpx.AsyncClient to return our fake client with 404 response
    monkeypatch.setattr(httpx, "AsyncClient", fake_create_client)
    monkeypatch.setattr(talk_deps, "consume_sse_stream", fail_if_stream_consumed)

    stream = io.StringIO()
    console = Console(file=stream, force_terminal=True, width=200)
    setattr(console, "input", _make_input_handler(iter(["test prompt", "/quit"])))
    output = _run_talk(console, stream, json_output=False)

    # Strip ANSI codes for assertion (Rich adds color codes)
    import re

    output_clean = re.sub(r"\x1b\[[0-9;]*m", "", output)

    # EXPLICIT ASSERTIONS FOR 404 FALLBACK SCENARIO:
    # 1. /stream endpoint was probed exactly once with GET request
    assert probe_get_call_count == 1, (
        f"Expected exactly one GET probe for /stream endpoint, got: {probe_get_call_count}"
    )

    # 2. Fallback to blocking interact happened (stream was not used)
    assert "404 fallback reply" in output_clean, (
        f"Expected fallback reply when /stream returns 404, got: {output_clean}"
    )
    assert calls.count("anima.interact") == 1, (
        f"Expected one interact call in 404 fallback, got: {calls}"
    )

    # 3. SSE stream consumer was never invoked
    # (assertion is implicit via fail_if_stream_consumed raising if called)


def test_history_replay_and_response_contract_remain_compatible(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from anima.shell.cli_commands import talk_deps
    from anima.shell.cli_commands.talk_helpers import (
        parse_talk_response,
    )
    from anima.shell.cli_commands.talk_session import replay_talk_history
    from returns.result import Success

    async def fake_call_tool(tool: str, _params: dict[str, Any]) -> _FakeToolResult:
        if tool == "anima.interact":
            return _FakeToolResult('{"response":"contract reply","turn_number":9}')
        raise AssertionError(f"Unexpected tool call: {tool}")

    async def fake_probe(*args: Any, **kwargs: Any) -> bool:
        return False

    _install_fake_mcp(monkeypatch, fake_call_tool)
    monkeypatch.setattr(talk_deps, "probe_streaming_support", fake_probe)

    monkeypatch.setattr(
        talk_deps,
        "read_log",
        lambda _agent_id: Success(
            [
                {
                    "event": "response",
                    "details": {
                        "response_text": "historical reply",
                        "turn_number": 1,
                    },
                }
            ]
        ),
    )

    history_stream = io.StringIO()
    history_console = Console(file=history_stream, force_terminal=True, width=200)
    replay_talk_history(
        agent_id="stream-test-agent",
        json_output=False,
        console=history_console,
        deps=cast(
            Any,
            types.SimpleNamespace(
                read_log=lambda _agent_id: Success(
                    [
                        {
                            "event": "response",
                            "details": {
                                "response_text": "historical reply",
                                "turn_number": 1,
                            },
                        }
                    ]
                )
            ),
        ),
    )
    history_output = history_stream.getvalue()

    stream = io.StringIO()
    console = Console(file=stream, force_terminal=True, width=200)
    setattr(console, "input", _make_input_handler(iter(["hi", "/quit"])))
    output = _run_talk(console, stream, json_output=True)

    parsed = parse_talk_response('{"response":"contract reply","turn_number":9}')
    assert parsed == {"response": "contract reply", "turn_number": "9"}
    assert "historical reply" in history_output and "live" in history_output.lower(), (
        f"Expected history replay output compatibility, got: {history_output}"
    )
    assert "contract reply" in output and "turn_number" in output, (
        f"Expected existing interact response contract in output, got: {output}"
    )


# ============================================================================
# Integration-style tests for MCP I/O boundary (session -> transport -> output)
# Real MCP seam: no monkeypatching of session or streaming transport in path
# ============================================================================


class _RunnerStub:
    def __init__(self, instance_id: str) -> None:
        self.instance_id = instance_id


@pytest.mark.asyncio
async def test_real_talk_mcp_stream_path_executes_consume_sse_stream(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    """Talk session crosses real MCP session + real stream transport path."""
    from anima.shell.cli_commands.talk_deps import default_deps
    from anima.shell.cli_commands.talk_runtime import run_talk_session

    state_dir = Path(f"/tmp/anima_stream_{uuid.uuid4().hex[:8]}")
    state_dir.mkdir(parents=True, exist_ok=True)
    monkeypatch.setenv("ANIMA_STATE", str(state_dir))

    instance_id = "inst_sr1"
    server = AgentMCPServer(instance_id, _RunnerStub(instance_id))
    start_result = await server.start()
    assert isinstance(start_result, Success), (
        f"Expected real MCP server start success, got: {start_result}"
    )
    assert server.socket_path is not None

    stream = io.StringIO()
    console = Console(file=stream, force_terminal=True, width=200)
    input_call_count = 0

    def _timed_input(_prompt: str) -> str:
        nonlocal input_call_count
        input_call_count += 1
        if input_call_count == 1:
            return "hello"
        time.sleep(0.5)
        return "/quit"

    setattr(console, "input", _timed_input)
    stream_consumer_seen = False

    async def _drive_agent_response_and_stream() -> None:
        nonlocal stream_consumer_seen
        message = await asyncio.wait_for(
            server.get_pending_message().get(), timeout=3.0
        )
        assert message == "hello"

        for _ in range(50):
            if server._stream_consumer_active:
                stream_consumer_seen = True
                break
            await asyncio.sleep(0.02)

        await server.start_stream_session(turn_number=1)
        await server.enqueue_stream_chunk(turn_number=1, chunk="real transport chunk")
        await server.enqueue_stream_done(
            turn_number=1,
            stop_reason="complete",
            final_response_chars=20,
            prompt_tokens=3,
            completion_tokens=4,
            total_tokens=7,
        )
        server.set_response("final response", turn_number=1)

    driver_task = asyncio.create_task(_drive_agent_response_and_stream())
    try:
        await asyncio.wait_for(
            run_talk_session(
                socket_path=server.socket_path,
                agent_id=instance_id,
                json_output=False,
                console=console,
                deps=default_deps(),
            ),
            timeout=8.0,
        )
        await asyncio.wait_for(driver_task, timeout=2.0)
    finally:
        stop_result = await server.stop()
        assert isinstance(stop_result, Success), (
            f"Expected real MCP server stop success, got: {stop_result}"
        )
        shutil.rmtree(state_dir, ignore_errors=True)

    assert stream_consumer_seen, "Expected consume_sse_stream to connect to /stream"
    output = stream.getvalue()
    assert "final response" not in output, (
        "Expected interact final output suppression after streamed chunk rendering"
    )


@pytest.mark.asyncio
async def test_consume_sse_stream_reports_concrete_transport_failure(
    tmp_path: Path,
) -> None:
    """Concrete transport failure yields diagnostics from real socket error."""
    from anima.shell.cli_commands.talk_streaming import consume_sse_stream

    missing_socket = tmp_path / "missing-mcp.sock"
    captured = io.StringIO()
    console = Console(file=captured, force_terminal=True, width=200)
    stream_complete = asyncio.Event()

    had_error, stop_reason = await consume_sse_stream(
        socket_path=missing_socket,
        agent_id="inst_missing_socket",
        json_output=False,
        console=console,
        stream_complete=stream_complete,
    )

    assert had_error is True
    assert stop_reason is None
    assert stream_complete.is_set(), (
        "Expected stream_complete to be signaled on failure"
    )

    output = captured.getvalue().lower()
    assert "stream connection error" in output, (
        f"Expected transport diagnostic prefix, got: {captured.getvalue()}"
    )
    assert any(
        fragment in output for fragment in ("no such file", "connect", "socket")
    ), "Expected concrete transport failure details from real exception"


if __name__ == "__main__":
    raise SystemExit(pytest.main([__file__, "-q"]))
