from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from django_anyllm.agent import ModelAgent
from django_anyllm.registry import ModelRegistry
from django_anyllm.stream_types import StreamEventType

from .models import Author


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_choice(message):
    choice = MagicMock()
    choice.message = message
    return choice


def _make_message(content: str = "", tool_calls=None, reasoning=None):
    msg = MagicMock()
    msg.role = "assistant"
    msg.content = content
    msg.tool_calls = tool_calls or []
    msg.reasoning = reasoning
    return msg


def _make_tool_call(id_: str, name: str, arguments: str):
    tc = MagicMock()
    tc.id = id_
    tc.type = "function"
    tc.function = MagicMock()
    tc.function.name = name
    tc.function.arguments = arguments
    return tc


def _make_response(content: str = "", tool_calls=None):
    resp = MagicMock()
    resp.choices = [_make_choice(_make_message(content, tool_calls))]
    return resp


def _make_chunk(content: str | None = None, reasoning=None, tool_calls=None):
    chunk = MagicMock()
    delta = MagicMock()
    delta.content = content
    delta.reasoning = reasoning
    delta.tool_calls = tool_calls
    choice = MagicMock()
    choice.delta = delta
    chunk.choices = [choice]
    return chunk


def _make_delta_tool_call(
    index: int, name: str = "", arguments: str = "", id_: str = ""
):
    tc = MagicMock()
    tc.index = index
    tc.id = id_
    tc.function = MagicMock()
    tc.function.name = name
    tc.function.arguments = arguments
    return tc


@pytest.fixture
def agent(db):
    registry = ModelRegistry()
    registry.register(Author)
    return ModelAgent(
        model="gpt-4",
        provider="openai",
        registry=registry,
        api_key="test-key",
    )


# ---------------------------------------------------------------------------
# Sync non-streaming
# ---------------------------------------------------------------------------


class TestChatSync:
    """Tests for ModelAgent.chat() non-streaming."""

    def test_simple_response(self, agent):
        with patch(
            "django_anyllm.agent.completion",
            return_value=_make_response("Hello!"),
        ) as mock_completion:
            result = agent.chat("Hi")

        assert result == "Hello!"
        mock_completion.assert_called_once()
        call = mock_completion.call_args
        assert call.kwargs["model"] == "openai:gpt-4"
        assert call.kwargs.get("provider") is None
        assert call.kwargs["tools"] == agent._tools

    def test_tool_call_loop(self, agent, db):
        author = Author.objects.create(name="Alice", email="alice@example.com")

        tool_call = _make_tool_call(
            "tc_1", "_query_instances", '{"app_label": "tests", "model_name": "author"}'
        )

        responses = [
            _make_response("", tool_calls=[tool_call]),
            _make_response(f"Found author: Alice (id={author.pk})"),
        ]

        with patch(
            "django_anyllm.agent.completion",
            side_effect=responses,
        ) as mock_completion:
            result = agent.chat("List authors")

        assert "Alice" in result
        assert mock_completion.call_count == 2

    def test_max_tool_rounds(self, agent):
        tool_call = _make_tool_call("tc_1", "_query_instances", "{}")
        responses = [_make_response("", tool_calls=[tool_call])] * 10

        with patch(
            "django_anyllm.agent.completion",
            side_effect=responses,
        ):
            result = agent.chat("Loop")

        # Should stop after max_tool_rounds (default 5).
        assert result == ""

    def test_reasoning_effort_forwarded(self, agent):
        with patch(
            "django_anyllm.agent.completion",
            return_value=_make_response("OK"),
        ) as mock_completion:
            agent.chat("Think hard", reasoning_effort="high")

        assert mock_completion.call_args.kwargs["reasoning_effort"] == "high"


# ---------------------------------------------------------------------------
# Sync streaming
# ---------------------------------------------------------------------------


