# Experimental

from typing import Dict, Any, Optional, List
import asyncio
import logging
from uuid import uuid4
from anyio import ClosedResourceError
from anyio.streams.memory import MemoryObjectSendStream

from llama_index.core.callbacks.base import BaseCallbackHandler, CallbackManager
from llama_index.core.callbacks import CBEventType, EventPayload
from llama_index.core.query_engine.sub_question_query_engine import (
    SubQuestionAnswerPair,
)
from llama_index.core.chat_engine.types import StreamingAgentChatResponse
from pydantic import BaseModel

from core.chat import schema

from db.db import MessageSubProcessSourceEnum
from core.chat.schema import SubProcessMetadataKeysEnum, SubProcessMetadataMap
from core.chat.engine import Engine
from script.vector_db import IndexManager
from service.dto import UserPromptRequest

logger = logging.getLogger(__name__)

class StreamedMessage(BaseModel):
    content: str

async def handle_chat_message(
    user_message: str,
    send_chan: MemoryObjectSendStream,
) -> None:
    async with send_chan:
        engine = Engine()
        
        index_manager = IndexManager()
        index = index_manager.load_existing_indexes()

        # Retrieve the chat engine with the loaded index
        chat_engine = await engine.get_chat_engine(index)

        logger.debug("Engine received")
        streaming_chat_response: StreamingAgentChatResponse = (
            await chat_engine.astream_chat(user_message)
        )
        response_str = ""
        async for text in streaming_chat_response.async_response_gen():
            response_str += text
            if send_chan._closed:
                logger.debug(
                    "Received streamed token after send channel closed. Ignoring."
                )
                return
            await send_chan.send(StreamedMessage(content=response_str))

        if response_str.strip() == "":
            await send_chan.send(
                StreamedMessage(
                    content="Sorry, I either wasn't able to understand your question or I don't have an answer for it."
                )
            )