import gradio as gr
from chatbot_simulator import ChatbotSimulation
from datasets import load_dataset
import json_repair
import random
import os
import re
import firebase_admin
from firebase_admin import credentials, firestore
import json
os.environ["TOKENIZERS_PARALLELISM"] = "false"

cred = credentials.Certificate(json.loads(os.getenv("Firebase_JSON")))
firebase_admin.initialize_app(cred)
db = firestore.client()

openai_api_key = os.getenv("OPENAI_API_KEY")


def find_random_incomplete_task(app_name):

    collection_ref = db.collection(app_name)
    incomplete_indices = []

    for doc in collection_ref.stream():
        task_data = doc.to_dict()
        if task_data['task_completed'] is None and task_data['task_completed_steps'] is None:
            _, idx = doc.id.split('_')
            idx = int(idx)
            incomplete_indices.append(idx)

    # Check if any incomplete tasks found
    if not incomplete_indices:
        return None

    # Pick a random incomplete index from the list
    random_idx = random.choice(incomplete_indices)
    return random_idx


def write_task_data(app_name, idx, task, task_complete, task_completed_step, trajectory):
    doc_ref = db.collection(app_name).document(f"{app_name}_{idx}")
    doc_ref.set({
        "task": task,
        "task_completed": task_complete,
        "task_completed_steps": task_completed_step,
        "trajectory": trajectory
    })


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

    def initialize_simulator(self, sitemap_url, progress=gr.Progress(track_tqdm=True)):
        """Initialize the simulator with retries and elapsed time tracking."""
        synthetic_sitemap = load_dataset(sitemap_url, "schema", split='train')
        app_name, app_description, sitemap, relevant_tables_per_page, jinjia_prerender_page = None, None, None, None, None

        for row in synthetic_sitemap:
            if row['name'] == 'app_name':
                app_name = row['value']
            elif row['name'] == 'app_description':
                app_description = row['value']
            elif row['name'] == 'sitemap':
                sitemap = json_repair.loads(row['value'])
            elif row['name'] == 'relevant_tables_per_page':
                relevant_tables_per_page = json_repair.loads(row['value'])
            elif row['name'] == 'jinjia_prerender_pages':
                jinjia_prerender_page = json_repair.loads(row['value'])

        self.app_name = app_name
        smallest_index = find_random_incomplete_task(app_name)
        if smallest_index is None:
            return "All tasks in this app have been completed!"
        self.smallest_index = smallest_index

        synthetic_tasks = load_dataset(sitemap_url, "tasks", split='train')
        random_index = random.randint(0, len(synthetic_tasks) - 1)  # Generate a random index
        random_row = synthetic_tasks[random_index]
        task = random_row['task']
        solution = random_row['solution']
        database = random_row['database']

        self.task = task

        self.simulation = ChatbotSimulation(
            app_name=app_name,
            app_description=app_description,
            site_map=sitemap,
            relevant_tables_per_page=relevant_tables_per_page,
            database=database,
            jinjia_prerender_page=jinjia_prerender_page,
            task=task,
            solution=solution,
            log_location=f'conversation_log_{app_name}.txt',
            openai_api_key=openai_api_key,
            agent='Human'
        )

        initial_message = self.simulation.start_conversation()
        progress.update("Initialization Successful")
        return initial_message  # Return the initial assistant message for chat

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


def chat(user_input, history, simulator_app):
    """Chat handler that validates input and interacts with the simulator."""
    response = simulator_app.chat_interaction(user_input, history)

    # Initialize variables for task completion and steps

    # Define the pattern for matching the response
    pattern = r"Task completed! You took (\d+) steps\."
    match = re.match(pattern, response)

    if match:
        task_complete = 1
        task_completed_step = int(match.group(1))
        app_name = simulator_app.app_name
        idx = simulator_app.smallest_index
        task = simulator_app.task
        trajectory = simulator_app.simulation.trajectory

        write_task_data(app_name, idx, task, task_complete, task_completed_step, trajectory)

    return response


def give_up(simulator_app):
    """Handle the Give-Up action by marking the first incomplete task as abandoned."""
    task_completed = 0
    task_completed_steps = 0

    app_name = simulator_app.app_name
    idx = simulator_app.smallest_index
    task = simulator_app.task
    trajectory = simulator_app.simulation.trajectory

    write_task_data(app_name, idx, task, task_completed, task_completed_steps, trajectory)

    return "Task marked as abandoned (Give-Up action)."


# Gradio Interface using ChatInterface
with gr.Blocks(fill_height=True) as demo:
    # Simulator instance stored in gr.State for each session
    simulator_state = gr.State(AppSimulator(openai_api_key=openai_api_key))

    gr.Markdown("## Simulator Setup")

    # Input fields for initialization
    sitemap_input = gr.Textbox(label="Sitemap", placeholder="Enter the Hugging Face link to sitemap... (eg.jjz5463/AppStore_synthetic_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",
        additional_inputs=[simulator_state]
    )

    give_up_button = gr.Button("Give Up")

    # Define the callback function to initialize the simulator and update status
    def initialize_and_start_chat(sitemap, simulator_app):
        return simulator_app.initialize_simulator(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=[sitemap_input, simulator_state],
        outputs=status  # Update only the status block
    )

    # Set up the Give-Up button click to update the dataset
    give_up_button.click(
        fn=give_up,
        inputs=[simulator_state],
        outputs=status  # Update the status with the give-up message
    )

# Launch the app
demo.launch()