from models import BaseModel
from .base_refiner import BaseRefiner
from utils.image_encoder import encode_image
import asyncio

class SimpleRefiner(BaseRefiner):
    def __init__(self,
                 sys_prompt: str, 
                 model: BaseModel,
        ) -> None:
        BaseRefiner.__init__(self, sys_prompt=sys_prompt, model=model)
    
    async def refine_async(self, message: str, memory, image_paths=None) -> str:
        if memory is None:
            memory = []
        else:
            memory = memory.messages[1:]
        
        user_context = [{"role": "user", "content": [
                        {"type": "text", "text": f"{message}"},]}]
        if image_paths:
            if not isinstance(image_paths, list):
                image_paths = [image_paths]
            for image_path in image_paths:
                user_context.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}})
            context = [{"role": "system", "content": self.sys_prompt}] + memory + [{"role": "user", "content": [
                            {"type": "text", "text": f"{message}"},
                        ]}]
        else:
            context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context

        respond_task = asyncio.create_task(self.model.respond_async(context))
        await respond_task
        response = respond_task.result()
        return response
    
    def refine(self, message: str, memory, image_paths=None) -> str:
        if memory is None:
            memory = []
        else:
            memory = memory.messages[1:]
        
        user_context = [{"role": "user", "content": [
                        {"type": "text", "text": f"{message}"},]}]
        if image_paths:
            if not isinstance(image_paths, list):
                image_paths = [image_paths]
            for image_path in image_paths:
                user_context.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}})
            context = [{"role": "system", "content": self.sys_prompt}] + memory + [{"role": "user", "content": [
                            {"type": "text", "text": f"{message}"},
                        ]}]
        else:
            context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context
        response = self.model.respond(context)

        return response