import os
import subprocess
import random
from typing import List, Dict, Tuple
from datetime import datetime
import logging

import gradio as gr
from huggingface_hub import InferenceClient

# --- Configuration ---
MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"  # Hugging Face model for text generation
MAX_HISTORY_TURNS = 5  # Number of previous turns to include in the prompt
MAX_TOKENS_PER_TURN = 2048  # Maximum number of tokens to generate per turn
VERBOSE_LOGGING = True  # Enable verbose logging for debugging
DEFAULT_AGENT = "WEB_DEV"  # Default agent to use

# --- Logging Setup ---
logging.basicConfig(
    filename="app.log",  # Name of the log file
    level=logging.INFO,  # Set the logging level (INFO, DEBUG, etc.)
    format="%(asctime)s - %(levelname)s - %(message)s",
)

# --- Agent Definitions ---
class Agent:
    """Base class for all agents."""

    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description

    def handle_action(self, action: str, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles an action from the user.

        Args:
            action: The action name.
            action_input: The input for the action.
            history: The conversation history.
            task: The current task.

        Returns:
            A tuple containing the next action name, action input, updated history, and updated task.
        """
        raise NotImplementedError("Agent subclasses must implement handle_action.")

    def get_prompt(self, message: str, history: List[Tuple[str, str]], task: str) -> str:
        """Generates a prompt for the language model.

        Args:
            message: The user's message.
            history: The conversation history.
            task: The current task.

        Returns:
            The prompt string.
        """
        now = datetime.now()
        date_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
        prompt = f"""
        {date_time_str}
        Agent: {self.name}
        Task: {task}
        History:
        {self.format_history(history)}
        Message: {message}
        """
        return prompt

    def format_history(self, history: List[Tuple[str, str]]) -> str:
        """Formats the conversation history for the prompt."""
        formatted_history = ""
        for user_message, agent_response in history[-MAX_HISTORY_TURNS:]:
            formatted_history += f"[INST] {user_message} [/INST]\n{agent_response}\n"
        return formatted_history

class WebDevAgent(Agent):
    """Agent for web development tasks."""

    def __init__(self):
        super().__init__(name="WEB_DEV", description="Agent specialized in web development tasks.")

    def handle_action(self, action: str, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        if action == "SEARCH":
            return self._handle_search_action(action_input, history, task)
        elif action == "GENERATE_HTML":
            return self._handle_generate_html_action(action_input, history, task)
        elif action == "GENERATE_CSS":
            return self._handle_generate_css_action(action_input, history, task)
        elif action == "GENERATE_JS":
            return self._handle_generate_js_action(action_input, history, task)
        elif action == "COMPLETE":
            return "COMPLETE", "COMPLETE", history, task
        else:
            return "MAIN", None, history, task

    def _handle_search_action(self, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles the SEARCH action."""
        if VERBOSE_LOGGING:
            logging.info(f"Calling SEARCH action with input: {action_input}")
        try:
            if "http" in action_input:
                if "<" in action_input:
                    action_input = action_input.strip("<")
                if ">" in action_input:
                    action_input = action_input.strip(">")
                response = i_s(action_input)  # Use i_search for web search
                history.append(("observation: search result is:", response))
            else:
                history.append(("observation: I need a valid URL for the SEARCH action.", ""))
        except Exception as e:
            history.append(("observation:", str(e)))
        return "MAIN", None, history, task

    def _handle_generate_html_action(self, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles the GENERATE_HTML action."""
        if VERBOSE_LOGGING:
            logging.info(f"Calling GENERATE_HTML action with input: {action_input}")
        # Simulate OpenAI's code generation capabilities using Hugging Face
        prompt = self.get_prompt(f"Generate HTML code for a web page that {action_input}", history, task)
        response = run_gpt(prompt, stop_tokens=["```", "```html"], max_tokens=MAX_TOKENS_PER_TURN)
        history.append(("observation: generated HTML code:", response))
        return "MAIN", None, history, task

    def _handle_generate_css_action(self, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles the GENERATE_CSS action."""
        if VERBOSE_LOGGING:
            logging.info(f"Calling GENERATE_CSS action with input: {action_input}")
        # Simulate OpenAI's code generation capabilities using Hugging Face
        prompt = self.get_prompt(f"Generate CSS code for a web page that {action_input}", history, task)
        response = run_gpt(prompt, stop_tokens=["```", "```css"], max_tokens=MAX_TOKENS_PER_TURN)
        history.append(("observation: generated CSS code:", response))
        return "MAIN", None, history, task

    def _handle_generate_js_action(self, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles the GENERATE_JS action."""
        if VERBOSE_LOGGING:
            logging.info(f"Calling GENERATE_JS action with input: {action_input}")
        # Simulate OpenAI's code generation capabilities using Hugging Face
        prompt = self.get_prompt(f"Generate JavaScript code for a web page that {action_input}", history, task)
        response = run_gpt(prompt, stop_tokens=["```", "```js"], max_tokens=MAX_TOKENS_PER_TURN)
        history.append(("observation: generated JavaScript code:", response))
        return "MAIN", None, history, task

class AiSystemPromptAgent(Agent):
    """Agent for generating system prompts."""

    def __init__(self):
        super().__init__(name="AI_SYSTEM_PROMPT", description="Agent specialized in generating system prompts.")

    def handle_action(self, action: str, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        if action == "GENERATE_PROMPT":
            return self._handle_generate_prompt_action(action_input, history, task)
        elif action == "COMPLETE":
            return "COMPLETE", "COMPLETE", history, task
        else:
            return "MAIN", None, history, task

    def _handle_generate_prompt_action(self, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles the GENERATE_PROMPT action."""
        if VERBOSE_LOGGING:
            logging.info(f"Calling GENERATE_PROMPT action with input: {action_input}")
        # Simulate OpenAI's prompt generation capabilities using Hugging Face
        prompt = self.get_prompt(f"Generate a system prompt for a language model that {action_input}", history, task)
        response = run_gpt(prompt, stop_tokens=["```", "```json"], max_tokens=MAX_TOKENS_PER_TURN)
        history.append(("observation: generated system prompt:", response))
        return "MAIN", None, history, task

class PythonCodeDevAgent(Agent):
    """Agent for Python code development tasks."""

    def __init__(self):
        super().__init__(name="PYTHON_CODE_DEV", description="Agent specialized in Python code development tasks.")

    def handle_action(self, action: str, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        if action == "GENERATE_CODE":
            return self._handle_generate_code_action(action_input, history, task)
        elif action == "RUN_CODE":
            return self._handle_run_code_action(action_input, history, task)
        elif action == "COMPLETE":
            return "COMPLETE", "COMPLETE", history, task
        else:
            return "MAIN", None, history, task

    def _handle_generate_code_action(self, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles the GENERATE_CODE action."""
        if VERBOSE_LOGGING:
            logging.info(f"Calling GENERATE_CODE action with input: {action_input}")
        # Simulate OpenAI's code generation capabilities using Hugging Face
        prompt = self.get_prompt(f"Generate Python code that {action_input}", history, task)
        response = run_gpt(prompt, stop_tokens=["```", "```python"], max_tokens=MAX_TOKENS_PER_TURN)
        history.append(("observation: generated Python code:", response))
        return "MAIN", None, history, task

    def _handle_run_code_action(self, action_input: str, history: List[Tuple[str, str]], task: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
        """Handles the RUN_CODE action."""
        if VERBOSE_LOGGING:
            logging.info(f"Calling RUN_CODE action with input: {action_input}")
        # Simulate OpenAI's code execution capabilities using Hugging Face
        prompt = self.get_prompt(f"Run the following Python code and provide the output: {action_input}", history, task)
        response = run_gpt(prompt, stop_tokens=["```", "```python"], max_tokens=MAX_TOKENS_PER_TURN)
        history.append(("observation: code output:", response))
        return "MAIN", None, history, task

# --- Action Handlers ---
def handle_main_action(action: str, action_input: str, history: List[Tuple[str, str]], task: str, agent: Agent) -> Tuple[str, str, List[Tuple[str, str]], str]:
    """Handles the MAIN action, which is the default action."""
    if VERBOSE_LOGGING:
        logging.info(f"Calling MAIN action with input: {action_input}")
    prompt = agent.get_prompt(action_input, history, task)
    response = run_gpt(prompt, stop_tokens=["observation:", "task:", "action:", "thought:"], max_tokens=MAX_TOKENS_PER_TURN)
    if VERBOSE_LOGGING:
        logging.info(f"Response from model: {response}")
    history.append((action_input, response))
    lines = response.strip().strip("\n").split("\n")
    for line in lines:
        if line == "":
            continue
        if line.startswith("thought: "):
            history.append((line, ""))
            if VERBOSE_LOGGING:
                logging.info(f"Thought: {line}")
        elif line.startswith("action: "):
            action_name, action_input = parse_action(line)
            history.append((line, ""))
            if VERBOSE_LOGGING:
                logging.info(f"Action: {action_name} - {action_input}")
            if "COMPLETE" in action_name or "COMPLETE" in action_input:
                task = "END"
                return action_name, action_input, history, task
            else:
                return action_name, action_input, history, task
        else:
            history.append((line, ""))
            if VERBOSE_LOGGING:
                logging.info(f"Other Output: {line}")
    return "MAIN", None, history, task

def handle_update_task_action(action: str, action_input: str, history: List[Tuple[str, str]], task: str, agent: Agent) -> Tuple[str, str, List[Tuple[str, str]], str]:
    """Handles the UPDATE-TASK action, which updates the current task."""
    if VERBOSE_LOGGING:
        logging.info(f"Calling UPDATE-TASK action with input: {action_input}")
    prompt = agent.get_prompt(action_input, history, task)
    task = run_gpt(prompt, stop_tokens=[], max_tokens=64).strip("\n")
    history.append(("observation: task has been updated to:", task))
    return "MAIN", None, history, task

def handle_search_action(action: str, action_input: str, history: List[Tuple[str, str]], task: str, agent: Agent) -> Tuple[str, str, List[Tuple[str, str]], str]:
    """Handles the SEARCH action, which performs a web search."""
    if VERBOSE_LOGGING:
        logging.info(f"Calling SEARCH action with input: {action_input}")
    try:
        if "http" in action_input:
            if "<" in action_input:
                action_input = action_input.strip("<")
            if ">" in action_input:
                action_input = action_input.strip(">")
            response = i_s(action_input)  # Use i_search for web search
            history.append(("observation: search result is:", response))
        else:
            history.append(("observation: I need a valid URL for the SEARCH action.", ""))
    except Exception as e:
        history.append(("observation:", str(e)))
    return "MAIN", None, history, task

def handle_complete_action(action: str, action_input: str, history: List[Tuple[str, str]], task: str, agent: Agent) -> Tuple[str, str, List[Tuple[str, str]], str]:
    """Handles the COMPLETE action, which ends the current task."""
    if VERBOSE_LOGGING:
        logging.info(f"Calling COMPLETE action.")
    task = "END"
    return "COMPLETE", "COMPLETE", history, task

# --- Action Mapping ---
ACTION_HANDLERS: Dict[str, callable] = {
    "MAIN": handle_main_action,
    "UPDATE-TASK": handle_update_task_action,
    "SEARCH": handle_search_action,
    "COMPLETE": handle_complete_action,
}

# --- Utility Functions ---
def run_gpt(prompt: str, stop_tokens: List[str], max_tokens: int) -> str:
    """Runs the language model and returns the generated text."""
    if VERBOSE_LOGGING:
        logging.info(f"Prompt: {prompt}")
    client = InferenceClient(MODEL_NAME)
    resp = client.text_generation(prompt, max_new_tokens=max_tokens, stop_sequences=stop_tokens, temperature=0.7, top_p=0.8, repetition_penalty=1.5)
    if VERBOSE_LOGGING:
        logging.info(f"Response: {resp}")
    return resp

def parse_action(line: str) -> Tuple[str, str]:
    """Parses an action line to get the action name and input."""
    parts = line.split(":", 1)
    if len(parts) == 2:
        action_name = parts[0].replace("action", "").strip()
        action_input = parts[1].strip()
    else:
        action_name = parts[0].replace("action", "").strip()
        action_input = ""
    return action_name, action_input

def run_agent(purpose: str, history: List[Tuple[str, str]], agent: Agent) -> List[Tuple[str, str]]:
    """Runs the agent and returns the updated conversation history."""
    task = None
    directory = "./"
    action_name = "UPDATE-TASK" if task is None else "MAIN"
    action_input = None
    while True:
        if VERBOSE_LOGGING:
            logging.info(f"---")
            logging.info(f"Purpose: {purpose}")
            logging.info(f"Task: {task}")
            logging.info(f"---")
            logging.info(f"History: {history}")
            logging.info(f"---")
        if VERBOSE_LOGGING:
            logging.info(f"Running action: {action_name} - {action_input}")
        try:
            if "RESPONSE" in action_name or "COMPLETE" in action_name:
                action_name = "COMPLETE"
                task = "END"
                return history
            if action_name not in ACTION_HANDLERS:
                action_name = "MAIN"
            if action_name == "" or action_name is None:
                action_name = "MAIN"
            action_handler = ACTION_HANDLERS[action_name]
            action_name, action_input, history, task = action_handler(action_name, action_input, history, task, agent)
            yield history
            if task == "END":
                return history
        except Exception as e:
            history.append(("observation: the previous command did not produce any useful output, I need to check the commands syntax, or use a different command", ""))
            logging.error(f"Error in run_agent: {e}")
            return history

# --- Gradio Interface ---
def main():
    with gr.Blocks() as demo:
        gr.Markdown("## FragMixt: Your No-Code Development Powerhouse")
        gr.Markdown("### Agents w/ Agents:  Mastering No-Code Development")

        # Chat Interface
        chatbot = gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel")

        # Input Components
        message = gr.Textbox(label="Enter your message", placeholder="Ask me anything!")
        purpose = gr.Textbox(label="Purpose", placeholder="What is the purpose of this interaction?")
        agent_name = gr.Dropdown(label="Agents", choices=[agent.name for agent in [WebDevAgent(), AiSystemPromptAgent(), PythonCodeDevAgent()]], value=DEFAULT_AGENT, interactive=True)
        sys_prompt = gr.Textbox(label="System Prompt", max_lines=1, interactive=True)
        temperature = gr.Slider(label="Temperature", value=0.9, minimum=0.0, maximum=1.0, step=0.05, interactive=True, info="Higher values produce more diverse outputs")
        max_new_tokens = gr.Slider(label="Max new tokens", value=1048*10, minimum=0, maximum=1048*10, step=64, interactive=True, info="The maximum numbers of new tokens")
        top_p = gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.0, maximum=1, step=0.05, interactive=True, info="Higher values sample more low-probability tokens")
        repetition_penalty = gr.Slider(label="Repetition penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05, interactive=True, info="Penalize repeated tokens")

        # Button to submit the message
        submit_button = gr.Button(value="Send")

        # Project Explorer Tab
        with gr.Tab("Project Explorer"):
            project_path = gr.Textbox(label="Project Path", placeholder="/home/user/app/current_project")
            explore_button = gr.Button(value="Explore")
            project_output = gr.Textbox(label="File Tree", lines=20)

        # Chat App Logic Tab
        with gr.Tab("Chat App"):
            history = gr.State([])
            examples = [
                ["What is the purpose of this AI agent?", "I am designed to assist with no-code development tasks."],
                ["Can you help me generate a Python function to calculate the factorial of a number?", "Sure! Here is a Python function to calculate the factorial of a number:"],
                ["Generate a web page with a heading that says 'Welcome to My Website!'", "action: GENERATE_HTML action_input=a heading that says 'Welcome to My Website!'"],
            ]

        def chat(purpose, message, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, history):
            if agent_name == "WEB_DEV":
                agent = WebDevAgent()
            elif agent_name == "AI_SYSTEM_PROMPT":
                agent = AiSystemPromptAgent()
            elif agent_name == "PYTHON_CODE_DEV":
                agent = PythonCodeDevAgent()
            else:
                agent = WebDevAgent()  # Default to WEB_DEV if agent_name is invalid
            history = list(run_agent(purpose, history, agent))
            return history, history

        submit_button.click(chat, inputs=[purpose, message, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, history], outputs=[chatbot, history])

    demo.launch()

if __name__ == "__main__":
    main()