import gradio as gr
from chatbot_simulator import ChatbotSimulation
from task_specific_data_population import DataPopulation
import re
import os

openai_api_key = os.getenv("OPENAI_API_KEY")


class AppSimulator:
    def __init__(self, openai_api_key):
        self.simulation = None
        self.openai_api_key = openai_api_key

    def initialize_simulator(self, task, app_name, sitemap, progress=gr.Progress(track_tqdm=True)):
        """Initialize the simulator with retries and elapsed time tracking."""
        retry_count = 0
        max_retries = 50

        while retry_count < max_retries:
            try:
                # Display current attempt with progress bar
                progress.update(f"Initializing... Attempt {retry_count + 1}/{max_retries}")

                # Process data and initialize the simulation
                data_population = DataPopulation(api_key=self.openai_api_key)
                sitemap_data, page_details, user_state = data_population.process_data(task, sitemap)

                self.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=self.openai_api_key,
                    agent='human'
                )

                # Successful initialization
                initial_message = self.simulation.start_conversation()
                progress.update("Initialization Successful")
                return initial_message  # Return the initial assistant message for chat
            except Exception as e:
                retry_count += 1
                print(f"Attempt {retry_count}/{max_retries} failed: {e}")

        # If all retries failed
        progress.update("Failed to initialize simulator after multiple retries.")
        return "Initialization failed."  # Error message for chat

    def chat_interaction(self, user_input, history):
        """Handle one round of conversation."""
        return self.simulation.one_conversation_round(user_input)


# Initialize the simulator
simulator_app = AppSimulator(openai_api_key=openai_api_key)

def is_valid_input(user_input):
    """Validate user input format."""
    pattern = r"^\d+\.\s+.*"
    return bool(re.match(pattern, user_input))

def chat(user_input, history):
    """Chat handler that validates input and interacts with the simulator."""
    if is_valid_input(user_input):
        valid_response = simulator_app.chat_interaction(user_input, history)
        return valid_response
    else:
        invalid_response = "Invalid input. Please use the format: <Number>. <Description>: query"
        return invalid_response


# Gradio Interface using ChatInterface
with gr.Blocks(fill_height=True) as demo:
    gr.Markdown("## Simulator Setup")

    # Input fields for initialization
    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")

    # Status block to display initialization progress with elapsed time
    status = gr.Textbox(label="Status", interactive=False)

    # Chat interface to handle user interactions
    chat_interface = gr.ChatInterface(fn=chat, type='messages')

    # Define the callback function to initialize the simulator and update status
    def initialize_and_start_chat(task, app_name, sitemap):
        return simulator_app.initialize_simulator(task, app_name, sitemap)  # Use progress tracking

    # Set up the button click to initialize simulator and update status only
    initialize_button.click(
        fn=initialize_and_start_chat,
        inputs=[task_input, app_name_input, sitemap_input],
        outputs=status  # Update only the status block
    )

# Launch the app
demo.launch()