import gradio as gr
import pdfplumber
from groq import Groq
import os
import json

def file_to_string(file):
    if file.name.endswith(".pdf"):
        text = ''
        with pdfplumber.open(file.name) as pdf:
            for page in pdf.pages:
                text += page.extract_text()
        return text
    elif file.name.endswith(".txt"):
        return file.read().decode('utf-8')
    else:
        return "Unsupported file format."

def generate_summary(file, prompt, model, output_format, temperature, top_p):
    if not prompt:
        return "Please provide a prompt."

    file_text = file_to_string(file) if file else ""
    
    # Check if the {PDF_TEXT} token is in the prompt
    if "{PDF_TEXT}" in prompt and file_text:
        full_prompt = prompt.replace("{PDF_TEXT}", file_text)
    else:
        full_prompt = f"{prompt}\n{file_text}" if file_text else prompt

    try:
        client = Groq(api_key=os.environ["GROQ_API_KEY"])
        completion = client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": full_prompt
                }
            ],
            temperature=temperature,
            max_tokens=8000,
            top_p=top_p,
            response_format={"type": "json_object"} if output_format == "JSON" else None,
            stop=None,
        )

        output_text = completion.choices[0].message.content

        if output_format == "JSON":
            try:
                parsed_json = json.loads(output_text)
                output_text = json.dumps(parsed_json, indent=2)
            except json.JSONDecodeError:
                output_text = "Error: The model did not return valid JSON."

    except Exception as e:
        output_text = f"An error occurred: {str(e)}"

    return output_text

def clear_output():
    return None, "", "", "Plain Text", 0, 0.9

with gr.Blocks() as iface:
    gr.Markdown("LLAMA 70B Groq API")
    
    with gr.Row():
        file_input = gr.File(label="Upload File (Optional)")
        prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here... Use {PDF_TEXT} to insert PDF content at a specific location.", lines=3)
    
    with gr.Row():
        model_dropdown = gr.Dropdown(label="Choose Model", choices=["llama-3.1-8b-instant", "llama-3.1-70b-versatile"], value="llama-3.1-70b-versatile")
        output_format = gr.Radio(["Plain Text", "JSON"], label="Output Format", value="Plain Text")
    
    with gr.Row():
        temperature_slider = gr.Slider(minimum=0, maximum=2, value=0, step=0.1, label="Temperature")
        top_p_slider = gr.Slider(minimum=0, maximum=1, value=0.9, step=0.05, label="Top-p")

    with gr.Row():
        clear_button = gr.Button("Clear", size="small")
        submit_button = gr.Button("Generate Output")

    output = gr.Textbox(label="Output", lines=10, placeholder="Output will appear here...")

    submit_button.click(
        fn=generate_summary,
        inputs=[file_input, prompt_input, model_dropdown, output_format, temperature_slider, top_p_slider],
        outputs=[output]
    )

    clear_button.click(
        fn=clear_output,
        inputs=[],
        outputs=[file_input, prompt_input, output, output_format, temperature_slider, top_p_slider]
    )

iface.launch()