import os

import gradio as gr
import requests
from langchain.chains import RetrievalQA
from langchain.document_loaders import PDFMinerLoader
from langchain.indexes import VectorstoreIndexCreator
from langchain.llms import OpenAI


def set_openai_key(raw_key):
    # Check if the API is valid
    headers = {"Authorization": f"Bearer {raw_key}"}
    response = requests.get("https://api.openai.com/v1/engines", headers=headers)
    if response.status_code != 200:
        raise gr.Error("API key is not valid. Check the key and try again.")

    os.environ["OPENAI_API_KEY"] = raw_key
    return gr.File.update(interactive=True), gr.Button.update(interactive=True)


def create_langchain(pdf_object):
    loader = PDFMinerLoader(pdf_object.name)
    index_creator = VectorstoreIndexCreator()
    docsearch = index_creator.from_loaders([loader])
    chain = RetrievalQA.from_chain_type(
        llm=OpenAI(),
        chain_type="stuff",
        retriever=docsearch.vectorstore.as_retriever(),
        input_key="question",
        verbose=True,
        return_source_documents=True,
    )
    return chain, gr.Button.update(interactive=True)


def ask_question(chain, question_text):
    return chain({"question": question_text})["result"]


with gr.Blocks() as demo:
    # Sate objects
    chain_state = gr.State()

    # Layout
    oai_token = gr.Textbox(
        label="OpenAI Token",
        placeholder="Lm-iIas452gaw3erGtPar26gERGSA5RVkFJQST23WEG524EWEl",
    )

    pdf_object = gr.File(
        label="Upload your CV in PDF format",
        file_count="single",
        type="file",
        interactive=False,
    )
    gr.Examples(
        examples=[
            os.path.join(os.path.abspath(""), "sample_data", "CV_AITOR_MIRA.pdf")
        ],
        inputs=pdf_object,
        label="Example CV",
    )
    create_chain_btn = gr.Button(value="Create CVchat", interactive=False)

    question_placeholder = """Enumerate the candidate's top 5 hard skills and rate them by importance from 0 to 5.
Example:
- Algebra 5/5"""
    question_box = gr.Textbox(label="Question", value=question_placeholder)
    qa_button = gr.Button(value="Submit question", interactive=False)

    # Actions
    oai_token.change(
        set_openai_key, inputs=oai_token, outputs=[pdf_object, create_chain_btn]
    )
    lchain = create_chain_btn.click(
        create_langchain, inputs=pdf_object, outputs=[chain_state, qa_button]
    )
    qa_button.click(
        ask_question,
        inputs=[chain_state, question_box],
        outputs=gr.Textbox(label="Answer"),
    )

demo.launch(debug=True)