import textwrap

import gradio

from source import meta
from source.api.interface import base
from source.model.base import BaseModel


class ChatInterface(base.BaseInterface):
    """
    An interface for Chat-like models.
    Use the OpenAI convention (list of dict with roles and content)
    """

    def __init__(self, model: "BaseModel"):
        # Function to send and receive chat messages
        super().__init__(model)

    async def send_message(self, user_message, old_messages: list[dict], system_message: str):
        # normalize the user message (the type can be wrong, especially when "edited")
        if isinstance(user_message, str):
            user_message: dict = {"files": [], "text": user_message}

        # copy the history to avoid modifying it
        messages: list[dict] = old_messages.copy()

        # add the system instruction
        if system_message:
            messages.insert(0, {"role": "system", "content": system_message})

        # add the user message
        # NOTE: gradio.ChatInterface add our message and the assistant message
        # TODO(Faraphel): add support for files
        messages.append({
            "role": "user",
            "content": user_message["text"],
        })

        # infer the message through the model
        chunks = [chunk async for chunk in await self.model.infer(messages=messages)]
        assistant_message: str = b"".join(chunks).decode("utf-8")

        # send back the messages, clear the user prompt, disable the system prompt
        return assistant_message

    def get_gradio_application(self):
        # create a gradio interface
        with gradio.Blocks(analytics_enabled=False) as application:
            # header
            gradio.Markdown(textwrap.dedent(f"""
                # {meta.name}
                ## {self.model.name}
            """))

            # additional settings
            with gradio.Accordion("Advanced Settings") as advanced_settings:
                system_prompt = gradio.Textbox(
                    label="System prompt",
                    placeholder="You are an expert in C++...",
                    lines=2,
                )

            # chat interface
            gradio.ChatInterface(
                fn=self.send_message,
                type="messages",
                multimodal=True,
                editable=True,
                save_history=True,
                additional_inputs=[system_prompt],
                additional_inputs_accordion=advanced_settings,
            )

        return application