from __future__ import annotations

from dataclasses import dataclass
from typing import AsyncIterator
from unittest.mock import AsyncMock, MagicMock, patch

import mistralai
import pytest


@dataclass
class _Chunk:
    data: object


class _FakeEventStream:
    """Mimics EventStreamAsync[ConversationEvents] with an async generator."""

    def __init__(self, chunks: list[_Chunk]) -> None:
        self._chunks = chunks

    @property
    async def generator(self) -> AsyncIterator[_Chunk]:
        for chunk in self._chunks:
            yield chunk


def _function_call_chunk(
    output_index: int = 0,
    tool_call_id: str | None = None,
    name: str | None = None,
    arguments: str | None = None,
    id: str | None = None,
) -> _Chunk:
    return _Chunk(
        data=mistralai.FunctionCallEvent.model_construct(
            output_index=output_index,
            tool_call_id=tool_call_id,
            name=name,
            arguments=arguments,
            id=id or "",
            type="function.call.delta",
        )
    )


def _tool_execution_chunk(
    output_index: int = 0,
    name: str | None = None,
    arguments: str | None = None,
    id: str | None = None,
) -> _Chunk:
    return _Chunk(
        data=mistralai.ToolExecutionDeltaEvent.model_construct(
            output_index=output_index,
            name=name,
            arguments=arguments,
            id=id,
            type="tool.execution.delta",
        )
    )


def _tool_execution_done_chunk(output_index: int = 0) -> _Chunk:
    return _Chunk(
        data=mistralai.ToolExecutionDoneEvent.model_construct(
            output_index=output_index,
            id="",
            name="",
            type="tool.execution.done",
        )
    )


def _response_started_chunk(conversation_id: str = "conv-1") -> _Chunk:
    return _Chunk(
        data=mistralai.ResponseStartedEvent.model_construct(
            conversation_id=conversation_id,
            type="conversation.response.started",
            created_at=0,
        )
    )


def _mock_task_cls():
    """Returns a class that mimics Task as a no-op async context manager."""

    class _FakeTask:
        def __class_getitem__(cls, item):
            return cls

        def __init__(self, **kwargs):
            pass

        async def __aenter__(self):
            mock = MagicMock()
            mock.set_state = AsyncMock()
            return mock

        async def __aexit__(self, *args):
            pass

    return _FakeTask


def _response_done_chunk() -> _Chunk:
    return _Chunk(
        data=mistralai.ResponseDoneEvent.model_construct(
            type="conversation.response.done",
            created_at=0,
            usage=None,
        )
    )


@pytest.mark.asyncio
class TestFunctionCallEventStreaming:
    """Tests that FunctionCallEvent arguments are correctly accumulated."""

    async def test_first_chunk_args_preserved(self) -> None:
        """First chunk carries partial args + subsequent chunks carry rest -> all concatenated."""
        chunks = [
            _response_started_chunk(),
            _function_call_chunk(name="my_tool", tool_call_id="tc-1", arguments='{"enabled'),
            _function_call_chunk(arguments='": true}'),
            _response_done_chunk(),
        ]
        stream = _FakeEventStream(chunks)

        with patch("mistralai_workflows.plugins.mistralai.utils.Task", _mock_task_cls()):
            from mistralai_workflows.plugins.mistralai.utils import handle_conversation_stream

            result = await handle_conversation_stream(stream)  # type: ignore[arg-type]

        fc_outputs = [o for o in (result.outputs or []) if isinstance(o, mistralai.FunctionCallEntry)]
        assert len(fc_outputs) == 1
        assert fc_outputs[0].arguments == '{"enabled": true}'

    async def test_single_chunk_full_args(self) -> None:
        """Single chunk with all args -> preserved in full."""
        chunks = [
            _response_started_chunk(),
            _function_call_chunk(name="my_tool", tool_call_id="tc-1", arguments='{"key": "value"}'),
            _response_done_chunk(),
        ]
        stream = _FakeEventStream(chunks)

        with patch("mistralai_workflows.plugins.mistralai.utils.Task", _mock_task_cls()):
            from mistralai_workflows.plugins.mistralai.utils import handle_conversation_stream

            result = await handle_conversation_stream(stream)  # type: ignore[arg-type]

        fc_outputs = [o for o in (result.outputs or []) if isinstance(o, mistralai.FunctionCallEntry)]
        assert len(fc_outputs) == 1
        assert fc_outputs[0].arguments == '{"key": "value"}'

    async def test_first_chunk_empty_args(self) -> None:
        """First chunk has empty args + subsequent chunks have args -> works (regression guard)."""
        chunks = [
            _response_started_chunk(),
            _function_call_chunk(name="my_tool", tool_call_id="tc-1", arguments=""),
            _function_call_chunk(arguments='{"enabled": true}'),
            _response_done_chunk(),
        ]
        stream = _FakeEventStream(chunks)

        with patch("mistralai_workflows.plugins.mistralai.utils.Task", _mock_task_cls()):
            from mistralai_workflows.plugins.mistralai.utils import handle_conversation_stream

            result = await handle_conversation_stream(stream)  # type: ignore[arg-type]

        fc_outputs = [o for o in (result.outputs or []) if isinstance(o, mistralai.FunctionCallEntry)]
        assert len(fc_outputs) == 1
        assert fc_outputs[0].arguments == '{"enabled": true}'


