import gradio as gr
import api
import utils
import os
import shutil
import json
import logging
from transformers import pipeline

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

FILE_DIR = os.path.dirname(os.path.abspath(__file__))
EXAMPLES_PATH = os.path.join(FILE_DIR, 'examples.json')
OUTPUT_DIR = os.path.join(os.path.dirname(FILE_DIR), "auto_gpt_workspace")
if not os.path.exists(OUTPUT_DIR):
    os.mkdir(OUTPUT_DIR)

CSS = """
#chatbot {font-family: monospace;}
#files .generating {display: none;}
#files .min {min-height: 0px;}
"""

def get_api_key():
    return gr.Textbox(label="Hugging Face API Key", type="password")

def get_ai_name():
    return gr.Textbox(label="AI Name", placeholder="e.g. Entrepreneur-GPT")

def get_ai_role():
    return gr.Textbox(label="AI Role", placeholder="e.g. an AI designed to autonomously develop and run businesses with the sole goal of increasing your net worth.")

def get_top_5_goals():
    return gr.Dataframe(row_count=(5, "fixed"), col_count=(1, "fixed"), headers=["AI Goals - Enter up to 5"], type="array")

def get_example_values():
    try:
        with open(EXAMPLES_PATH, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        logger.error(f"Error loading examples: {e}")
        return []

def get_chatbot():
    return gr.Chatbot(elem_id="chatbot", type='messages')  # Specify type as 'messages'

def get_yes_btn():
    return gr.Button("Yes", variant="primary", interactive=False)

def get_consecutive_yes():
    return gr.Slider(1, 10, 1, step=1, label="Consecutive Yes", interactive=False)

def get_custom_response():
    return gr.Textbox(label="Custom Response", placeholder="Press 'Enter' to Submit.", interactive=False)

def get_progress():
    return gr.Progress()  # Corrected to remove the label argument

def get_generated_files():
    return gr.HTML(lambda: f"Generated Files<pre><code style='overflow-x: auto'>{utils.format_directory(OUTPUT_DIR)}</pre></code>", every=3, elem_id="files")
    
def get_download_btn():
    return gr.Button("Download All Files")  # Ensure consistent indentation

def start(huggingface_key, ai_name, ai_role, top_5_goals):
    try:
        from api import AutoAPI
        auto_api = AutoAPI(huggingface_key, ai_name, ai_role, top_5_goals)
        logger.info("AutoAPI started with AI Name: %s, AI Role: %s", ai_name, ai_role)
        return gr.Column.update(visible=False), gr.Column.update(visible=True), auto_api
    except Exception as e:
        logger.error("Failed to start AutoAPI: %s", str(e))
        return gr.Column.update(visible=True), gr.Column.update(visible=False), None

def bot_response(chat, api):
    messages = []
    for message in api.get_chatbot_response():
        messages.append(message)
        chat[-1][1] = "\n".join(messages) + "..."
        yield chat
    chat[-1][1] = "\n".join(messages)
    yield chat

def send_message(count, chat, api, message="Y"):
    if message != "Y":
        count = 1
    for i in range(count):
        chat.append([message, None])
        yield chat, count - i
        api.send_message(message)
        for updated_chat in bot_response(chat, api):
            yield updated_chat, count - i

def activate_inputs():
    return {
        get_yes_btn(): gr.Button.update(interactive=True),
        get_consecutive_yes(): gr.Slider.update(interactive=True),
        get_custom_response(): gr.Textbox.update(interactive=True),
    }

def deactivate_inputs():
    return {
        get_yes_btn(): gr.Button.update(interactive=False),
        get_consecutive_yes(): gr.Slider.update(interactive=False),
        get_custom_response(): gr.Textbox.update(interactive=False),
    }

def download_all_files():
    try:
        shutil.make_archive("outputs", "zip", OUTPUT_DIR)
        logger.info("All files downloaded successfully.")
    except Exception as e:
        logger.error("Failed to download files: %s", str(e))

with gr.Blocks(css=CSS) as app:
    with gr.Column() as setup_pane:
        gr.Markdown(f"""# Auto-GPT
        1. Duplicate this Space: <a href="https://huggingface.co/spaces/{os.getenv('SPACE_ID')}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a> This will **NOT** work without duplication!
        2. Enter your <a href="https://huggingface.co/settings/tokens">Hugging Face API Key</a> below.
        """)
        huggingface_key = get_api_key()
        gr.Markdown(
            "3. Fill the values below, then click 'Start'. There are example values you can load at the bottom of this page."
        )
        ai_name = get_ai_name()
        ai_role = get_ai_role()
        top_5_goals = get_top_5_goals()
        start_btn = gr.Button("Start", variant="primary")
        example_values = get_example_values()
        gr.Examples(
            example_values,
            [ai_name, ai_role, top_5_goals],
        )
    with gr.Column(visible=False) as main_pane:
        with gr.Row():
            with gr.Column(scale=2):
                chatbot = get_chatbot()
                with gr.Row():
                    yes_btn = get_yes_btn()
                    consecutive_yes = get_consecutive_yes()
                custom_response = get_custom_response()
                progress = get_progress()  # Use the corrected progress function
            with gr.Column(scale=1):
                generated_files = get_generated_files()
                download_btn = get_download_btn()
    chat_history = gr.State([[None, None]])
    api = gr.State(None)

    start_btn.click(
        start,
        [huggingface_key, ai_name, ai_role, top_5_goals],
        [setup_pane, main_pane, api],
    ).then(bot_response, [chat_history, api], chatbot).then(
        activate_inputs, None, [yes_btn, consecutive_yes, custom_response]
    )

    yes_btn.click(
        deactivate_inputs, None, [yes_btn, consecutive_yes, custom_response]
    ).then(
        send_message, [consecutive_yes, chat_history, api], [chatbot, consecutive_yes]
    ).then(
        activate_inputs, None, [yes_btn, consecutive_yes, custom_response]
    )

    custom_response.submit(
        deactivate_inputs, None, [yes_btn, consecutive_yes, custom_response]
    ).then(
        send_message,
        [consecutive_yes, chat_history, api, custom_response],
        [chatbot, consecutive_yes],
    ).then(
        activate_inputs, None, [yes_btn, consecutive_yes, custom_response]
    )

    download_btn.click(download_all_files)

app.queue().launch(max_threads=20)