import gradio as gr
from chatbot_simulator import ChatbotSimulation
from huggingface_hub import HfApi, create_repo
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_smallest_incomplete_task(app_name):

    collection_ref = db.collection(app_name)
    docs = collection_ref.stream()

    smallest_incomplete_idx = None
    for doc in docs:
        doc_id = doc.id
        _, idx = doc_id.split('_')
        idx = int(idx)

        task_data = doc.to_dict()
        if task_data['task_completed'] is None and task_data['task_completed_steps'] is None:
            if smallest_incomplete_idx is None or idx < smallest_incomplete_idx:
                smallest_incomplete_idx = idx

    return smallest_incomplete_idx


def write_task_data(app_name, idx, task, task_complete, task_completed_step):
    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
    })


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, "sitemap", split='train')
        app_name, sitemap, page_details, user_state, system_data = None, None, None, None, None
        for row in synthetic_sitemap:
            if row['name'] == 'app_name':
                app_name = row['value']  # Use `eval` to convert the string to a list
            elif row['name'] == 'sitemap':
                sitemap = json_repair.loads(row['value'])
            elif row['name'] == 'page_details':
                page_details = json_repair.loads(row['value'])
            elif row['name'] == 'user_state':
                user_state = json_repair.loads(row['value'])
            elif row['name'] == 'system_data':
                system_data = json_repair.loads(row['value'])

        self.app_name = app_name
        smallest_index = find_smallest_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')
        incomplete_task = synthetic_tasks[smallest_index]
        task = incomplete_task["tasks"]
        solution = incomplete_task["steps"]
        user_data = incomplete_task["attributes"]["user_data"]

        self.task = task

        self.simulation = ChatbotSimulation(
            app_name=app_name,
            site_map=sitemap,
            page_details=page_details,
            user_state=user_state,
            system_data=system_data,
            user_data=user_data,
            task=task,
            solution=solution,
            log_location=f'conversation_log_{app_name}.txt',
            openai_api_key=openai_api_key,
            agent='llm'
        )

        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)


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


def chat(user_input, history):
    """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
        write_task_data(app_name, idx, task, task_complete, task_completed_step)

    return response


def give_up():
    """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

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

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


# Gradio Interface using ChatInterface
with gr.Blocks(fill_height=True) as demo:
    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')

    give_up_button = gr.Button("Give Up")

    # Define the callback function to initialize the simulator and update status
    def initialize_and_start_chat(sitemap):
        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],
        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=[],
        outputs=status  # Update the status with the give-up message
    )

# Launch the app
demo.launch()