import gradio as gr
import asyncio
from core import chatbot, estimate_costs
import openai
import PyPDF2
from file_operations import open_file

prompts = [
    "Can you give me a very clear explanation of the core assertions, mechanics mentioned in this paper?",
    "Can you give a short summary with bullet points and key takeaways",
    "Can you give me an analogy or metaphor that will help explain this to a broad audience.",
]


async def process_prompt(prompt, ALL_MESSAGES, model):
    report = ""
    ALL_MESSAGES.append({"role": "user", "content": prompt})
    response, tokens = await chatbot(ALL_MESSAGES, model)
    ALL_MESSAGES.append({"role": "assistant", "content": response})
    report += "\n\n\n\nQ: %s\n\nA: %s" % (prompt, response)
    return report


async def process_pdf_content(text, prompts):
    model = "gpt-3.5-turbo-16k"
    if len(text) > 22000:
        text = text[:22000]
        model = "gpt-4-32k"
    prompt_tokens = len(text) / 0.75
    for p in prompts:
        prompt_tokens += len(p) / 0.75
    costs = estimate_costs(prompt_tokens, model)
    if costs > 2:
        return f"THIS IS WAY TO MUCH {costs}"
    else:
        ALL_MESSAGES = [{"role": "system", "content": text}]
        prompt_tasks = [process_prompt(p, ALL_MESSAGES, model) for p in prompts]
        results = await asyncio.gather(*prompt_tasks)
        return " ".join(results).strip()


def process_pdf(pdf_file, prompt1, prompt2, prompt3, key):
    
    openai.api_key = key
    # Open the PDF file
    prompts = [
        prompt for prompt in [prompt1, prompt2, prompt3] if prompt
    ]  # Only include prompts that are not empty
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    paper = ""
    for page_num in range(len(pdf_reader.pages)):
        page = pdf_reader.pages[page_num]
        paper += page.extract_text()
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    task = process_pdf_content(paper, prompts)
    result = loop.run_until_complete(task)

    return result




iface = gr.Interface(
    fn=process_pdf,
    inputs=[
        gr.inputs.File(),
        gr.inputs.Textbox(
            lines=2, placeholder="Enter Prompt 1 Here...", label="Prompt 1"
        ),
        gr.inputs.Textbox(
            lines=2, placeholder="Enter Prompt 2 Here...", label="Prompt 2"
        ),
        gr.inputs.Textbox(
            lines=2, placeholder="Enter Prompt 3 Here...", label="Prompt 3"
        ),
        gr.inputs.Textbox(
            lines=2, placeholder="Enter Key starts with sk", label="OPENAI API KEY"
        )
    ],
    outputs="text",
    title="Paper Analyser",
    description="This tool analyse your academic papers and returns key findings",
)

iface.launch()