import json
from enum import StrEnum
from typing import TypeAlias
from uuid import UUID

from pydantic import BaseModel

import api.context as ctx
from api.agent import Agent
from banortegpt.database.mongo_memory import crud


class ChunkType(StrEnum):
    START = "start"
    TEXT = "text"
    REFERENCE = "reference"
    IMAGE = "image"
    TOOL = "tool"
    END = "end"
    ERROR = "error"


ContentType: TypeAlias = str | int


class ResponseChunk(BaseModel):
    type: ChunkType
    content: ContentType | list[ContentType] | None


async def stream(agent: Agent, prompt: str, conversation_id: UUID, with_deep_research: bool = False):
    yield ResponseChunk(type=ChunkType.START, content="")

    conversation = await crud.get_conversation(conversation_id)

    if conversation is None:
        raise ValueError("Conversation not found")

    conversation.add(role="user", content=prompt)

    history = conversation.to_openai_format(agent.message_limit, langchain_compat=True)
    async for content in agent.stream(history, with_deep_research):
        yield ResponseChunk(type=ChunkType.TEXT, content=content)

    if (tool_id := ctx.tool_id.get()) is not None:
        tool_buffer = ctx.tool_buffer.get()
        assert tool_buffer is not None

        tool_name = ctx.tool_name.get()
        assert tool_name is not None

        yield ResponseChunk(type=ChunkType.TOOL, content=None)

        buffer_dict = json.loads(tool_buffer)

        result = await agent.tool_map[tool_name](**buffer_dict)

        conversation.add(
            role="assistant",
            tool_calls=[
                {
                    "id": tool_id,
                    "type": "function",
                    "function": {
                        "name": tool_name,
                        "arguments": tool_buffer,
                    },
                }
            ],
        )
        conversation.add(role="tool", content=result, tool_call_id=tool_id)

        history = conversation.to_openai_format(agent.message_limit, langchain_compat=True)
        async for content in agent.stream(history, with_deep_research, {"tools": None}):
            yield ResponseChunk(type=ChunkType.TEXT, content=content)

    conversation.add(role="assistant", content=ctx.buffer.get())

    await conversation.replace()

    yield ResponseChunk(type=ChunkType.END, content="")