"""Tests for the session streaming engine."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from unittest.mock import AsyncMock, MagicMock

import pytest
from icarus.session.events import (
    StatusUpdate,
    TextComplete,
    TextDelta,
    TokenUsage,
    ToolCallStart,
    ToolResult,
)
from icarus.session.stream import (
    _build_interrupted_ai_message,
    _is_summarization_chunk,
    _preprocess_input,
    build_stream_config,
    stream_turn,
)
from langchain_core.messages import ToolMessage

# ---------------------------------------------------------------------------
# Helpers for building mock LangGraph chunks
# ---------------------------------------------------------------------------

@dataclass
class FakeAIMessageChunk:
    """Minimal stand-in for AIMessageChunk with content_blocks."""

    content_blocks: list[dict]
    usage_metadata: dict | None = None
    chunk_position: str | None = None


@dataclass
class FakeInterrupt:
    id: str
    value: dict


class MockAgent:
    """Mock LangGraph agent that yields pre-configured chunks from astream."""

    def __init__(self, chunks: list[tuple]):
        self._chunks = chunks
        self.aupdate_state = AsyncMock()

    async def astream(self, *args, **kwargs):
        for chunk in self._chunks:
            yield chunk


def _text_chunk(text: str, position: str | None = None) -> tuple:
    """Build a main-agent text chunk."""
    msg = FakeAIMessageChunk(
        content_blocks=[{"type": "text", "text": text}],
        chunk_position=position,
    )
    return ([], "messages", (msg, {}))


def _text_chunk_with_tokens(text: str, total: int) -> tuple:
    """Build a text chunk that also carries token usage."""
    msg = FakeAIMessageChunk(
        content_blocks=[{"type": "text", "text": text}],
        usage_metadata={"total_tokens": total},
        chunk_position="last",
    )
    return ([], "messages", (msg, {}))


def _tool_call_chunk(name: str, args: dict, tool_id: str) -> tuple:
    """Build a complete tool_call block (non-streaming)."""
    msg = FakeAIMessageChunk(
        content_blocks=[{
            "type": "tool_call",
            "name": name,
            "args": args,
            "id": tool_id,
            "index": None,
        }],
    )
    return ([], "messages", (msg, {}))


def _tool_result_chunk(
    content: str,
    tool_id: str = "tool-1",
    name: str = "some_tool",
    status: str = "success",
) -> tuple:
    """Build a ToolMessage result chunk."""
    msg = ToolMessage(
        content=content, name=name, tool_call_id=tool_id, status=status,
    )
    return ([], "messages", (msg, {}))


def _interrupt_chunk(interrupt_id: str, action_requests: list[dict]) -> tuple:
    """Build an updates chunk with an __interrupt__."""
    # HITLRequest requires review_configs alongside action_requests
    review_configs = [
        {
            "action_name": ar["name"],
            "allowed_decisions": ["approve", "reject"],
        }
        for ar in action_requests
    ]
    interrupt = FakeInterrupt(
        id=interrupt_id,
        value={
            "action_requests": action_requests,
            "review_configs": review_configs,
        },
    )
    return ([], "updates", {"__interrupt__": [interrupt]})


# ---------------------------------------------------------------------------
# Tests: build_stream_config
# ---------------------------------------------------------------------------

class TestBuildStreamConfig:
    def test_basic(self):
        cfg = build_stream_config("thread-1", "agent-1")
        assert cfg["configurable"]["thread_id"] == "thread-1"
        assert cfg["metadata"]["assistant_id"] == "agent-1"

    def test_no_assistant(self):
        cfg = build_stream_config("thread-1", None)
        assert cfg["configurable"]["thread_id"] == "thread-1"
        assert cfg["metadata"] == {}


# ---------------------------------------------------------------------------
# Tests: _is_summarization_chunk
# ---------------------------------------------------------------------------

class TestIsSummarizationChunk:
    def test_none(self):
        assert not _is_summarization_chunk(None)

    def test_summarization(self):
        assert _is_summarization_chunk({"lc_source": "summarization"})

    def test_normal(self):
        assert not _is_summarization_chunk({"lc_source": "other"})


# ---------------------------------------------------------------------------
# Tests: _build_interrupted_ai_message
# ---------------------------------------------------------------------------

class TestBuildInterruptedAIMessage:
    def test_empty(self):
        assert _build_interrupted_ai_message("", {}) is None

    def test_text_only(self):
        msg = _build_interrupted_ai_message("hello world", {})
        assert msg is not None
        assert msg.content == "hello world"
        assert msg.tool_calls == []

    def test_with_tools(self):
        tools = {"t1": ("read_file", {"path": "/a"})}
        msg = _build_interrupted_ai_message("", tools)
        assert msg is not None
        assert len(msg.tool_calls) == 1
        assert msg.tool_calls[0]["name"] == "read_file"

    def test_whitespace_only(self):
        assert _build_interrupted_ai_message("   \n  ", {}) is None


# ---------------------------------------------------------------------------
# Tests: _preprocess_input
# ---------------------------------------------------------------------------

class TestPreprocessInput:
    def test_plain_text(self):
        result = _preprocess_input("hello world")
        assert result == "hello world"

    def test_with_image_tracker(self):
        tracker = MagicMock()
        tracker.get_images.return_value = []
        result = _preprocess_input("hello", image_tracker=tracker)
        assert result == "hello"
        tracker.clear.assert_called_once()

    def test_no_image_tracker(self):
        result = _preprocess_input("test message", image_tracker=None)
        assert result == "test message"


# ---------------------------------------------------------------------------
# Tests: stream_turn — text streaming
# ---------------------------------------------------------------------------

class TestStreamTurnText:
    @pytest.fixture
    def agent_with_text(self):
        return MockAgent([
            _text_chunk("Hello "),
            _text_chunk("world!", "last"),
        ])

    async def test_text_events(self, agent_with_text):
        events = []
        async for event in stream_turn(
            agent=agent_with_text,
            user_input="hi",
            thread_id="t1",
        ):
            events.append(event)

        # StatusUpdate("Thinking") first
        assert isinstance(events[0], StatusUpdate)
        assert events[0].status == "Thinking"

        # StatusUpdate(None) when text starts
        assert isinstance(events[1], StatusUpdate)
        assert events[1].status is None

        # TextDelta for "Hello "
        assert isinstance(events[2], TextDelta)
        assert events[2].text == "Hello "

        # TextDelta for "world!"
        assert isinstance(events[3], TextDelta)
        assert events[3].text == "world!"

        # TextComplete flush (triggered by chunk_position="last")
        assert isinstance(events[4], TextComplete)
        assert events[4].full_text == "Hello world!"

    async def test_message_ids_consistent(self, agent_with_text):
        text_events = []
        async for event in stream_turn(
            agent=agent_with_text,
            user_input="hi",
            thread_id="t1",
        ):
            if isinstance(event, (TextDelta, TextComplete)):
                text_events.append(event)

        # All text events for one message share the same ID
        ids = {e.message_id for e in text_events}
        assert len(ids) == 1
        assert all(id_.startswith("asst-") for id_ in ids)


# ---------------------------------------------------------------------------
# Tests: stream_turn — tool calls
# ---------------------------------------------------------------------------

class TestStreamTurnToolCalls:
    async def test_tool_call_and_result(self):
        agent = MockAgent([
            _text_chunk("Let me check.", "last"),
            _tool_call_chunk("read_file", {"path": "/a.txt"}, "tc-1"),
            _tool_result_chunk("file contents", tool_id="tc-1", name="read_file"),
            _text_chunk("Done!", "last"),
        ])

        events = []
        async for event in stream_turn(
            agent=agent,
            user_input="read file",
            thread_id="t1",
        ):
            events.append(event)

        types = [type(e).__name__ for e in events]

        assert "TextComplete" in types
        assert "ToolCallStart" in types
        assert "ToolResult" in types

        # Find the ToolCallStart
        tc_start = next(e for e in events if isinstance(e, ToolCallStart))
        assert tc_start.tool_name == "read_file"
        assert tc_start.tool_id == "tc-1"
        assert tc_start.args == {"path": "/a.txt"}

        # Find the ToolResult
        tc_result = next(e for e in events if isinstance(e, ToolResult))
        assert tc_result.tool_id == "tc-1"
        assert tc_result.status == "success"

    async def test_tool_error(self):
        agent = MockAgent([
            _tool_call_chunk("bad_tool", {}, "tc-err"),
            _tool_result_chunk("not found", tool_id="tc-err", status="error"),
        ])

        events = []
        async for event in stream_turn(
            agent=agent,
            user_input="do thing",
            thread_id="t1",
        ):
            events.append(event)

        result = next(e for e in events if isinstance(e, ToolResult))
        assert result.status == "error"


# ---------------------------------------------------------------------------
# Tests: stream_turn — token usage
# ---------------------------------------------------------------------------

class TestStreamTurnTokens:
    async def test_token_usage_emitted(self):
        agent = MockAgent([
            _text_chunk_with_tokens("answer", total=1500),
        ])

        events = []
        async for event in stream_turn(
            agent=agent,
            user_input="question",
            thread_id="t1",
        ):
            events.append(event)

        token_events = [e for e in events if isinstance(e, TokenUsage)]
        assert len(token_events) == 1
        assert token_events[0].total_tokens == 1500


# ---------------------------------------------------------------------------
# Tests: stream_turn — HITL (auto-approve)
# ---------------------------------------------------------------------------

class TestStreamTurnHITLAutoApprove:
    async def test_auto_approve(self):
        """When auto_approve=True, ToolsApproved is yielded and stream resumes."""
        action_requests = [{"name": "execute", "args": {"command": "ls"}}]

        # First astream returns text + interrupt, second returns more text
        call_count = 0

        class ResumeAgent:
            aupdate_state = AsyncMock()

            async def astream(self, stream_input, **kwargs):
                nonlocal call_count
                if call_count == 0:
                    call_count += 1
                    yield _text_chunk("Running...", "last")
                    yield _tool_call_chunk("execute", {"command": "ls"}, "tc-1")
                    yield _interrupt_chunk("int-1", action_requests)
                else:
                    yield _tool_result_chunk("output", tool_id="tc-1", name="execute")
                    yield _text_chunk("Done.", "last")

        events = []
        async for event in stream_turn(
            agent=ResumeAgent(),
            user_input="run ls",
            thread_id="t1",
            auto_approve=True,
        ):
            events.append(event)

        types = [type(e).__name__ for e in events]
        assert "ToolsApproved" in types
        # Stream should have resumed after approval
        assert types.count("TextComplete") >= 2


# ---------------------------------------------------------------------------
# Tests: stream_turn — HITL (manual reject)
# ---------------------------------------------------------------------------

class TestStreamTurnHITLReject:
    async def test_reject(self):
        """When resolve_hitl returns reject, ToolsRejected is yielded."""
        action_requests = [{"name": "write_file", "args": {"file_path": "/x"}}]

        class InterruptAgent:
            aupdate_state = AsyncMock()

            async def astream(self, stream_input, **kwargs):
                yield _tool_call_chunk("write_file", {"file_path": "/x"}, "tc-1")
                yield _interrupt_chunk("int-1", action_requests)

        async def reject_hitl(action_requests, assistant_id):
            return {"type": "reject"}

        events = []
        async for event in stream_turn(
            agent=InterruptAgent(),
            user_input="write file",
            thread_id="t1",
            resolve_hitl=reject_hitl,
        ):
            events.append(event)

        types = [type(e).__name__ for e in events]
        assert "ToolsRejected" in types
        # No further events after rejection
        rejected_idx = types.index("ToolsRejected")
        assert rejected_idx == len(types) - 1


# ---------------------------------------------------------------------------
# Tests: stream_turn — HITL (manual approve)
# ---------------------------------------------------------------------------

class TestStreamTurnHITLApprove:
    async def test_approve(self):
        """When resolve_hitl returns approve, ToolsApproved is yielded."""
        action_requests = [{"name": "execute", "args": {"command": "echo hi"}}]
        call_count = 0

        class ApproveAgent:
            aupdate_state = AsyncMock()

            async def astream(self, stream_input, **kwargs):
                nonlocal call_count
                if call_count == 0:
                    call_count += 1
                    yield _tool_call_chunk("execute", {"command": "echo hi"}, "tc-1")
                    yield _interrupt_chunk("int-1", action_requests)
                else:
                    yield _tool_result_chunk("hi", tool_id="tc-1", name="execute")
                    yield _text_chunk("Done.", "last")

        async def approve_hitl(action_requests, assistant_id):
            return {"type": "approve"}

        events = []
        async for event in stream_turn(
            agent=ApproveAgent(),
            user_input="run",
            thread_id="t1",
            resolve_hitl=approve_hitl,
        ):
            events.append(event)

        types = [type(e).__name__ for e in events]
        assert "ToolsApproved" in types
        assert "ToolResult" in types

    async def test_auto_approve_all(self):
        """auto_approve_all yields AutoApproveEnabled + ToolsApproved."""
        action_requests = [{"name": "execute", "args": {"command": "pwd"}}]
        call_count = 0

        class Agent:
            aupdate_state = AsyncMock()

            async def astream(self, stream_input, **kwargs):
                nonlocal call_count
                if call_count == 0:
                    call_count += 1
                    yield _tool_call_chunk("execute", {"command": "pwd"}, "tc-1")
                    yield _interrupt_chunk("int-1", action_requests)
                else:
                    yield _tool_result_chunk("/home", tool_id="tc-1")
                    yield _text_chunk("here", "last")

        async def auto_all_hitl(action_requests, assistant_id):
            return {"type": "auto_approve_all"}

        events = []
        async for event in stream_turn(
            agent=Agent(),
            user_input="where",
            thread_id="t1",
            resolve_hitl=auto_all_hitl,
        ):
            events.append(event)

        types = [type(e).__name__ for e in events]
        assert "AutoApproveEnabled" in types
        assert "ToolsApproved" in types


# ---------------------------------------------------------------------------
# Tests: stream_turn — subagent filtering
# ---------------------------------------------------------------------------

class TestStreamTurnSubagentFiltering:
    async def test_subagent_messages_filtered(self):
        """Messages from subagents (non-empty namespace) are not emitted."""
        agent = MockAgent([
            # Main agent text
            _text_chunk("main", "last"),
            # Subagent text (namespace = ["subagent:research"])
            (["subagent:research"], "messages", (
                FakeAIMessageChunk(
                    content_blocks=[{"type": "text", "text": "sub text"}],
                    chunk_position="last",
                ),
                {},
            )),
        ])

        events = []
        async for event in stream_turn(
            agent=agent,
            user_input="test",
            thread_id="t1",
        ):
            events.append(event)

        text_deltas = [e for e in events if isinstance(e, TextDelta)]
        assert len(text_deltas) == 1
        assert text_deltas[0].text == "main"


# ---------------------------------------------------------------------------
# Tests: stream_turn — summarization filtering
# ---------------------------------------------------------------------------

class TestStreamTurnSummarizationFiltering:
    async def test_summarization_chunks_filtered(self):
        """Chunks with lc_source=summarization metadata are skipped."""
        msg = FakeAIMessageChunk(
            content_blocks=[{"type": "text", "text": "summary"}],
            chunk_position="last",
        )
        agent = MockAgent([
            ([], "messages", (msg, {"lc_source": "summarization"})),
            _text_chunk("real text", "last"),
        ])

        events = []
        async for event in stream_turn(
            agent=agent,
            user_input="test",
            thread_id="t1",
        ):
            events.append(event)

        text_deltas = [e for e in events if isinstance(e, TextDelta)]
        assert len(text_deltas) == 1
        assert text_deltas[0].text == "real text"


# ---------------------------------------------------------------------------
# Tests: stream_turn — cancellation
# ---------------------------------------------------------------------------

class TestStreamTurnCancellation:
    async def test_cancelled_error_propagates(self):
        """CancelledError saves state and re-raises."""

        class CancelAgent:
            aupdate_state = AsyncMock()

            async def astream(self, *args, **kwargs):
                yield _text_chunk("start")
                raise asyncio.CancelledError()

        with pytest.raises(asyncio.CancelledError):
            async for _ in stream_turn(
                agent=CancelAgent(),
                user_input="test",
                thread_id="t1",
            ):
                pass

        # Verify state was saved
        CancelAgent.aupdate_state.assert_called()
