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


class AppSimulator:
    def __init__(self, openai_api_key):
        self.simulation = None
        self.conversation = []  # Stores full conversation for logging
        self.display_conversation = []  # Stores last 2 messages for display
        self.openai_api_key = openai_api_key

    def initialize_simulator(self, task, app_name, sitemap):
        """Initialize the simulator with retries in case of failure."""
        success = False
        retry_count = 0
        max_retries = 50

        while not success and retry_count < max_retries:
            try:
                # Process data
                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'
                )

                # Start the conversation and update the logs and display
                text = self.simulation.start_conversation()
                self._log_message("assistant", text)
                self.display_conversation = [('Start Simulator', text)]

                return self.display_conversation
            except Exception as e:
                retry_count += 1
                print(f"Attempt {retry_count}/{max_retries}: An error occurred: {e}. Retrying...")

    def chatbot_interaction(self, user_input):
        """Handle one round of conversation."""
        if self.simulation is None:
            return [("system", "Simulation is not initialized. Please start the simulator.")]

        try:
            # Get the response from the simulator
            response = self.simulation.one_conversation_round(user_input)

            # Log both user input and assistant's response
            self._log_message("user", user_input)
            self._log_message("assistant", response)

            # Update display conversation (keep last 2 entries)
            self.display_conversation.extend([(user_input, response)])
            while len(self.display_conversation) > 2:
                self.display_conversation.pop(0)
            print(len(self.display_conversation))
            return self.display_conversation
        except Exception as e:
            return [("system", f"An error occurred: {e}")]

    def _log_message(self, role, content):
        """Log conversation messages to both memory and file."""
        self.conversation.append({"role": role, "content": content})
        self._write_log_to_file(content)

    def _write_log_to_file(self, content):
        """Append conversation to a log file."""
        log_location = self.simulation.log_location if self.simulation else "conversation_log.txt"
        try:
            with open(log_location, 'a') as f:
                for message in self.conversation:
                    f.write(f"{message['role']}: {message['content']}\n\n")
        except Exception as e:
            print(f"Error logging conversation: {e}")


# JavaScript function to refresh theme to dark mode
js_func = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""

simulator_app = AppSimulator(openai_api_key=openai_api_key)

# Gradio Interface
with gr.Blocks(js=js_func) 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")
    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
    initialize_button.click(
        simulator_app.initialize_simulator,
        inputs=[task_input, app_name_input, sitemap_input],
        outputs=chatbot
    )

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

# Launch the Gradio app
demo.launch()