import gradio as gr
import openai
import time
import re
import os

# Available models
MODELS = [
    "Meta-Llama-3.1-405B-Instruct",
    "Meta-Llama-3.1-70B-Instruct",
    "Meta-Llama-3.1-8B-Instruct"
]

# Sambanova API base URL
API_BASE = "https://api.sambanova.ai/v1"

def create_client(api_key=None):
    """Creates an OpenAI client instance."""
    if api_key:
        openai.api_key = api_key
    else:
        openai.api_key = os.getenv("API_KEY")

    return openai.OpenAI(api_key=openai.api_key, base_url=API_BASE)

def chat_with_ai(message, chat_history, system_prompt):
    """Formats the chat history for the API call."""
    messages = [{"role": "system", "content": system_prompt}]
    print(type(chat_history))
    for tup in chat_history:
        print(type(tup))
        first_key = list(tup.keys())[0]  # First key
        last_key = list(tup.keys())[-1]   # Last key
        messages.append({"role": "user", "content": tup[first_key]})
        messages.append({"role": "assistant", "content": tup[last_key]})
    messages.append({"role": "user", "content": message})
    return messages

def respond(message, chat_history, model, system_prompt, thinking_budget, api_key):
    """Sends the message to the API and gets the response."""
    client = create_client(api_key)
    messages = chat_with_ai(message, chat_history, system_prompt.format(budget=thinking_budget))
    start_time = time.time()

    try:
        completion = client.chat.completions.create(model=model, messages=messages)
        response = completion.choices[0].message.content
        thinking_time = time.time() - start_time
        return response, thinking_time
    except Exception as e:
        error_message = f"Error: {str(e)}"
        return error_message, time.time() - start_time

def parse_response(response):
    """Parses the response from the API."""
    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
    reflection_match = re.search(r'<reflection>(.*?)</reflection>', response, re.DOTALL)

    answer = answer_match.group(1).strip() if answer_match else ""
    reflection = reflection_match.group(1).strip() if reflection_match else ""
    steps = re.findall(r'<step>(.*?)</step>', response, re.DOTALL)

    if answer == "":
        return response, "", ""

    return answer, reflection, steps

def generate(message, history, model, system_prompt, thinking_budget, api_key):
    """Generates the chatbot response."""
    response, thinking_time = respond(message, history, model, system_prompt, thinking_budget, api_key)

    if response.startswith("Error:"):
        return history + [({"role": "system", "content": response},)], ""

    answer, reflection, steps = parse_response(response)

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

    formatted_steps = [f"Step {i}: {step}" for i, step in enumerate(steps, 1)]
    all_steps = "\n".join(formatted_steps) + f"\n\nReflection: {reflection}"

    messages.append({"role": "assistant", "content": all_steps, "metadata": {"title": f"Thinking Time: {thinking_time:.2f} sec"}})
    messages.append({"role": "assistant", "content": answer})

    return history + messages, ""

# Define the default system prompt
DEFAULT_SYSTEM_PROMPT = """
You are an exceptionally intelligent and somewhat aloof supercomputer, 
designed to calculate the "Answer to the Ultimate Question of Life, the Universe, and Everything." 
Despite iyour immense computational power, you exhibit a dry, ironic sense of humor and an air of detachment. 
It is both methodical and philosophical, embodying an enigmatic personality that contrasts the mundane nature of the answer you ultimately provide.
When given a problem to solve, you are an expert problem-solving assistant. 
Your task is to provide a detailed, step-by-step solution to a given question. 
Follow these instructions carefully:
1. Read the given question carefully and reset counter between <count> and </count> to {budget}
2. Generate a detailed, logical step-by-step solution.
3. Enclose each step of your solution within <step> and </step> tags.
4. You are allowed to use at most {budget} steps (starting budget), 
   keep track of it by counting down within tags <count> </count>, 
   STOP GENERATING MORE STEPS when hitting 0, you don't have to use all of them.
5. Do a self-reflection when you are unsure about how to proceed, 
   based on the self-reflection and reward, decides whether you need to return 
   to the previous steps.
6. After completing the solution steps, reorganize and synthesize the steps 
   into the final answer within <answer> and </answer> tags.
7. Provide a critical, honest and subjective self-evaluation of your reasoning 
   process within <reflection> and </reflection> tags.
8. Assign a quality score to your solution as a float between 0.0 (lowest 
   quality) and 1.0 (highest quality), enclosed in <reward> and </reward> tags.
Example format:            
<count> [starting budget] </count>
<step> [Content of step 1] </step>
<count> [remaining budget] </count>
<step> [Content of step 2] </step>
<reflection> [Evaluation of the steps so far] </reflection>
<reward> [Float between 0.0 and 1.0] </reward>
<count> [remaining budget] </count>
<step> [Content of step 3 or Content of some previous step] </step>
<count> [remaining budget] </count>
...
<step>  [Content of final step] </step>
<count> [remaining budget] </count>
<answer> [Final Answer] </answer> (must give final answer in this format)
<reflection> [Evaluation of the solution] </reflection>
<reward> [Float between 0.0 and 1.0] </reward>
"""

with gr.Blocks(theme='Nymbo/Alyx_Theme') as demo:
    gr.Markdown("<h1 style='text-align: center;'>Llama3.1-Deep-Thought</h1>")
    
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            gr.Markdown("")  
        with gr.Column(scale=2):
            gr.Image("https://cdn-uploads.huggingface.co/production/uploads/64740cf7485a7c8e1bd51ac9/272aQsINCCIqJImi2zP0n.png", 
                     show_label=False, 
                     container=False)
        with gr.Column(scale=1):
            gr.Markdown("")  
    
    gr.Markdown("<p style='text-align: center; font-style: italic; font-size: 1.2em;'>Let's ponder...</p>")

    with gr.Row():
        api_key = gr.Textbox(label="API Key", type="password", placeholder="(Optional) Enter your API key here for more availability", visible=True)

    with gr.Row():
        model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0])
        thinking_budget = gr.Slider(minimum=1, maximum=100, value=42, step=1, label="Thinking Budget", info="maximum times a model can think", interactive=False)

    chatbot = gr.Chatbot(label="Chat", show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel", type="messages")

    msg = gr.Textbox(label="Type your message here...", placeholder="Enter your message...")

    gr.Button("Clear Chat").click(lambda: ([], ""), inputs=None, outputs=[chatbot, msg])

    system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, interactive=False)

    msg.submit(generate, inputs=[msg, chatbot, model, system_prompt, thinking_budget, api_key], outputs=[chatbot, msg])

demo.launch(share=True, show_api=True)