from typing import Any, Coroutine

from fastapi import Request, Response, status
from fastapi.responses import JSONResponse
from .blackbox import Blackbox

from ..utils import chroma_setting
from .fastchat import Fastchat
from .chroma_query import ChromaQuery

DEFAULT_COLLECTION_ID = "123"

from injector import singleton,inject
@singleton
class ChromaChat(Blackbox):

    @inject
    def __init__(self, fastchat: Fastchat, chroma_query: ChromaQuery):
        self.fastchat = fastchat
        self.chroma_query = chroma_query

    def __call__(self, *args, **kwargs):
        return self.processing(*args, **kwargs)

    def valid(self, *args, **kwargs) -> bool:
        data = args[0]
        return isinstance(data, list)

    def processing(self, question, setting: chroma_setting) -> str:

        # load or create collection
        if setting is None:
            collection_id = DEFAULT_COLLECTION_ID
        else:
            collection_id = setting.ChromaSetting.collection_ids[0]
        # query it
        chroma_result = self.chroma_query(question, collection_id)

        fast_question = "问题: "+ question + "。根据问题，总结以下内容和来源：" + chroma_result
        response = self.fastchat(model_name="Qwen1.5-14B-Chat", prompt=fast_question, template='回答限制50字.', context=None, temperature=0.8, top_p=0.8, top_k=-1, n=1, max_tokens=1024)

        return response


    async def fast_api_handler(self, request: Request) -> Response:
        try:
            data = await request.json()
        except:
            return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)

        user_question = data.get("question")
        user_context = data.get("context")

        if user_question is None:
            return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)

        return JSONResponse(
            content={"response": self.processing(user_question, user_context)},
            status_code=status.HTTP_200_OK)