import aiohttp
import asyncio
import json,os
import yaml
import google.generativeai as palm
from App.Embedding.utils.Initialize import search
PALM_API = ""
API_KEY = os.environ.get("PALM_API", PALM_API)
palm.configure(api_key=API_KEY)
class GenerativeAIAssistant:
    def __init__(self, api_key=API_KEY, model='chat-bison-001', temperature=0.85,
                 candidate_count=1, top_k=40, top_p=0.95):
        self.api_key = api_key
        self.model = model
        self.temperature = temperature
        self.candidate_count = candidate_count
        self.top_k = top_k
        self.top_p = top_p
        self.examples=[{"input": {"content": "hello"}, "output": {"content": "Hello to you too! How can I help you today?"}}]
        self.context = "You are a helpful assistant"

    def generate_template(self,question,task_id,summary=None):
        if summary == None:
            self.context = "You are a helpful assistant"
        else:
            self.context = summary
        contexts=search(question,task_id=task_id)
        context_yaml = ""        
        for context in contexts:
            context_yaml += "\n"+ yaml.dump(context)        
        Template =f''' 
      #Instructions
      You are given the following context in yaml of a transcript of a youtube video, the start and end times are indicated and the text that was said is also given. You are also given a question, use the context to answer the question in a consise manner, make it short and to the point, don't provide additional details.
     

      #Context
      {context_yaml}

      #Your response be as short as possible and to the point.
      
      #Question
      {question}
        '''
        return Template

    async def generate_message(self,  messages,task_id):
        latest_message = messages[-1]['content']
        latest_message={"content":self.generate_template(latest_message,task_id)}
        messages[-1]=latest_message
        url = f'https://generativelanguage.googleapis.com/v1beta3/models/{self.model}:generateMessage?key={self.api_key}'
        data = {
            "prompt": {
                "context": self.context,
                "examples": self.examples,
                "messages": messages
            },
            "temperature": self.temperature,
            "top_k": self.top_k,
            "top_p": self.top_p,
            "candidate_count": self.candidate_count
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(url, json=data, headers={'Content-Type': 'application/json'}) as response:
                try:
                    temp=  await response.json()
                    return temp['candidates']
                except Exception as e:
                    return f"Error ⚠️ {e}"