import os
from typing import Optional, Tuple

import gradio as gr
from dotenv import load_dotenv
from langchain.chains.summarize import load_summarize_chain
from langchain_core.prompts import PromptTemplate
from langchain_community.callbacks import get_openai_callback
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import ChatOpenAI


os.makedirs("data", exist_ok=True)
load_dotenv()
OPENAI_API_KEY: Optional[str] = os.getenv("OPENAI_API_KEY")


def summarize_pdf(
    pdf_file: bytes, custom_prompt: str = "", openai_api_key: Optional[str] = None
) -> Tuple[str, str]:
    """
    Summarizes the content of a PDF file using a custom prompt.

    Args:
        pdf_file (bytes): The uploaded PDF file as bytes.
        custom_prompt (str): The prompt for summarization.
        openai_api_key (Optional[str]): User-provided OpenAI API key.

    Returns:
        Tuple[str, str]: Summary in markdown format and the cost in USD.
    """
    pdf_path: str = os.path.join("data", "tmp.pdf")
    try:
        with open(pdf_path, "wb") as f:
            f.write(pdf_file)
    except IOError as e:
        return f"Failed to write PDF file: {e}", "N/A"

    api_key: Optional[str] = openai_api_key or OPENAI_API_KEY

    if not api_key:
        return "Error: No OpenAI API key provided.", "N/A"

    with get_openai_callback() as callback:
        try:
            model = ChatOpenAI(
                model="gpt-4o-mini",  # Verify the correct model name
                temperature=0.0,
                openai_api_key=api_key,
            )

            loader = PyPDFLoader(pdf_path)
            documents = loader.load_and_split()

            prompt_text: str = custom_prompt.strip() or default_prompt
            prompt_template: str = f"{prompt_text}\n\n{{text}}\n\nSUMMARY:"
            prompt = PromptTemplate(template=prompt_template, input_variables=["text"])

            summarize_chain = load_summarize_chain(
                llm=model,
                chain_type="map_reduce",
                map_prompt=prompt,
                combine_prompt=prompt,
            )

            chain_input = {"input_documents": documents}
            result = summarize_chain(chain_input, return_only_outputs=True)
            summary: str = result.get("output_text", "No summary generated.")
            total_cost: float = callback.total_cost

            return summary, f"${total_cost:.4f}"

        except Exception as e:
            return f"An error occurred during summarization: {str(e)}", "N/A"


default_prompt: str = (
    "Summarize this paper. Return markdown, keep it in a language that scientists understand, "
    "but the purpose is to highlight the key takeaways, so that we save time for the reader."
)
with gr.Blocks() as demo:
    gr.Markdown("# PDF Summarizer 📝")
    gr.Markdown(
        "Upload a PDF, customize your summarization prompt, and get a concise summary along with the processing cost."
    )

    with gr.Row():
        with gr.Column():
            api_key_label: str
            placeholder_text: str

            if OPENAI_API_KEY is None:
                api_key_label = "OpenAI API Key"
                placeholder_text = "Enter your OpenAI API key."
            else:
                api_key_label = "OpenAI API Key (Optional)"
                placeholder_text = (
                    "Enter your OpenAI API key if you want to override the global key."
                )

            api_key_input = gr.Textbox(
                label=api_key_label,
                type="password",
                placeholder=placeholder_text,
            )
            prompt_input = gr.Textbox(
                label="Custom Prompt",
                lines=4,
                value=default_prompt,
                placeholder="Enter your custom summarization prompt here...",
            )
            pdf_input = gr.File(
                label="Upload PDF",
                type="binary",
                file_types=[".pdf"],
            )
            summarize_btn = gr.Button("Summarize")

        with gr.Column():
            cost_output = gr.Textbox(label="Approximate Cost (USD)", interactive=False)
            summary_output = gr.Markdown(label="Summary")

    summarize_btn.click(
        fn=summarize_pdf,
        inputs=[pdf_input, prompt_input, api_key_input],
        outputs=[summary_output, cost_output],
    )

    gr.Markdown("---")
    gr.Markdown("Created by [Daniel Herman](https://www.hermandaniel.com), check out the code [detrin/llm-pdf-summarization](https://github.com/detrin/llm-pdf-summarization).")

if __name__ == "__main__":
    demo.launch()