import json
import gradio as gr
from huggingface_hub import InferenceClient
from huggingface_hub.utils._errors import HfHubHTTPError
import os
from huggingface_hub.utils._errors import HfHubHTTPError
print("HfHubHTTPError is available.")

# Load the specific questions and answers from the JSON file
with open('promptlist.json', 'r') as file:
    prompt_data = json.load(file)
# Use the provided access token
HF_TOKEn = os.getenv('HF_TOKEN')
print(os.getenv('HF_TOKEN'))
print({HF_TOKEn})
print(f"Your Hugging Face token is: {HF_TOKEn}")

# Initialize the Hugging Face Inference Client with the access token
client = InferenceClient(
    # "meta-llama/Llama-3.2-3B-Instruct",
    token=HF_TOKEn
)

def chat_mem(message, chat_history):
    chat_history_role = [{"role": "system", "content": "You are a helpful assistant."}]

    if chat_history:
        for user_msg, assistant_msg in chat_history:
            chat_history_role.append({"role": "user", "content": user_msg})
            chat_history_role.append({"role": "assistant", "content": assistant_msg})

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

    # Check for specific questions from prompt.json
    specific_question_found = False
    for item in prompt_data:
        if message.strip().lower() in [q.strip().lower() for q in item["prompt"]]:
            assistant_reply = item["completion"]
            specific_question_found = True
            break

    if not specific_question_found:
        try:
            chat_completion = client.chat_completion(
                messages=chat_history_role,
                max_tokens=500,
            )
            assistant_reply = chat_completion.choices[0].message.content
        except HfHubHTTPError as e:
            if e.response.status_code == 429:  # Rate limit error
                assistant_reply = "Rate limit reached. Please try again later."
            else:
                assistant_reply = "An error occurred. Please try again."

    chat_history.append((message, assistant_reply))

    return "", chat_history, chat_history  # Return the message, state, and chatbot history

with gr.Blocks() as demo:
    with gr.Column():
        gr.HTML("""
        <style>
            .send-button {
                background-color: #6f0389;
                color: white;
                border: none;
                padding: 10px 20px;
                font-size: 16px;
                cursor: pointer;
                border-radius: 5px;
            }
            .send-button:hover {
                background-color: #5a026e;
            }
            .header {
                text-align: center;
                margin-bottom: 20px;
            }
            .header h1 {
                font-family: 'Arial', sans-serif;
                color: #333;
            }
            .header p {
                font-family: 'Arial', sans-serif;
                color: #555;
            }
        </style>
        <div class="header">
            <h1>Meta-Llama3 (FAME)</h1>
            <p>FAME AI ASSISTANT</p>
        </div>
        """)
        chatbot = gr.Chatbot()
        state = gr.State([])  # Initialize state to store chat history
        msg = gr.Textbox(interactive=True, placeholder="Type your message here...")
        with gr.Row():
            clear = gr.ClearButton([msg, chatbot, state], icon="https://img.icons8.com/?size=100&id=Xnx8cxDef16O&format=png&color=000000")
            send_btn = gr.Button("Send", variant='primary', elem_classes=["send-button"], icon="https://img.icons8.com/?size=100&id=g8ltXTwIfJ1n&format=png&color=ffffff")
        msg.submit(fn=chat_mem, inputs=[msg, state], outputs=[msg, state, chatbot])
        send_btn.click(fn=chat_mem, inputs=[msg, state], outputs=[msg, state, chatbot])
print(os.getenv('HF_TOKEN'))
if __name__ == "__main__":
    demo.launch(share=True)