class TestChatStream:
    """Tests for ModelAgent.chat(stream=True)."""

    def test_content_streaming(self, agent):
        chunks = [
            _make_chunk("Hello"),
            _make_chunk(", "),
            _make_chunk("world!"),
        ]

        with patch(
            "django_anyllm.agent.completion",
            return_value=iter(chunks),
        ):
            events = list(agent.chat("Hi", stream=True))

        content_events = [e for e in events if e.type == StreamEventType.CONTENT]
        texts = [e.text for e in content_events]
        assert texts == ["Hello", ", ", "world!"]
        assert events[-1].type == StreamEventType.DONE

    def test_reasoning_streaming(self, agent):
        reasoning = MagicMock()
        reasoning.content = "Let me think..."
        chunks = [
            _make_chunk("OK", reasoning=reasoning),
        ]

        with patch(
            "django_anyllm.agent.completion",
            return_value=iter(chunks),
        ):
            events = list(agent.chat("Hi", stream=True))

        reasoning_events = [e for e in events if e.type == StreamEventType.REASONING]
        assert len(reasoning_events) == 1
        assert reasoning_events[0].text == "Let me think..."

    def test_tool_call_streaming(self, agent, db):
        # First round: assistant requests a tool call via streaming deltas.
        round1 = [
            _make_chunk("Let me check"),
            _make_chunk(
                None,
                tool_calls=[
                    _make_delta_tool_call(0, id_="tc_1"),
                ],
            ),
            _make_chunk(
                None,
                tool_calls=[
                    _make_delta_tool_call(0, name="_query_instances"),
                ],
            ),
            _make_chunk(
                None,
                tool_calls=[
                    _make_delta_tool_call(
                        0, arguments='{"app_label": "tests", "model_name": "author"}'
                    ),
                ],
            ),
        ]
        # Second round: assistant provides final answer.
        round2 = [
            _make_chunk("Done"),
        ]

        with patch(
            "django_anyllm.agent.completion",
            side_effect=[iter(round1), iter(round2)],
        ):
            events = list(agent.chat("List authors", stream=True))

        types = [e.type for e in events]
        assert StreamEventType.CONTENT in types
        assert StreamEventType.TOOL_CALL in types
        assert StreamEventType.TOOL_RESULT in types
        assert StreamEventType.DONE in types

        tool_call_event = next(e for e in events if e.type == StreamEventType.TOOL_CALL)
        assert tool_call_event.tool_name == "_query_instances"


# ---------------------------------------------------------------------------
# Async
# ---------------------------------------------------------------------------


class TestAChatSync:
    """Tests for ModelAgent.achat() non-streaming."""

    @pytest.mark.asyncio
    @pytest.mark.django_db(transaction=True)
    async def test_simple_response(self, agent):
        with patch(
            "django_anyllm.agent.acompletion",
            new_callable=AsyncMock,
            return_value=_make_response("Hello!"),
        ) as mock_completion:
            result = await agent.achat("Hi")

        assert result == "Hello!"
        mock_completion.assert_awaited_once()

    @pytest.mark.asyncio
    @pytest.mark.django_db(transaction=True)
    async def test_tool_call_loop(self, agent, db):
        from asgiref.sync import sync_to_async

        await sync_to_async(Author.objects.create)(
            name="Alice", email="alice@example.com"
        )
        tool_call = _make_tool_call(
            "tc_1", "_query_instances", '{"app_label": "tests", "model_name": "author"}'
        )
        responses = [
            _make_response("", tool_calls=[tool_call]),
            _make_response("Found Alice"),
        ]

        with patch(
            "django_anyllm.agent.acompletion",
            new_callable=AsyncMock,
            side_effect=responses,
        ):
            result = await agent.achat("List authors")

        assert "Alice" in result


