import gradio as gr
from openai import OpenAI
import os
import logging
import time

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

css = '''
.gradio-container{max-width: 1000px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''

ACCESS_TOKEN = os.getenv("HF_TOKEN")


start_time = time.time()
logger.info("Loading Client....")

client = OpenAI(
    base_url="https://api-inference.huggingface.co/v1/",
    api_key=ACCESS_TOKEN,
)

end_time = time.time()
logger.info(f"Client Loaded. Time taken : {end_time - start_time} seconds.")

#interact with API
def respond(
    message,
    history,
    temperature,
    max_tokens,
):
    SYS_PROMPT = """
                  Extract the following information from the given text:
            Identify the specific areas where the work needs to be done and Add the furniture that has to be changed.
            Do not specify the work that has to be done.
            Format the extracted information in the following JSON structure:

            {
              "Area Type1": {
                  "Furnture1",
                  "Furnture2",
                  ...
              }
              "Area Type2": {
                   "Furnture1",
                  "Furnture2",
                  ...
              }
            }

            Requirements:
            1. Each area type (e.g., lobby, bar, etc.) should have its own node.
            3. List the furniture on which the work needs to be performed without specifying the work or units of items.
            4. Ignore any personal information or irrelevant details.
            5. Follow the JSON pattern strictly and ensure clarity and accuracy in the extracted information.

            Example:

            Given the paragraph: "In the lobby, replace 5 light fixtures and remove 2 old carpets. In the bar,
            install 3 new tables and remove 4 broken chairs."

            The JSON output should be:
            {
              "Lobby": {
                  "Light fixtures"
                  "Old carpets"
              },
              "Bar": {
                  "New tables"
                  "Broken chairs"
                }
              }
            }

            Please ensure that the output JSON is well-structured and includes only relevant details about the work to be done.
                  """
    messages = [{"role": "system", "content": SYS_PROMPT}]
    
    if len(history) == 0:
      pass
    else:
      history.pop()

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    start_time = time.time()
    logger.info("Generating Response....")
    
    for message in  client.chat.completions.create(
        model="meta-llama/Meta-Llama-3.1-8B-Instruct",
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        messages=messages,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

    end_time = time.time()
    logger.info(f"Response Generated. Time taken : {end_time - start_time} seconds.")


DESCRIPTION = '''
<div>
<h1 style="text-align: center;">ContenteaseAI custom trained model</h1>
</div>
'''

LICENSE = """
<p/>
---
For more information, visit our [website](https://contentease.ai).
"""

PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">ContenteaseAI Custom AI trained model</h1>
   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Enter the text extracted from the PDF:</p>
</div>
"""

css = """
h1 {
  text-align: center;
  display: block;
}
"""

chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')

with gr.Blocks(fill_height=True, css=css) as demo:
    gr.Markdown(DESCRIPTION)

    gr.ChatInterface(
        fn=respond,
        chatbot=chatbot,
        fill_height=True,
        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Slider(minimum=0, maximum=1, step=0.1, value=0.9, label="Temperature", render=False),
            gr.Slider(minimum=128, maximum=2000, step=1, value=2000, label="Max new tokens", render=False),
        ]
    )

    gr.Markdown(LICENSE)

if __name__ == "__main__":
    try:
        demo.launch(show_error=True, debug = True)
    except Exception as e:
        logger.error(f"Error launching Gradio demo: {e}")