@pytest.mark.asyncio
class TestToolExecutionDeltaEventStreaming:
    """Tests that ToolExecutionDeltaEvent arguments are correctly accumulated."""

    async def test_first_chunk_args_preserved(self) -> None:
        """First chunk carries partial args + subsequent chunks carry rest -> all concatenated."""
        chunks = [
            _response_started_chunk(),
            _tool_execution_chunk(name="web_search", id="te-1", arguments='{"enabled'),
            _tool_execution_chunk(arguments='": true}'),
            _tool_execution_done_chunk(),
            _response_done_chunk(),
        ]
        stream = _FakeEventStream(chunks)

        with patch("mistralai_workflows.plugins.mistralai.utils.Task", _mock_task_cls()):
            from mistralai_workflows.plugins.mistralai.utils import handle_conversation_stream

            result = await handle_conversation_stream(stream)  # type: ignore[arg-type]

        te_outputs = [o for o in (result.outputs or []) if isinstance(o, mistralai.ToolExecutionEntry)]
        assert len(te_outputs) == 1
        assert te_outputs[0].arguments == '{"enabled": true}'

    async def test_single_chunk_full_args(self) -> None:
        """Single chunk with all args -> preserved in full."""
        chunks = [
            _response_started_chunk(),
            _tool_execution_chunk(name="web_search", id="te-1", arguments='{"key": "value"}'),
            _tool_execution_done_chunk(),
            _response_done_chunk(),
        ]
        stream = _FakeEventStream(chunks)

        with patch("mistralai_workflows.plugins.mistralai.utils.Task", _mock_task_cls()):
            from mistralai_workflows.plugins.mistralai.utils import handle_conversation_stream

            result = await handle_conversation_stream(stream)  # type: ignore[arg-type]

        te_outputs = [o for o in (result.outputs or []) if isinstance(o, mistralai.ToolExecutionEntry)]
        assert len(te_outputs) == 1
        assert te_outputs[0].arguments == '{"key": "value"}'

    async def test_first_chunk_empty_args(self) -> None:
        """First chunk has empty args + subsequent chunks have args -> works (regression guard)."""
        chunks = [
            _response_started_chunk(),
            _tool_execution_chunk(name="web_search", id="te-1", arguments=""),
            _tool_execution_chunk(arguments='{"enabled": true}'),
            _tool_execution_done_chunk(),
            _response_done_chunk(),
        ]
        stream = _FakeEventStream(chunks)

        with patch("mistralai_workflows.plugins.mistralai.utils.Task", _mock_task_cls()):
            from mistralai_workflows.plugins.mistralai.utils import handle_conversation_stream

            result = await handle_conversation_stream(stream)  # type: ignore[arg-type]

        te_outputs = [o for o in (result.outputs or []) if isinstance(o, mistralai.ToolExecutionEntry)]
        assert len(te_outputs) == 1
        assert te_outputs[0].arguments == '{"enabled": true}'
