import gradio as gr
import os
import json
import time
import groq
from g1 import generate_response

def format_steps(steps, total_time):
    md_content = ""
    for title, content, thinking_time in steps:
        if title == "Final Answer":
            md_content += f"### {title}\n"
            md_content += f"{content}\n"
        else:
            md_content += f"#### {title}\n"
            md_content += f"{content}\n"
            md_content += f"_Thinking time for this step: {thinking_time:.2f} seconds_\n"
            md_content += "\n---\n"
    if total_time != 0:
        md_content += f"\n**Total thinking time: {total_time:.2f} seconds**"
    return md_content

def main(api_key, user_query, mode):
    if mode == "private" and not api_key:
        yield "Please enter your Groq API key to proceed."
        return
    
    if not user_query:
        yield "Please enter a query to get started."
        return
    
    try:
        # Initialize the Groq client with the provided API key or the environment variable
        if mode == "public":
            client = groq.Groq(api_key=os.getenv("GROQ_API_KEY"))
        else:
            client = groq.Groq(api_key=api_key)
    except Exception as e:
        yield f"Failed to initialize Groq client. Error: {str(e)}"
        return
    
    try:
        for steps, total_time in generate_response(user_query, custom_client=client):
            formatted_steps = format_steps(steps, total_time if total_time is not None else 0)
            yield formatted_steps
    except Exception as e:
        yield f"An error occurred during processing. Error: {str(e)}"
        return

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 g1: Использование Llama-3.1 70b на Groq для создания O1-подобных цепочек рассуждений")
    
    gr.Markdown("""
    Это не идеально, и точность еще предстоит оценить. Он питается от Groq, так что шаг рассуждения быстрый!
    """)
    
    with gr.Row():
        with gr.Column():
            mode_toggle = gr.Radio(["public", "private"], label="Режим ключа API", value="public")
            api_input = gr.Textbox(
                label="Введите свой ключ Groq API:",
                placeholder="Ваш ключ API Groq",
                type="password",
                visible=False  # Initially hidden
            )
            user_input = gr.Textbox(
                label="Введите свой запрос:",
                placeholder="Например, сколько 'х в словом охуевшийхомячек?",
                lines=2
            )
            submit_btn = gr.Button("Генерировать ответ")
            gr.Markdown("\n")
    
    with gr.Row():
        with gr.Column():
            output_md = gr.Markdown()
    
    # Show/hide the API key input based on the mode toggle
    mode_toggle.change(lambda mode: gr.update(visible=mode == "private"), mode_toggle, api_input)
    
    submit_btn.click(fn=main, inputs=[api_input, user_input, mode_toggle], outputs=output_md)

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()