import os
import subprocess
import random
import json
from datetime import datetime
from huggingface_hub import InferenceClient, cached_download, hf_hub_url
import gradio as gr
from safe_search import safe_search
from i_search import google, i_search as i_s
from agent import (
    ACTION_PROMPT, ADD_PROMPT, COMPRESS_HISTORY_PROMPT, LOG_PROMPT,
    LOG_RESPONSE, MODIFY_PROMPT, PRE_PREFIX, SEARCH_QUERY, READ_PROMPT,
    TASK_PROMPT, UNDERSTAND_TEST_RESULTS_PROMPT
)
from utils import parse_action, parse_file_content, read_python_module_structure

# Global Variables
terminal_history = ""

# Component Library
components_registry = {
    "Button": {
        "properties": {"label": "Click Me", "onclick": ""},
        "description": "A clickable button",
        "code_snippet": 'gr.Button(value="{label}", variant="primary")',
    },
    "Text Input": {
        "properties": {"value": "", "placeholder": "Enter text"},
        "description": "A field for entering text",
        "code_snippet": 'gr.Textbox(label="{placeholder}")',
    },
    "Image": {
        "properties": {"src": "#", "alt": "Image"},
        "description": "Displays an image",
        "code_snippet": 'gr.Image(label="{alt}")',
    },
    "Dropdown": {
        "properties": {"choices": ["Option 1", "Option 2"], "value": ""},
        "description": "A dropdown menu for selecting options",
        "code_snippet": 'gr.Dropdown(choices={choices}, label="Dropdown")',
    },
    # Add more components here...
}

# NLP Model (Example using Hugging Face)
nlp_model_names = [
    "google/flan-t5-small",
    "Qwen/CodeQwen1.5-7B-Chat-GGUF",
    "bartowski/Codestral-22B-v0.1-GGUF",
    "bartowski/AutoCoder-GGUF"
]
nlp_models = []

for nlp_model_name in nlp_model_names:
    try:
        cached_download(hf_hub_url(nlp_model_name, revision="main"))
        nlp_models.append(InferenceClient(nlp_model_name))
    except:
        nlp_models.append(None)

# Function to get NLP model response
def get_nlp_response(input_text, model_index):
    if nlp_models[model_index]:
        response = nlp_models[model_index].text_generation(input_text)
        return response.generated_text
    else:
        return "NLP model not available."

# Component Class
class Component:
    def __init__(self, type, properties=None, id=None):
        self.id = id or random.randint(1000, 9999)
        self.type = type
        self.properties = properties or components_registry[type]["properties"].copy()

    def to_dict(self):
        return {
            "id": self.id,
            "type": self.type,
            "properties": self.properties,
        }

    def render(self):
        # Properly format choices for Dropdown
        if self.type == "Dropdown":
            self.properties["choices"] = (
                str(self.properties["choices"])
                .replace("[", "")
                .replace("]", "")
                .replace("'", "")
            )
        return components_registry[self.type]["code_snippet"].format(**self.properties)

# App Creation Process Class
class AppCreationProcess:
    def __init__(self):
        self.current_step = 1
        self.app_name = ""
        self.components = []

    def next_step(self):
        self.current_step += 1

    def previous_step(self):
        if self.current_step > 1:
            self.current_step -= 1

    def get_current_step_info(self):
        steps = {
            1: "App Initialization",
            2: "Component Addition",
            3: "Property Configuration",
            4: "Code Generation",
            5: "Deployment"
        }
        return f"Step {self.current_step}: {steps[self.current_step]}"

    def add_component(self, component_type):
        new_component = Component(component_type)
        self.components.append(new_component.to_dict())
        return self.update_app_canvas()

    def set_component_property(self, component_id, property_name, property_value):
        for component in self.components:
            if component['id'] == component_id:
                if property_name in component['properties']:
                    component['properties'][property_name.strip()] = property_value.strip()
                    return self.update_app_canvas(), f"Property '{property_name}' set to '{property_value}' for component {component_id}"
                else:
                    return None, f"Error: Property '{property_name}' not found in component {component_id}"
        return None, f"Error: Component with ID {component_id} not found."

    def update_app_canvas(self):
        components_html = "".join([
            f"<div>Component ID: {component['id']}, Type: {component['type']}, Properties: {component['properties']}</div>"
            for component in self.components
        ])
        return components_html

    def generate_python_code(self):
        code = f"""import gradio as gr\n\nwith gr.Blocks() as {self.app_name}:\n"""
        for component in self.components:
            code += "    " + Component(**component).render() + "\n"
        code += f"\n{self.app_name}.launch()\n"
        return code

    def deploy_to_huggingface(self):
        # Generate Python code
        code = self.generate_python_code()
        # Create requirements.txt
        with open("requirements.txt", "w") as f:
            f.write("gradio==3.32.0\n")
        # Create the app.py file
        with open("app.py", "w") as f:
            f.write(code)
        # Execute the deployment command
        try:
            subprocess.run(["huggingface-cli", "repo", "create", "--type", "space", "--space_sdk", "gradio", self.app_name], check=True)
            subprocess.run(["git", "init"], cwd=f"./{self.app_name}", check=True)
            subprocess.run(["git", "add", "."], cwd=f"./{self.app_name}", check=True)
            subprocess.run(["git", "commit", "-m", "Initial commit"], cwd=f"./{self.app_name}", check=True)
            subprocess.run(["git", "push", "https://huggingface.co/spaces/" + self.app_name, "main"], cwd=f"./{self.app_name}", check=True)
            return f"Successfully deployed to Hugging Face Spaces: https://huggingface.co/spaces/{self.app_name}"
        except Exception as e:
            return f"Error deploying to Hugging Face Spaces: {e}"

