import asyncio
from collections.abc import AsyncIterator

from pydantic_ai import Agent

from tests.fixtures.agents.prompt import SYSTEM_PROMPT
from tests.fixtures.datasource import datasource


class PydanticAgent:
    agent: Agent

    def __init__(self):
        self.agent = Agent(
            model="openai:gpt-4o-mini",
            system_prompt=SYSTEM_PROMPT,
            tools=datasource,
        )

    async def stream(self, question: str) -> AsyncIterator[str]:
        """
        Stream responses from Pydantic AI with tool call support.

        Args:
            question: The question to ask

        Yields:
            Text chunks from the streaming response (tool calls are handled automatically)
        """
        from pydantic_ai.messages import PartDeltaEvent

        async with self.agent.run_stream(question) as stream_result:
            async for event in stream_result.stream_output():
                if isinstance(event, PartDeltaEvent):
                    if hasattr(event, "delta") and hasattr(event.delta, "content"):
                        content = event.delta.content
                        if content:
                            yield str(content)
                elif isinstance(event, str):
                    yield event

    def ask(self, question: str) -> str:
        response = asyncio.run(self.agent.run(question))
        return response.output
