import gradio as gr
import os
import json
import logging
from transformers import pipeline
import utils  # Ensure this import is present

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

# Create output directory if it doesn't exist
if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

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

# UI Components
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.")

def get_description():
    return gr.Textbox(label="Project Description", placeholder="Enter a brief description of the project.")

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_inferred_tasks():
    return gr.Textbox(label="Inferred Tasks", interactive=False)

def get_generated_files():
    """Get HTML element to display generated files."""
    files_list = utils.format_directory(OUTPUT_DIR)  # This function should list files in the directory
    return gr.HTML(f"<h3>Generated Files:</h3><pre><code style='overflow-x: auto'>{files_list}</pre></code>")

def get_download_btn():
    """Get download all files button."""
    return gr.Button("Download All Files", elem_id="download-btn")

class AutoAPI:
    def __init__(self, huggingface_key, ai_name, ai_role, top_5_goals):
        self.huggingface_key = huggingface_key
        self.ai_name = ai_name
        self.ai_role = ai_role
        self.top_5_goals = top_5_goals
        # Replace 'your-model-name' with a valid model identifier
        self.nlp_model = pipeline("text2text-generation", model="your-actual-model-name")

    def infer_tasks(self, description):
        # Use the NLP model to generate tasks based on the description
        tasks = self.nlp_model(description)
        return tasks[0]['generated_text'].split(',')

def start(huggingface_key, ai_name, ai_role, top_5_goals, description):
    try:
        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)
        
        # Infer tasks based on the role and description
        tasks = auto_api.infer_tasks(description)
        logger.info("Inferred tasks: %s", tasks)
        
        return gr.Column(visible=False), gr.Column(visible=True), gr.update(value=tasks)
    except Exception as e:
        logger.error("Failed to start AutoAPI: %s", str(e))
        return gr.Column(visible=True), gr.Column(visible=False), gr.update(value=[])

# Main Gradio Interface
with gr.Blocks(css=CSS) as demo:
    gr.Markdown("# AutoGPT Task Inference")
    
    with gr.Row():
        api_key = get_api_key()
        ai_name = get_ai_name()
        ai_role = get_ai_role()
    
    description = get_description()
    top_5_goals = get_top_5_goals()
    
    start_btn = gr.Button("Start")
    main_pane = gr.Column(visible=False)
    setup_pane = gr.Column(visible=True)
    
    inferred_tasks = get_inferred_tasks()
    
    start_btn.click(
        start,
        inputs=[api_key, ai_name, ai_role, top_5_goals, description],
        outputs=[setup_pane, main_pane, inferred_tasks]
    )

    with main_pane:
        get_generated_files()
        get_download_btn()

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()