from .inference import GGMLModel
from ..user_interface import GradioUserInference
import logging
from typing import List, Literal, Any
from ..utils import BaseClassAgent, DEFAULT_SYSTEM_PROMPT, prompt_model

logging.basicConfig(
    level=logging.INFO
)


class PixelyAIServeGGML(GradioUserInference, BaseClassAgent):
    def __init__(
            self,
            model: GGMLModel,
    ):
        self.model = model
        self.gradio_app_custom = self.create_gradio_pixely_ai()

    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,
                prompting_style="Llama2",
        ):
            yield holder, response

    @classmethod
    def from_pretrained(
            cls,
            pretrained_model_name_or_path: str,
            quantize_type: Literal[
                "Q2_K",
                "Q3_K_S",
                "Q3_K_M",
                "Q3_K_L",
                "Q4_K_S",
                "Q4_K_M",
                "Q5_K_S",
                "Q5_K_M",
                "Q6_K"
            ] = "Q4_K_S"
    ):
        return cls(
            GGMLModel.from_pretrained(
                pretrained_model_name_or_path=pretrained_model_name_or_path,
                quantize_type=quantize_type
            )
        )

    def process(
            self,
            prompt: str,
            history: List[List[str]],
            system_prompt: str,
            prompting_style: str | Literal["Llama2", "OpenChat"],
            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 system_prompt: prompt for system 
        :param prompting_style: str: Prompting style to prompt Model
        :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
        """
        template = self.model.get_chat_template(
            prompting_style
        )
        if self.model is not None:
            string = template(
                prompt,
                history,
                None if system_prompt == "" else system_prompt
            )
            history.append([prompt, ""])
            total_response = ""
            for response in self.model(
                    string,
                    top_k=top_k,
                    top_p=top_p,
                    max_position_embedding=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 process_gradio_custom(self, prompt, user_id, data, system, max_new_tokens, greedy):
        history = data.split("<|END_OF_MESSAGE|>") if data != "" else []
        system = system if system != "" else DEFAULT_SYSTEM_PROMPT
        his = []
        for hs in history:
            if hs != "":
                his.append(hs.split("<|END_OF_MESSAGE_TURN_HUMAN|>"))
        history = his

        for holder, response in self.process(
                prompt=prompt,
                history=history,
                system_prompt=system,
                max_tokens=max_new_tokens,
                prompting_style="Llama2",
                top_p=0.95,
                top_k=50,
                temperature=0.8,
        ):
            yield holder, response[-1][-1]

    def launch(
            self,
            launch_custom: bool = True,
            launch_chat: bool = True,
            share: bool = True,
            *args,
            **kwargs
    ):
        urls = {}
        if launch_custom:
            self.gradio_app_custom.launch(share)
            urls["custom"] = self.gradio_app_custom.share_url
        if launch_chat:
            grad_chat = self.build_inference(
                self.process_gradio, 2048, 2048, 1
            )
            grad_chat.launch(share)
            urls["chat"] = grad_chat.share_url
        return urls

    @staticmethod
    def format_chat(history: List[List[str]], prompt: str, system: str = None) -> str:
        return prompt_model(message=prompt, system_prompt=system, chat_history=history)

    @staticmethod
    def prompt_model(message: str, chat_history, system_prompt: str):
        return prompt_model(message=message, chat_history=chat_history, system_prompt=system_prompt)
