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"
]

def create_client(api_key=None):
    openai.api_base = "https://api.sambanova.ai/v1"  # Fixed Base URL
    if api_key:
        openai.api_key = api_key
    else:
        openai.api_key = os.getenv("API_KEY")

def chat_with_ai(message, chat_history, system_prompt):
    messages = [
        {"role": "system", "content": system_prompt},
    ]

    for human, ai in chat_history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": ai})

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

    return messages

def respond(message, chat_history, model, system_prompt, thinking_budget, api_key):
    print("Starting respond function...")
    create_client(api_key)  # Sets api_key and api_base globally
    messages = chat_with_ai(message, chat_history, system_prompt.format(budget=thinking_budget))
    start_time = time.time()

    try:
        print("Calling OpenAI API...")
        completion = openai.ChatCompletion.create(
            model=model,
            messages=messages,
            stream=False  # Set to False for synchronous response
        )
        response = completion.choices[0].message['content']
        thinking_time = time.time() - start_time
        print("Response received from OpenAI API.")
        return response, thinking_time
    except Exception as e:
        error_message = f"Error: {str(e)}"
        print(error_message)
        return error_message, time.time() - start_time

def parse_response(response):
    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 is not "":
        return answer, reflection, steps
    else:
        return response, "", ""

def process_chat(message, history, model, system_prompt, thinking_budget, api_key):
    print(f"Received message: {message}")
    #if not api_key:
    #    print("API key missing")
    #    return history + [("System", "Please provide your API Key before starting the chat.")]

    try:
        formatted_system_prompt = system_prompt.format(budget=thinking_budget)
    except KeyError as e:
        error_msg = f"System prompt missing placeholder: {str(e)}"
        print(error_msg)
        return history + [("System", error_msg)]

    response, thinking_time = respond(message, history, model, formatted_system_prompt, thinking_budget, api_key)

    if response.startswith("Error:"):
        return history + [("System", response)]

    answer, reflection, steps = parse_response(response)

    formatted_response = f"**Answer:** {answer}\n\n**Reflection:** {reflection}\n\n**Thinking Steps:**\n"
    for i, step in enumerate(steps, 1):
        formatted_response += f"**Step {i}:** {step}\n"

    formatted_response += f"\n**Thinking time:** {thinking_time:.2f} s"

    print(f"Appended response: {formatted_response}")
    return history + [(message, formatted_response)]

# Define the default system prompt
default_system_prompt = """
You are a helpful assistant in normal conversation.
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>

<reflection> [Evaluation of the solution] </reflection>

<reward> [Float between 0.0 and 1.0] </reward>
"""

with gr.Blocks() as demo:
    gr.Markdown("# Llama3.1-Instruct-O1")
    gr.Markdown("[Powered by Llama3.1 models through SambaNova Cloud API](https://sambanova.ai/fast-api?api_ref=907266)")

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

    with gr.Row():
        model = gr.Dropdown(
            choices=MODELS,
            label="Select Model",
            value=MODELS[0]
        )
        thinking_budget = gr.Slider(
            minimum=1,
            maximum=100,
            value=10,
            step=1,
            label="Thinking Budget"
        )

    system_prompt = gr.Textbox(
        label="System Prompt",
        value=default_system_prompt,
        lines=15,
        interactive=True
    )

    with gr.Row():
        msg = gr.Textbox(
            label="Type your message here...",
            placeholder="Enter your message..."
        )
        submit = gr.Button("Submit")
        clear = gr.Button("Clear Chat")

    chatbot = gr.Chatbot(
        label="Chat History"
    )

    # Initialize chat history as a Gradio state
    chat_history = gr.State([])

    def handle_submit(message, history, model, system_prompt, thinking_budget, api_key):
        updated_history = process_chat(message, history, model, system_prompt, thinking_budget, api_key)
        return updated_history, ""

    def handle_clear():
        return [], ""

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

    clear.click(
        handle_clear,
        inputs=None,
        outputs=[chatbot, msg]
    )

demo.launch()