import asyncio
import time
from .prompts import *
from .langchain_memory import BotConversationSummerBufferWindowMemory

from langchain import PromptTemplate
from langchain.chains import LLMChain


import logging

logger = logging.getLogger(__name__)


class AI(object):

    def __init__(self, bot, text_wrapper, image_wrapper):
        self.name = bot.name
        self.bot = bot

        from ..wrappers.langchain_koboldcpp import KoboldCpp
        self.llm_chat = KoboldCpp(temperature=self.bot.temperature, endpoint_url="http://172.16.85.10:5001/api/latest/generate", stop=['<|endoftext|>'])
        self.llm_summary = KoboldCpp(temperature=0.2, endpoint_url="http://172.16.85.10:5001/api/latest/generate", stop=['<|endoftext|>'])
        self.text_wrapper = text_wrapper
        self.image_wrapper = image_wrapper

        self.memory = BotConversationSummerBufferWindowMemory(llm=self.llm_summary, max_token_limit=1200, min_token_limit=200)




    async def generate(self, input_text):
        prompt_template = "{input}"
        chain = LLMChain(
            llm=self.llm_chat,
            prompt=PromptTemplate.from_template(prompt_template),
        )
        output = chain.run(input_text)
        return output.strip()

    async def generate_roleplay(self, message, reply_fn, typing_fn):
        prompt = PromptTemplate(
            input_variables=["ai_name", "persona", "scenario", "chat_history", "human_name", "ai_name_chat", "human_input"],
            template=prompt_template_alpaca,
        )
        template_roleplay = prompt.format(
            ai_name = self.bot.name,
            persona = self.bot.persona,
            scenario = self.bot.scenario,
            chat_history = "{history}",
            human_name = message.user_name,
            ai_name_chat = self.bot.name,
            human_input = "{input}",
        )
        chain = LLMChain(
            llm=self.llm_chat,
            prompt=PromptTemplate.from_template(template_roleplay),
            verbose=True,
            memory=self.memory,
            #stop=['<|endoftext|>', '\nYou:', f"\n{message.user_name}:"],
        )
        output = chain.run(message.message)
        return output.strip()