import gradio as gr
from .theme import seafoam
from ..inference import (
    InferenceSession,
    download_model_gguf,
    AVAILABLE_MODELS,
    AVAILABLE_FORMATS,
    VERIFIED_PROVIDER,
    PROMPTING_STYLES,
    CHAT_MODE
)
from ._info import MARKDOWN_TEXT_INFO
from typing import Optional, List, Union
from llama_cpp import Llama
from gradio.helpers import log_message


class UserInference:
    """ 
    The `UserInference` class is a Python class that represents a user's interaction with a language
    model. It has an `__init__` method that initializes the class with an `inference_session` object,
    `max_tokens`, and `max_length` parameters. The `inference_session` object is used to perform
    inference with the language model. The `max_tokens` parameter sets the maximum number of tokens that
    can be used in a single query, and the `max_length` parameter sets the maximum length of a sentence.
    """

    def __init__(
            self,
            inference_session: InferenceSession,
            max_tokens: int,
            max_length: int,
    ):
        """
        The __init__ function is called when the class is instantiated.
        It sets up the inference session and defines some parameters for how long a sentence can be.

        :param self: Represent the instance of the class
        :param inference_session: InferenceSession: Store the inference session object
        :param max_tokens: int: Set the maximum number of tokens that can be used in a single query
        :param max_length: int: Set the maximum length of a sentence
        """
        self.inference_session = inference_session
        self.max_tokens = max_tokens
        self.max_length = max_length

    def sample(
            self,
            prompt: str,
            history: List[List[str]],
            system_prompt: str,
            prompting_style: str,
            mode: CHAT_MODE = CHAT_MODE[-1],
            max_tokens: int = 4096,
            temperature: float = 0.8,
            top_p: float = 0.9,
            top_k: int = 50
    ):
        """
        The sample function is the main entry point for a user to interact with the model.
        It takes in a prompt, which can be any string, and returns an iterator over
        strings that are generated by the model.
        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 prompting_style: str: Prompting style to prompt Model
        :param mode: str: represent the mode that model infrence be used in (e.g chat or instruction)
        :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
        """
        temeplate = self.inference_session.get_chat_template(
            prompting_style
        )
        assert mode in CHAT_MODE, "Requested Mode is not in Available Modes"
        if mode == "Instruction":
            history = []
        if self.inference_session is not None:

            string = temeplate(
                prompt,
                history,
                None if system_prompt == "" else system_prompt
            )
            history.append([prompt, ""])
            total_response = ""
            for response in self.inference_session(
                    string,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    temperature=temperature,
            ):
                total_response += response.predictions.text
                history[-1][-1] = total_response
                yield '', history
        else:
            return [
                [prompt, "Opps Seems like you forgot to load me first ;\\"]
            ]

    def _load_new_model(

        self,
        model_name,
        model_quantization_format,
        model_provider,
        hf_token,
        model_context_length_load,
        offload_qkv_button,
    ):
        """
        The function `_load_new_model` initializes a new inference session with a specified model and
        configuration parameters.

        :param model_name: The name of the model to be loaded
        :param model_quantization_format: The `model_quantization_format` parameter specifies the format
        in which the model is quantized. Quantization is a technique used to reduce the memory footprint
        and improve the inference speed of deep learning models. Common quantization formats include
        INT8, INT16, and FP16
        :param model_provider: The `model_provider` parameter is the provider of the model. It could be
        a specific organization or individual that has developed or provided the model for use
        :param model_context_length_load: The model_context_length_load parameter represents the number of tokens to load into
        the model's context. It determines the size of the input context for the model during inference
        :param offload_qkv_button: The `offload_qkv_button` parameter is a boolean value that determines
        whether the key, query, and value computations should be offloaded to a separate device or not.
        If `offload_qkv_button` is set to `True`, the computations will be offloaded. If it is set
        :param prompting_style: The `prompting_style` parameter is used to specify the type of prompting
        strategy to be used in the model. It is likely a string or an enum that determines how the model
        will generate responses based on the given input
        """
        self.inference_session = None
        log_message(
            f"Model is being Loaded", level="info"
        )
        self.inference_session = InferenceSession(
            model=Llama(
                model_path=download_model_gguf(
                    model_name,
                    model_quantization_format,
                    model_provider,
                    hf_token
                ),
                n_ctx=model_context_length_load,
                offload_kqv=offload_qkv_button,

            ),
            max_tokens=model_context_length_load
        )
        log_message(
            f"Model is Loaded", level="info"
        )

    def model_tab_interface_components(self):
        """
        The function `model_tab_interface_components` creates a user interface for selecting and loading
        a model with various parameters.
        """
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown(
                    MARKDOWN_TEXT_INFO
                )
            with gr.Column(scale=2):
                model_quantization_format = gr.Dropdown(
                    AVAILABLE_FORMATS,
                    value="Q4_K_S",
                    max_choices=1,
                    label="Model Quantization Format",
                )
                model_name = gr.Dropdown(
                    AVAILABLE_MODELS,
                    max_choices=1,
                    value="LinguaMatic-Tiny-GGUF",
                    label="Model ID",
                    allow_custom_value=True
                )

                model_provider = gr.Dropdown(
                    VERIFIED_PROVIDER,
                    value="erfanzar",
                    max_choices=1,
                    label="Model Provider",
                    allow_custom_value=True
                )
                hf_token = gr.Text(
                    label="HuggingFace Access Token",
                    placeholder="Leave Empty in case of using Public Models"
                )
                model_context_length_load = gr.Slider(
                    minimum=128,
                    maximum=32000,
                    value=2048,
                    step=1,
                    label="Model Context Length",
                )
                offload_qkv_button = gr.Checkbox(
                    value=False,
                    label="Offload QKV for Attention"
                )
                load_button = gr.Button(
                    value="Load Model",
                    variant="primary"
                )

        load_button.click(
            self._load_new_model,
            inputs=[
                model_name,
                model_quantization_format,
                model_provider,
                hf_token,
                model_context_length_load,
                offload_qkv_button,
            ]
        )

    def chat_interface_components(self):
        """
        The function `chat_interface_components` creates the components for a chat interface, including
        a chat history, message box, buttons for submitting, stopping, and clearing the conversation,
        and sliders for advanced options.
        """
        with gr.Row():
            with gr.Column(scale=4):
                history = gr.Chatbot(
                    elem_id="EasyDel",
                    label="EasyDel",
                    container=True,
                    height=800
                )
            with gr.Column(scale=1):
                gr.Markdown(
                    "# <h1><center>Powered by [EasyDeL](https://github.com/erfanzar/EasyDel)</center></h1>"
                )
                system_prompt = gr.Textbox(
                    value="",
                    show_label=True,
                    label="System Prompt",
                    placeholder='System Prompt',
                    container=False
                )
                max_tokens = gr.Slider(
                    value=self.max_tokens,
                    maximum=10000,
                    minimum=0,
                    label='Max Tokens',
                    step=1
                )
                temperature = gr.Slider(
                    value=0.8,
                    maximum=1,
                    minimum=0.1,
                    label='Temperature',
                    step=0.01
                )
                top_p = gr.Slider(
                    value=0.9,
                    maximum=1,
                    minimum=0.1,
                    label='Top P',
                    step=0.01
                )
                top_k = gr.Slider(
                    value=50,
                    maximum=100,
                    minimum=1,
                    label='Top K',
                    step=1
                )
                prompting_style = gr.Dropdown(
                    choices=PROMPTING_STYLES,
                    value=PROMPTING_STYLES[0],
                    label="Prompting Style",
                    max_choices=1,
                )
                mode = gr.Dropdown(
                    choices=CHAT_MODE,
                    value=CHAT_MODE[1],
                    label="Mode",
                    max_choices=1,
                )
                submit = gr.Button(
                    value="Run",
                    variant="primary"
                )
                stop = gr.Button(
                    value='Stop'
                )
                clear = gr.Button(
                    value='Clear Conversation'
                )
        with gr.Row():
            with gr.Column(scale=4):
                prompt = gr.Textbox(
                    show_label=False, placeholder='Message Box', container=False)
            with gr.Column(scale=1):
                ...
        inputs = [
            prompt,
            history,
            system_prompt,
            prompting_style,
            mode,
            max_tokens,
            temperature,
            top_p,
            top_k
        ]

        clear.click(fn=lambda: [], outputs=[history])
        sub_event = submit.click(
            fn=self.sample, inputs=inputs, outputs=[prompt, history])
        txt_event = prompt.submit(
            fn=self.sample, inputs=inputs, outputs=[prompt, history])

        stop.click(fn=None, inputs=None, outputs=None,
                   cancels=[txt_event, sub_event])

    def build_model_tab_interface(self) -> gr.Blocks:
        """
        The function "build_model_tab_interface" returns a gr.Blocks object that contains the components
        of a model tab interface.
        :return: a gr.Blocks object.
        """

        with gr.Blocks(
                theme=seafoam
        ) as block:
            self.model_tab_interface_components()
        return block

    def build_chat_interface(self) -> gr.Blocks:
        """
        The build function is the main entry point for your app.
        It should return a single gr.Block instance that will be displayed in the browser.

        :param self: Make the class methods work with an instance of the class
        :return: A block, which is then queued
        """
        with gr.Blocks(
                theme=seafoam
        ) as block:
            self.chat_interface_components()
        block.queue()
        return block

    def build_inference(self) -> gr.Blocks:
        """
        The function "build_inference" returns a gr.Blocks object that contains two tabs, one for model
        interface components and one for chat interface components.
        :return: a gr.Blocks object.
        """
        with gr.Blocks(
                theme=seafoam
        ) as block:
            with gr.Tab("Model",):
                self.model_tab_interface_components()
            with gr.Tab("Chat"):
                self.chat_interface_components()
        return block
