import json

from .base import Base


class Session(Base):
    def __init__(self, rag, res_dict):
        self.id = None
        self.name = "New session"
        self.messages = [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
        self.assistant_id = None
        super().__init__(rag, res_dict)

    def chat(self, question: str, stream: bool = False):
        for message in self.messages:
            if "reference" in message:
                message.pop("reference")
        res = self.post("/session/completion",
                        {"session_id": self.id, "question": question, "stream": True}, stream=stream)
        for line in res.iter_lines():
            line = line.decode("utf-8")
            if line.startswith("{"):
                json_data = json.loads(line)
                raise Exception(json_data["retmsg"])
            if line.startswith("data:"):
                json_data = json.loads(line[5:])
                if json_data["data"] != True:
                    answer = json_data["data"]["answer"]
                    reference = json_data["data"]["reference"]
                    temp_dict = {
                        "content": answer,
                        "role": "assistant"
                    }
                    if "chunks" in reference:
                        chunks = reference["chunks"]
                        chunk_list = []
                        for chunk in chunks:
                            new_chunk = {
                                "id": chunk["chunk_id"],
                                "content": chunk["content_with_weight"],
                                "document_id": chunk["doc_id"],
                                "document_name": chunk["docnm_kwd"],
                                "knowledgebase_id": chunk["kb_id"],
                                "image_id": chunk["img_id"],
                                "similarity": chunk["similarity"],
                                "vector_similarity": chunk["vector_similarity"],
                                "term_similarity": chunk["term_similarity"],
                                "positions": chunk["positions"],
                            }
                            chunk_list.append(new_chunk)
                        temp_dict["reference"] = chunk_list
                    message = Message(self.rag, temp_dict)
                    yield message

    def save(self):
        res = self.post("/session/save",
                        {"id": self.id, "assistant_id": self.assistant_id, "name": self.name})
        res = res.json()
        if res.get("retmsg") == "success": return True
        raise Exception(res.get("retmsg"))

    def delete(self):
        res = self.rm("/session/delete", {"id": self.id})
        res = res.json()
        if res.get("retmsg") == "success": return True
        raise Exception(res.get("retmsg"))


class Message(Base):
    def __init__(self, rag, res_dict):
        self.content = "Hi! I am your assistant，can I help you?"
        self.reference = None
        self.role = "assistant"
        self.prompt = None
        self.id = None
        super().__init__(rag, res_dict)


class Chunk(Base):
    def __init__(self, rag, res_dict):
        self.id = None
        self.content = None
        self.document_id = ""
        self.document_name = ""
        self.knowledgebase_id = ""
        self.image_id = ""
        self.similarity = None
        self.vector_similarity = None
        self.term_similarity = None
        self.positions = None
        super().__init__(rag, res_dict)