import aiohttp
import asyncio,pprint
import google.generativeai as palm
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain import PromptTemplate
import os
PALM_API = ''
API_KEY=os.environ.get("PALM_API",PALM_API)
palm.configure(api_key=API_KEY)




text_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n","."], chunk_size=40_000, chunk_overlap=500)


map_prompt = """
Write a verbose summary  like a masters student of the following:
"{text}"
CONCISE SUMMARY:
"""


combine_prompt = """
Write a concise summary of the following text delimited by triple backquotes.
Return your response in a detailed verbose paragraph which covers the text. Make it as insightful to the reader as possible, write like a masters student.

```{text}```

SUMMARY:
"""
def count_tokens(text):
    return palm.count_message_tokens(prompt=text)['token_count']


async def PalmTextModel(text,candidates=1):
    url = f"https://generativelanguage.googleapis.com/v1beta2/models/text-bison-001:generateText?key={API_KEY}"

    headers = {
        "Content-Type": "application/json",
    }

    data = {
        "prompt": {
            "text": text
        },
        "temperature": 0.95,
        "top_k": 100,
        "top_p": 0.95,
        "candidate_count": candidates,
        "max_output_tokens": 1024,
        "stop_sequences": ["</output>"],
        "safety_settings": [
            {"category": "HARM_CATEGORY_DEROGATORY", "threshold": 4},
            {"category": "HARM_CATEGORY_TOXICITY", "threshold": 4},
            {"category": "HARM_CATEGORY_VIOLENCE", "threshold": 4},
            {"category": "HARM_CATEGORY_SEXUAL", "threshold": 4},
            {"category": "HARM_CATEGORY_MEDICAL", "threshold": 4},
            {"category": "HARM_CATEGORY_DANGEROUS", "threshold": 4},
        ],
    }


    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=data, headers=headers) as response:
            if response.status == 200:
                result = await response.json()
                # print(result)
                if candidates>1:
                    temp = [candidate["output"] for candidate in result["candidates"]]
                    return temp                
                temp = result["candidates"][0]["output"]
                return temp
            else:
                print(f"Error: {response.status}\n{await response.text()}")


async def Summarizer(essay):

    docs = text_splitter.create_documents([essay])

    #for 1 large document
    if len(docs) == 1:
        tasks = [PalmTextModel(combine_prompt.format(text=doc.page_content)) for doc in docs]
        # Gather and execute the tasks concurrently
        responses = await asyncio.gather(*tasks)
        ans=" ".join(responses)
        return ans 

    tasks = [PalmTextModel(map_prompt.format(text=doc.page_content)) for doc in docs]
    # Gather and execute the tasks concurrently
    responses = await asyncio.gather(*tasks)
    main=" ".join(responses)
    ans=await PalmTextModel(combine_prompt.format(text=main),candidates=1)
    return ans