import gradio as gr
import os
import shutil
import json
import logging
import utils
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")

# 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():
    """Get Hugging Face API key input."""
    return gr.Textbox(label="Hugging Face API Key", type="password")

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

def get_ai_role():
    """Get AI role input."""
    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_description():
    """Get project description input."""
    return gr.Textbox(label="Description", placeholder="Enter a brief description of the project.")

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

def get_example_values():
    """Load example values from JSON file."""
    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():
    """Get chatbot UI element."""
    return gr.Chatbot(elem_id="chatbot", type='messages')

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

def get_consecutive_yes():
    """Get slider for consecutive yes count."""
    return gr.Slider(1, 10, 1, step=1, label="Consecutive Yes", interactive=False)

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

def get_progress():
    """Get progress bar."""
    return gr.Progress()

def get_generated_files():
    """Get HTML element to display 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():
    """Get download all files button."""
    return gr.Button("Download All Files")

def get_inferred_tasks():
    """Get inferred tasks textbox."""
    return gr.Textbox(label="Inferred Tasks", interactive=False)

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

    def infer_tasks(self, description):
        # Placeholder for actual task inference logic
        # Simulate task inference based on the description
        tasks = []
        
        # Define keywords and corresponding tasks
        keyword_tasks = {
            "business": ["Analyze market trends", "Create business plan"],
            "technology": ["Research latest technology", "Prototype development"],
            "startup": ["Identify target audience", "Develop marketing strategy"],
            "product": ["Design product", "Test product"],
            "finance": ["Budget planning", "Financial forecasting"],
            "team": ["Recruit team members", "Team building activities"],
            "strategy": ["Develop strategic plan", "Set milestones"]
        }
        
        # Split the description into words and check for keywords
        words = description.lower().split()
        for keyword, task_list in keyword_tasks.items():
            if any(keyword in word for word in words):
                tasks.extend(task_list)
        
        # Ensure the list always has 5 tasks
        while len(tasks) < 5:
            tasks.append(f"Generic Task {len(tasks) + 1}")
        
        return tasks

def start(huggingface_key, ai_name, ai_role, top_5_goals, description):
    """Start AutoAPI and infer tasks."""
    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)
        
        # Infer tasks based on the role and description
        tasks = auto_api.infer_tasks(description)
        logger.info("Inferred tasks: %s", tasks)
        
        return gr.Column.update(visible=False), gr.Column.update(visible=True), auto_api, gr.update(value=tasks)
    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, 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()