app_process = AppCreationProcess()

# Function to handle terminal input
def run_terminal_command(command, history):
    global terminal_history
    output = ""
    try:
        # Basic command parsing (expand with NLP)
        if command.startswith("add "):
            component_type = command.split("add ", 1)[1].strip()
            output = app_process.add_component(component_type)
        elif command.startswith("set "):
            _, output = set_component_property(command)
        elif command.startswith("search "):
            search_query = command.split("search ", 1)[1].strip()
            output = i_s(search_query)
        elif command.startswith("deploy "):
            output = app_process.deploy_to_huggingface()
        else:
            # Attempt to execute command as Python code
            try:
                result = subprocess.check_output(
                    command, shell=True, stderr=subprocess.STDOUT, text=True
                )
                output = result
            except Exception as e:
                output = f"Error executing Python code: {str(e)}"
    except Exception as e:
        output = f"Error: {str(e)}"
    finally:
        terminal_history += f"User: {command}\n{output}\n"
    return terminal_history

def set_component_property(command):
    try:
        # Improved 'set' command parsing
        set_parts = command.split(" ", 2)[1:]
        if len(set_parts) != 2:
            raise ValueError("Invalid 'set' command format.")
        component_id = int(set_parts[0])  # Use component ID
        property_name, property_value = set_parts[1].split("=", 1)
        return app_process.set_component_property(component_id, property_name, property_value)
    except Exception as e:
        return None, f"Error: {str(e)}\n"

# Function to handle chat interaction
def run_chat(message, history):
    global terminal_history
    if message.startswith("!"):
        command = message[1:]
        terminal_history = run_terminal_command(command, history)
    else:
        model_index = 0  # Select the model to use for chat response
        response = get_nlp_response(message, model_index)
        if response:
            return history, terminal_history + f"User: {message}\nAssistant: {response}"
        else:
            return history, terminal_history + f"User: {message}\nAssistant: I'm sorry, I couldn't generate a response. Please try again.\n"

# Gradio Interface
with gr.Blocks() as iface:
    gr.Markdown("# Sequential App Builder")
    
    with gr.Row():
        current_step = gr.Markdown(app_process.get_current_step_info())
    
    with gr.Row():
        prev_button = gr.Button("Previous Step")
        next_button = gr.Button("Next Step")

    # Step 1: App Initialization
    with gr.Group() as step1:
        app_name_input = gr.Textbox(label="Enter App Name")
        init_app_button = gr.Button("Initialize App")

    # Step 2: Component Addition
    with gr.Group() as step2:
        component_type = gr.Dropdown(choices=list(components_registry.keys()), label="Select Component Type")
        add_component_button = gr.Button("Add Component")
        components_display = gr.HTML()

    # Step 3: Property Configuration
    with gr.Group() as step3:
        component_id = gr.Number(label="Component ID")
        property_name = gr.Textbox(label="Property Name")
        property_value = gr.Textbox(label="Property Value")
        set_property_button = gr.Button("Set Property")

    # Step 4: Code Generation
    with gr.Group() as step4:
        generated_code = gr.Code(language="python")
        generate_code_button = gr.Button("Generate Code")

    # Step 5: Deployment
    with gr.Group() as step5:
        deploy_button = gr.Button("Deploy to Hugging Face Spaces")
        deployment_status = gr.Markdown()

    # Chat and Terminal (optional, can be hidden or shown based on preference)
    with gr.Accordion("Advanced", open=False):
        chat_history = gr.Chatbot(label="Chat with Agent")
        chat_input = gr.Textbox(label="Your Message")
        chat_button = gr.Button("Send")
        
        terminal_output = gr.Textbox(lines=8, label="Terminal", value=terminal_history)
        terminal_input = gr.Textbox(label="Enter Command")
        terminal_button = gr.Button("Run")

    # Function to update visibility based on current step
    def update_visibility(step):
        return {
            step1: gr.update(visible=(step == 1)),
            step2: gr.update(visible=(step == 2)),
            step3: gr.update(visible=(step == 3)),
            step4: gr.update(visible=(step == 4)),
            step5: gr.update(visible=(step == 5)),
        }

    # Event handlers
    def next_step():
        app_process.next_step()
        return app_process.get_current_step_info(), update_visibility(app_process.current_step)

    def prev_step():
        app_process.previous_step()
        return app_process.get_current_step_info(), update_visibility(app_process.current_step)

    next_button.click(next_step, outputs=[current_step, step1, step2, step3, step4, step5])
    prev_button.click(prev_step, outputs=[current_step, step1, step2, step3, step4, step5])

    # Step 1: Initialize App
    def init_app(name):
        app_process.app_name = name
        return f"App '{name}' initialized."

    init_app_button.click(init_app, inputs=[app_name_input], outputs=[components_display])

    # Step 2: Add Component
    add_component_button.click(app_process.add_component, inputs=[component_type], outputs=[components_display])

    # Step 3: Set Property
    set_property_button.click(app_process.set_component_property, inputs=[component_id, property_name, property_value], outputs=[components_display])

    # Step 4: Generate Code
    generate_code_button.click(app_process.generate_python_code, outputs=[generated_code])

    # Step 5: Deploy
    deploy_button.click(app_process.deploy_to_huggingface, outputs=[deployment_status])

    # Existing chat and terminal functionality
    chat_button.click(run_chat, inputs=[chat_input, chat_history], outputs=[chat_history, terminal_output])
    terminal_button.click(run_terminal_command, inputs=[terminal_input, terminal_output], outputs=[terminal_output])

iface.launch()