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

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from huggingface_hub import InferenceClient, cached_download

# --- Configuration ---
VERBOSE = True 
MAX_HISTORY = 5  
MAX_TOKENS = 2048  
TEMPERATURE = 0.7  
TOP_P = 0.8  
REPETITION_PENALTY = 1.5 
MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"
API_KEY = "YOUR_API_KEY" 

# --- Logging Setup ---
logging.basicConfig(
    filename="app.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)

# --- Agents ---
agents = [
    "WEB_DEV",
    "AI_SYSTEM_PROMPT",
    "PYTHON_CODE_DEV",
    "DATA_SCIENCE",
    "UI_UX_DESIGN",
]

# --- Prompts ---
PREFIX = """
{date_time_str}
Purpose: {purpose}
Agent: {agent_name}
"""

LOG_PROMPT = """
PROMPT: {content}
"""

LOG_RESPONSE = """
RESPONSE: {resp}
"""

# --- Functions ---
def format_prompt(message: str, history: List[Tuple[str, str]], max_history_turns: int = 2) -> str:
    prompt = ""
    for user_prompt, bot_response in history[-max_history_turns:]:
        prompt += f"Human: {user_prompt}\nAssistant: {bot_response}\n"
    prompt += f"Human: {message}\nAssistant:"
    return prompt

def generate(
    prompt: str,
    history: List[Tuple[str, str]],
    agent_name: str = agents[0],
    sys_prompt: str = "",
    temperature: float = TEMPERATURE,
    max_new_tokens: int = MAX_TOKENS,
    top_p: float = TOP_P,
    repetition_penalty: float = REPETITION_PENALTY,
) -> str:
    # Load model and tokenizer
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    
    # Create a text generation pipeline
    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
    
    # Prepare the full prompt
    date_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    full_prompt = PREFIX.format(
        date_time_str=date_time_str,
        purpose=sys_prompt,
        agent_name=agent_name
    ) + format_prompt(prompt, history)
    
    if VERBOSE:
        logging.info(LOG_PROMPT.format(content=full_prompt))
    
    # Generate response
    response = generator(
        full_prompt,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True
    )[0]['generated_text']
    
    # Extract the assistant's response
    assistant_response = response.split("Assistant:")[-1].strip()
    
    if VERBOSE:
        logging.info(LOG_RESPONSE.format(resp=assistant_response))
    
    return assistant_response

def main():
    with gr.Blocks() as demo:
        gr.Markdown("## FragMixt: The No-Code Development Powerhouse")
        gr.Markdown("###  Your AI-Powered Development Companion")

        # 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=[s for s in agents], value=agents[0], interactive=True)
        sys_prompt = gr.Textbox(label="System Prompt", max_lines=1, interactive=True)
        temperature = gr.Slider(label="Temperature", value=TEMPERATURE, 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=MAX_TOKENS, 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=TOP_P, 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=REPETITION_PENALTY, 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 simple HTML page with a heading and a paragraph.", "