import gradio as gr
from rag import RAG, ServiceContextModule
from llama_index.core import set_global_service_context
import json
from prompts import general_prompt
from gradio_pdf import PDF
import requests
import os

service_context_module = None
current_model = None


def initialize(api_key, model_name):
    global service_context_module, current_model
    gr.Info("Initializing app")
    url = "https://api.groq.com/openai/v1/models"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    try:
        response = requests.get(url, headers=headers)
        data = response.json()
        models = [model["id"] for model in data["data"]]

    except Exception:
        gr.Error("Invalid API KEY")
        return gr.update(choices=[])

    if not service_context_module or current_model != model_name:
        try:
            service_context_module = ServiceContextModule(api_key, model_name)
        except Exception as e:
            print(e)
        current_model = model_name
        gr.Info("App started")
        set_global_service_context(
            service_context=service_context_module.service_context
        )
    else:
        gr.Info("App is already running")

    return gr.update(choices=models)


def process_document(file, query):
    if file.endswith(".pdf"):
        return process_pdf(file, query=query)
    else:
        return "Unsupported file format"


def postprocess_json_string(json_string: str) -> dict:
    json_string = json_string.replace("'", '"')
    json_string = json_string[json_string.rfind("{") : json_string.rfind("}") + 1]
    try:
        json_data = json.loads(json_string)
    except Exception as e:
        print("Error parsing output, invalid json format", e)
    return json_data


def process_pdf(file, query):
    rag_module = RAG(filepaths=[file])
    fields = [field for field in query.split(",")]
    formatted_prompt = general_prompt(fields=fields)
    response = rag_module.run_query_engine(prompt=formatted_prompt)
    extracted_json = postprocess_json_string(json_string=response)
    return extracted_json


with gr.Blocks(title="Document Information Extractor.") as app:
    gr.Markdown(
        value="""
# Welcome to Document Information Extractor.
Created by [@rajsinghparihar](https://huggingface.co/rajsinghparihar) for extracting useful information from pdf documents like invoices, salary slips, etc.
## Usage:
- In the Init Section, Enter your `GROQ_API_KEY` in the corresponding labeled textbox.
- choose the model from the list of available models.
- click `Initialize` to start the app.

- In the app section, you can upload a document (pdf files: currently works for readable pdfs only, will add ocr functionality later)
- Enter the entities you wanna extract as a comma seperated string. (check the examples for more info)
- Click Submit to see the extracted entities as a JSON object.
"""
    )
    with gr.Tab(label="Init Section") as init_tab:
        with gr.Row():
            api_key = gr.Text(
                label="Enter your Groq API KEY",
                type="password",
            )
            if api_key == "" or not api_key:
                api_key = os.getenv("GROQ_API_KEY")
            available_models = gr.Dropdown(
                value="llama3-70b-8192",
                label="Choose your LLM",
                choices=[
                    "gemma-7b-it",
                    "llama3-70b-8192",
                    "llama3-8b-8192",
                    "mixtral-8x7b-32768",
                    "whisper-large-v3",
                ],
            )
        init_btn = gr.Button(value="Initialize")
        init_btn.click(
            fn=initialize,
            inputs=[api_key, available_models],
            outputs=available_models,
        )
    with gr.Tab(label="App Section") as app_tab:
        iface = gr.Interface(
            fn=process_document,
            inputs=[
                PDF(label="Document"),
                gr.Text(
                    label="Entities you wanna extract in comma separated string format"
                ),
            ],
            outputs=gr.JSON(label="Extracted Entities"),
            description="Upload a PDF document and extract specified entities from it.",
            examples=[
                [
                    "examples/Commerce Bank Statement Sample.pdf",
                    "Customer Name, Account Number, Statement Date, Ending Balance, Total Deposits, Checks Paid",
                ],
                [
                    "examples/Salary-Slip-pdf.pdf",
                    "Employee Name, Bank Name, Location, Total Salary, Total Deductions",
                ],
            ],
        )
    gr.Markdown("""
## Pros of LLMs as information extractors over current extraction solutions:
- LLMs are able to understand the scope of the problem from the context and are more robust to typos or extraction failure

## Cons
- Higher Inference Cost
- Can't use free APIs for Sensitive documents.
""")

app.launch(server_name="0.0.0.0", server_port=7860)