import gradio as gr
from chatbot_simulator import ChatbotSimulation
from task_specific_data_population import DataPopulation
import os
openai_api_key = os.getenv("OPENAI_API_KEY")


simulation = None
conversation = []
display_conversation = []

def initialize_simulator(task, app_name, sitemap):
    """Initialize the simulator."""
    success = False  # Track if the operation succeeds
    retry_count = 0  # Track the number of retries
    max_retries = 50  # Set the maximum number of retries

    while not success and retry_count < max_retries:
        try:
            # Process data (simulating data loading)
            data_population = DataPopulation(api_key=openai_api_key)
            sitemap_data, page_details, user_state = data_population.process_data(task, sitemap)

            global simulation
            simulation = ChatbotSimulation(
                site_map=sitemap_data,
                page_details=page_details,
                user_state=user_state,
                task=task,
                app_name=app_name,
                log_location=f'conversation_log_{app_name}_human.txt',
                openai_api_key=openai_api_key,
                agent='human'
            )
            text = simulation.start_conversation()
            global conversation
            conversation.append({"role": "assistant", "content": text})
            log_conversation(simulation.log_location)
            display_conversation.append(('Start Simulator', text))
            display_conversation.pop(0)
            return display_conversation
        except Exception as e:
            # Handle the exception and increment retry count
            retry_count += 1
            print(f"Attempt {retry_count}/{max_retries}: An error occurred: {e}. Retrying...")


def log_conversation(log_location):
    """
    Append the conversation to the specified log file location.
    """
    try:
        with open(log_location, 'a') as f:  # Use 'a' for append mode
            for message in conversation:
                f.write(f"{message['role']}: {message['content']}\n\n")
    except Exception as e:
        print(f"Error logging conversation: {e}")


def chatbot_interaction(user_input):
    """Handle the conversation."""
    if simulation is None:
        return "Simulation is not initialized. Please start the simulator."

    try:
        # Perform one round of conversation
        response = simulation.one_conversation_round(user_input)
        global conversation
        conversation.append({"role": "user", "content": user_input})
        conversation.append({"role": "assistant", "content": response})
        log_conversation(simulation.log_location)
        display_conversation.append((user_input, response))
        display_conversation.pop(0)
        return display_conversation
    except Exception as e:
        return f"An error occurred: {e}"


# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Simulator Setup")

    task_input = gr.Textbox(label="Task", placeholder="Describe your task...")
    app_name_input = gr.Textbox(label="App Name", placeholder="Enter the app name...")
    sitemap_input = gr.Textbox(label="Sitemap", placeholder="Enter the Hugging Face link to sitemap...")

    initialize_button = gr.Button("Initialize Simulator")
    #setup_output = gr.Textbox(label="Setup Status", interactive=False)

    chatbot = gr.Chatbot(label="Simulator Chat", height=800)
    user_message = gr.Textbox(label="Enter your message", placeholder="Type your message here...")
    submit_button = gr.Button("Send")

    # Initialize simulator and display the welcome message in chatbot
    initialize_button.click(
        initialize_simulator,
        inputs=[task_input, app_name_input, sitemap_input],
        outputs=chatbot  # Show setup message in the chatbot
    )

    # Handle conversation
    submit_button.click(
        chatbot_interaction,
        inputs=user_message,
        outputs=chatbot
    )

# Launch the app
demo.launch()