from typing import Literal
from chatAPI import chat
from typing import Union

class ChatbotBase:
    def __init__(self,system_prompt) -> None:
        self.history=[]
        self.system_prompt=''
        self.set_system_prompt(system_prompt)

    def set_system_prompt(self,prompt):
        '''set system prompt. MUST BE used after reset_bot()'''
        self.system_prompt=prompt
        self.add_message(prompt,image=None,role='system')
    
    def reset_bot(self):
        self.history=[]
        self.system_prompt=''

    def add_message(self,message,image:Union[str,list[str]],role:Literal['user','assistant','system']):
        '''add new message to history
        image: base64 string of image
        '''
        if isinstance(image,list):
            content=[]
            content.append({'type':'text','text':message})
            for img in image:
                content.append({
                    'type':'image_url',
                    "image_url": {
                     "url": f"data:image/jpeg;base64,{image}",
                     "detail": "high"
                    }})
            
            self.history.append({
                'role':role,
                'content':content})
        elif image:
            self.history.append({
            'role':role,
            'content':[
                {'type':'text','text':message},
                {'type':'image_url',"image_url": {
              "url": f"data:image/jpeg;base64,{image}",
              "detail": "high"
               }}
            ]})
        else:
            self.history.append({
            'role':role,
            'content':message})

    def pop_history(self):
        '''remove last history message'''
        self.history.pop()

    def get_history(self):
        return self.history
    
    def generate(self):
        '''generate new response and store to history'''
        print('code generating')
        rsp=chat(self.history)
        self.add_message(rsp,image=None,role='assistant')
        print(self.history)
        return rsp