class TestAChatStream:
    """Tests for ModelAgent.achat(stream=True)."""

    @pytest.mark.asyncio
    async def test_content_streaming(self, agent):
        chunks = [_make_chunk("Hello"), _make_chunk("!")]

        async def async_iter():
            for c in chunks:
                yield c

        with patch(
            "django_anyllm.agent.acompletion",
            new_callable=AsyncMock,
            return_value=async_iter(),
        ):
            events = [e async for e in await agent.achat("Hi", stream=True)]

        content_events = [e for e in events if e.type == StreamEventType.CONTENT]
        assert [e.text for e in content_events] == ["Hello", "!"]
        assert events[-1].type == StreamEventType.DONE


# ---------------------------------------------------------------------------
# PersistentModelAgent
# ---------------------------------------------------------------------------


class TestPersistentModelAgent:
    """Tests for PersistentModelAgent."""

    @pytest.fixture
    def persistent_agent(self, db):
        from django_anyllm.models import Chat
        from django_anyllm.agent import PersistentModelAgent

        chat = Chat.objects.create(model="gpt-4", provider="openai")
        registry = ModelRegistry()
        registry.register(Author)
        return PersistentModelAgent(chat=chat, registry=registry, api_key="test-key")

    def test_skips_empty_assistant_messages(self, db):
        from django_anyllm.models import Chat
        from django_anyllm.agent import PersistentModelAgent

        chat = Chat.objects.create(model="gpt-4", provider="openai")
        chat.messages.create(role="user", content="Hello")
        chat.messages.create(role="assistant", content="")
        chat.messages.create(role="user", content="Again")

        agent = PersistentModelAgent(chat=chat)
        messages = agent._build_messages()

        roles = [m["role"] for m in messages]
        assert roles == ["system", "user", "user"]

    def test_keeps_assistant_with_tool_calls(self, db):
        from django_anyllm.models import Chat
        from django_anyllm.agent import PersistentModelAgent

        chat = Chat.objects.create(model="gpt-4", provider="openai")
        chat.messages.create(role="user", content="Hello")
        chat.messages.create(
            role="assistant",
            content="",
            tool_calls=[{"id": "tc_1", "type": "function", "function": {"name": "foo", "arguments": "{}"}}],
        )
        chat.messages.create(role="user", content="Again")

        agent = PersistentModelAgent(chat=chat)
        messages = agent._build_messages()

        roles = [m["role"] for m in messages]
        assert roles == ["system", "user", "assistant", "user"]

    def test_skips_assistant_with_reasoning_but_no_content(self, db):
        from django_anyllm.models import Chat
        from django_anyllm.agent import PersistentModelAgent

        chat = Chat.objects.create(model="gpt-4", provider="openai")
        chat.messages.create(role="user", content="Hello")
        chat.messages.create(role="assistant", content="", reasoning="Let me think...")
        chat.messages.create(role="user", content="Again")

        agent = PersistentModelAgent(chat=chat)
        messages = agent._build_messages()

        roles = [m["role"] for m in messages]
        # Assistant with reasoning but no content is filtered out because
        # most providers (including OpenAI) reject it.
        assert roles == ["system", "user", "user"]

    def test_chat_sync_saves_assistant_message(self, persistent_agent, db):
        with patch(
            "django_anyllm.agent.completion",
            return_value=_make_response("Hello!"),
        ):
            result = persistent_agent.chat("Hi")

        assert result == "Hello!"
        assert persistent_agent._chat.messages.filter(role="assistant", content="Hello!").exists()

    def test_chat_stream_saves_content(self, persistent_agent, db):
        chunks = [
            _make_chunk("Hello"),
            _make_chunk(" world"),
        ]

        with patch(
            "django_anyllm.agent.completion",
            return_value=iter(chunks),
        ):
            events = list(persistent_agent.chat("Hi", stream=True))

        content_events = [e for e in events if e.type == StreamEventType.CONTENT]
        assert [e.text for e in content_events] == ["Hello", " world"]

        assistant_msg = persistent_agent._chat.messages.filter(role="assistant").first()
        assert assistant_msg is not None
        assert assistant_msg.content == "Hello world"
