from logging import getLogger
from unittest import IsolatedAsyncioTestCase
from unittest.mock import MagicMock

from avalan.entities import (
    GenerationSettings,
    ReasoningSettings,
    ReasoningToken,
    Token,
    TokenDetail,
    ToolCallToken,
)
from avalan.model.response.parsers.reasoning import ReasoningParser
from avalan.model.response.parsers.tool import ToolCallResponseParser
from avalan.model.response.text import TextGenerationResponse
from avalan.tool.parser import ToolCallParser


async def _complex_generator():
    rp = ReasoningParser(
        reasoning_settings=ReasoningSettings(), logger=getLogger()
    )
    tm = MagicMock()
    tm.is_potential_tool_call.return_value = True
    tm.get_calls.return_value = None
    base_parser = ToolCallParser()
    tm.tool_call_status.side_effect = base_parser.tool_call_status
    tp = ToolCallResponseParser(tm, None)

    sequence = [
        "X",
        "<think>",
        "ra",
        "rb",
        "</think>",
        "Y",
        "<tool_call>",
        "foo",
        "bar",
        "</tool_call>",
        "Z",
    ]

    for s in sequence:
        items = await rp.push(s)
        for item in items:
            parsed = await tp.push(item) if isinstance(item, str) else [item]
            for p in parsed:
                if isinstance(p, str):
                    if p == "</think>":
                        yield TokenDetail(id=3, token=p, probability=0.5)
                    elif p in {"X", "Y"}:
                        yield Token(id=1, token=p)
                    else:
                        yield p
                elif isinstance(p, ToolCallToken):
                    if p.token == "</tool_call>":
                        yield TokenDetail(id=4, token=p.token, probability=0.5)
                    else:
                        yield p
                else:
                    yield p


class TextGenerationResponseParsersTestCase(IsolatedAsyncioTestCase):
    async def test_mixed_tokens(self):
        settings = GenerationSettings()
        resp = TextGenerationResponse(
            lambda **_: _complex_generator(),
            logger=getLogger(),
            use_async_generator=True,
            generation_settings=settings,
            settings=settings,
        )

        tokens = []
        async for t in resp:
            tokens.append(t)

        self.assertEqual(
            len([t for t in tokens if isinstance(t, ReasoningToken)]),
            4,
        )
        self.assertEqual(
            len([t for t in tokens if isinstance(t, ToolCallToken)]),
            3,
        )
        self.assertEqual(
            len([t for t in tokens if isinstance(t, TokenDetail)]),
            1,
        )
        self.assertGreaterEqual(
            len([t for t in tokens if type(t) is Token]),
            2,
        )
        self.assertEqual(len([t for t in tokens if isinstance(t, str)]), 1)
