import os
from langchain.chains.summarize import load_summarize_chain
from utils import openai_llm
from dotenv import load_dotenv
from langchain_core.prompts import PromptTemplate

from langfuse.callback import CallbackHandler
load_dotenv()

os.environ["LANGFUSE_PUBLIC_KEY"] = os.getenv("LANGFUSE_PUBLIC_KEY")
os.environ["LANGFUSE_SECRET_KEY"] = os.getenv("LANGFUSE_SECRET_KEY")
os.environ["LANGFUSE_HOST"] = os.getenv("LANGFUSE_HOST") 


 
langfuse_handler = CallbackHandler()

map_prompt = """
Write a detailed summary of the following text.  
The summary should be in the language of the text.
Do not miss any important details.
Mention all the important entities at the end of the summary with a heading 'IMPORTANT ENTITIES' and bullets underneath the heading


"{text}"
DETAILED SUMMARY:
"""
map_prompt_template = PromptTemplate(template=map_prompt, input_variables=["text"])

combine_prompt = """
Write a detailed summary of the following text delimited by triple backquotes.
Return your response as a detailed paragraph.
The Response should be in the language of the text
Mention all the important entities at the end of the summary with a heading 'IMPORTANT ENTITIES' and bullets underneath the heading

```{text}```
PARAGRAPH SUMMARY:
"""
combine_prompt_template = PromptTemplate(template=combine_prompt, input_variables=["text"])

def setup_summary_chain(api_key):
    """Set up a summary chain with a specified LLM."""
    llm = openai_llm(api_key=api_key)
    return load_summarize_chain(llm=llm, chain_type='map_reduce',
                                map_prompt=map_prompt_template,
                                combine_prompt=combine_prompt_template)


def summarize_documents(documents, api_key):
    """Generate summaries for provided documents."""
    summary_chain = setup_summary_chain(api_key)
    return summary_chain.run(documents, callbacks=[langfuse_handler])