import gradio as gr

from ..user_interface import GradioUserInference
from agentx import ServeEngine
import logging
from typing import List, Optional
import transformers

logging.basicConfig(
    level=logging.INFO
)


class AgentXServer(GradioUserInference):
    tokenizer: transformers.PreTrainedTokenizerBase | None = None

    def __init__(
            self,
            engine: ServeEngine,
    ):
        self.engine = engine
        self.gradio_pixely_ai = self.create_gradio_pixely_ai()

    def create_gradio_pixely_ai(self):
        with gr.Blocks(theme=gr.themes.Soft()) as block:
            gr.Markdown("# <h1> <center>Powered by [EasyDeL](https://github.com/erfanzar/EasyDel) </center> </h1>")

            with gr.Row():
                with gr.Column():
                    prompt = gr.Textbox(
                        show_label=True,
                        placeholder="Message Box",
                        container=True,
                        label="Message Box"
                    )
                    user_id = gr.Textbox(
                        show_label=True,
                        placeholder="UserId",
                        container=True,
                        value="",
                        label="UserId"
                    )
                    data = gr.Textbox(
                        show_label=True,
                        placeholder="Data",
                        container=True,
                        value="",
                        label="Data"
                    )
                    system = gr.Textbox(
                        show_label=True,
                        placeholder="System",
                        container=True,
                        value="",
                        label="System"
                    )
                    response = gr.TextArea(
                        show_label=True,
                        placeholder="Response",
                        container=True,
                        label="Response"
                    )
                    submit = gr.Button(variant="primary")
            with gr.Row():
                with gr.Accordion("Advanced Options", open=False):
                    max_new_tokens = gr.Slider(
                        value=self.engine.sample_config.max_new_tokens,
                        maximum=10000,
                        minimum=1,
                        label="Max New Tokens",
                        step=1
                    )

                    greedy = gr.Checkbox(value=False, label="Greedy Search")

            inputs = [
                prompt,
                user_id,
                data,
                system,
                max_new_tokens,
                greedy
            ]
            _ = prompt.submit(
                fn=self.process_gradio_custom,
                inputs=inputs,
                outputs=[prompt, response]
            )
            _ = submit.click(
                fn=self.process_gradio_custom,
                inputs=inputs,
                outputs=[prompt, response]
            )

            block.queue()
        return block

    def process_gradio(
            self,
            prompt: str,
            history: List[List[str]],
            system_prompt: str | None,
            mode: str,
            max_length: int,
            max_new_tokens: int,
            max_compile_tokens: int,
            greedy: bool,
            temperature: float,
            top_p: float,
            top_k: int
    ):
        for holder, response in self.process(
                prompt=prompt,
                history=history,
                system_prompt=system_prompt,
                max_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
        ):
            yield holder, response

    def process(
            self,
            prompt: str,
            history: List[List[str]],
            system_prompt: str = "",
            max_tokens: Optional[int] = None,
            temperature: Optional[float] = None,
            top_p: Optional[float] = None,
            top_k: Optional[int] = None
    ):
        """
        The sample function is the main entry point for a user to interact with the engine.
        It takes in a prompt, which can be any string, and returns an iterator over
        strings that are generated by the engine.
        The sample function also takes in some optional arguments:


        :param self: Refer to the current object
        :param prompt: str: Pass in the text that you want to generate a response for
        :param history: List[List[str]]: Keep track of the conversation history
        :param system_prompt: prompt for system
        :param max_tokens: int: Limit the number of tokens in a response
        :param temperature: float: Control the randomness of the generated text
        :param top_p: float: Control the probability of sampling from the top k tokens
        :param top_k: int: Control the number of candidates that are considered for each token
        :return: A generator that yields the next token in the sequence
        """

        if self.engine is not None:

            conversation = [{"role": "system", "content": system_prompt}] if system_prompt != "" else []
            for perv in history:
                conversation.append({"role": "user", "content": perv[0]})
                conversation.append({"role": "assistant", "content": perv[1]})
            conversation.append({"role": "user", "content": prompt})

            string = self.engine.prompt_template.render(
                conversation
            )
            history.append([prompt, ""])
            total_response = ""
            for response in self.engine.process(
                    string,
                    top_k=top_k,
                    top_p=top_p,
                    max_sequence_length=max_tokens,
                    temperature=temperature,
                    max_new_tokens=max_tokens
            ):
                total_response += response
                history[-1][-1] = total_response
                yield "", history
        else:
            return [
                [prompt, "Opps Seems like you forgot to load me first ;\\"]
            ]

    def process_gradio_custom(
            self,
            prompt: str,
            user_id: int,
            data: str,
            system: str,
            max_new_tokens: int,
            greedy: bool
    ):
        history = data.split("<|END_OF_MESSAGE|>") if data != "" else []
        conversations = []
        for conversation in history:
            if conversation != "":
                conversations.append(conversation.split("<|END_OF_MESSAGE_TURN_HUMAN|>"))

        for holder, response in self.process(
                prompt=prompt,
                history=conversations,
                system_prompt=system,
                max_tokens=max_new_tokens,
        ):
            yield holder, response[-1][-1]

    def build(
            self,
    ):
        with gr.Blocks(
                theme=gr.themes.Soft()
        ) as block:
            with gr.Tab("Pixely Serve"):
                self.create_gradio_pixely_ai()
            with gr.Tab("Pixely Chat"):
                self.build_inference(
                    self.process_gradio,
                    self.engine.sample_config.max_sequence_length,
                    self.engine.sample_config.max_new_tokens,
                    1
                )
        return block
