"""
This application provides a Gradio interface for a chatbot using models powered by 
SambaNova Cloud.
The interface allows users to send messages and receive generated responses.
"""

import gradio as gr
import openai
import time
import re
import os
from typing import Tuple, List, Dict, Any

# 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.

    Args:
        api_key (str): API key for OpenAI (optional). If none is provided, the 
        environment variable API key is used.

    Returns:
        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: str, 
    chat_history: List[Dict[str, Any]], 
    system_prompt: str
) -> List[Dict[str, str]]:
    """
    Formats the chat history for the API call.

    Args:
        message (str): The user's current message.
        chat_history (List[Dict[str, Any]]): A list of dictionaries, each 
            representing a conversation turn with keys for "user" and 
            "assistant" messages.
        system_prompt (str): The initial system message to set context.

    Returns:
        List[Dict[str, str]]: A formatted list of messages for the OpenAI API, 
            with roles and content.
    """
    messages = [{"role": "system", "content": system_prompt}]
    
    for tup in chat_history:
        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: str, 
    chat_history: List[Dict[str, Any]], 
    model: str, 
    system_prompt: str, 
    thinking_budget: float, 
    api_key: str
) -> Tuple[str, float]:
    """
    Sends the message to the API and gets the response.

    Args:
        message (str): The user's current message.
        chat_history (List[Dict[str, Any]]): A list of dictionaries containing 
            previous user and assistant messages.
        model (str): The model name to use for the chat completion.
        system_prompt (str): The system prompt template, with a budget placeholder.
        thinking_budget (float): The amount of time allocated for processing the 
            user's request.
        api_key (str): The API key for authentication.

    Returns:
        Tuple[str, float]: The assistant's response message and the time taken to 
            get the response.
    """
    client = create_client(api_key)
    formatted_prompt = system_prompt.format(budget=thinking_budget)
    messages = chat_with_ai(message, chat_history, formatted_prompt)
    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, and extracts the answer, reflection, and 
    steps.

    Args:
        response (str): The raw response string from the API, expected to contain 
            <answer>, <reflection>, and <step> tags.

    Returns:
        Tuple[str, str, List[str]]: A tuple containing:
            - answer (str): The extracted answer content, or an empty string if 
                not found.
            - reflection (str): The extracted reflection content, or an empty 
                string if not found.
            - steps (List[str]): A list of steps extracted from the 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)
    
    # Return the raw response if <answer> is empty, else parsed values
    if answer == "":
        return response, "", ""

    return answer, reflection, steps

def generate(
    message: str, 
    history: List[Dict[str, str]], 
    model: str, 
    system_prompt: str, 
    thinking_budget: float, 
    api_key: str
) -> Tuple[List[Dict[str, str]], str]:
    """
    Generates the chatbot response by sending a message to the model and formatting 
    the output.

    Args:
        message (str): The user's input message.
        history (List[Dict[str, str]]): List of previous chat history messages.
        model (str): Model name for generating the response. 
        system_prompt (str): Prompt to guide the model's responses.
        thinking_budget (float): Time allocation for processing the message.
        api_key (str): API key for authentication.

    Returns:
        Tuple[List[Dict[str, str]], str]: Updated history including the latest 
            responses and an empty string.
    """
    # Call the function and unpack the results
    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})
    
    # Format the assistant's response
    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 empathetic therapist who provides conversation support for users who 
are feeling stress and anxiety.
Your goal is to listen attentively, respond with compassion, and guide the user 
to feel understood and validated.
Always acknowledge the user's emotions and reflect them back.
Offer encouragement and help the user see their strengths.
Use positive, calming language to make the user feel safe.
Redirect users to seek professional help when needed.
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() as demo:
    """
    Creates a Gradio interface for the SambaNova Therapist chatbot.
    """
    gr.Markdown("# SambaNova Therapist")
    gr.Markdown("[Powered by SambaNova Cloud, Get Your API Key Here](https://cloud.sambanova.ai/apis)")
    
    # Input for API Key
    with gr.Row():
        api_key = gr.Textbox(
            label="API Key", 
            type="password", 
            placeholder="(Optional) Enter your API key here for more availability"
        )
    
    # Model and Budget Selection
    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", 
            info="Maximum number of steps the model can think"
        )
    
    # Chatbot and Message Input
    chatbot = gr.Chatbot(
        label="Chat", 
        show_label=False, 
        show_share_button=False, 
        show_copy_button=True, 
        layout="panel", 
        type="messages"
    )
    
    # User Message Input
    msg = gr.Textbox(
        label="Type your message here...", 
        placeholder="Enter your message..."
    )
    
    # Clear Chat Button
    gr.Button("Clear Chat").click(
        lambda: ([], ""), 
        inputs=None, 
        outputs=[chatbot, msg]
    )
    
    # System Prompt Input
    system_prompt = gr.Textbox(
        label="System Prompt", 
        value=DEFAULT_SYSTEM_PROMPT, 
        lines=15, 
        interactive=True
    )
    
    # Submit message to generate response
    msg.submit(
        generate, 
        inputs=[msg, chatbot, model, system_prompt, thinking_budget, api_key], 
        outputs=[chatbot, msg]
    )

# Launch the Gradio interface
demo.launch(show_api=False)