import os
import random
from huggingface_hub import InferenceClient
from datetime import datetime
import yaml
import logging
import gradio as gr
from transformers import pipeline, AutoTokenizer

# Create a directory for logs if it doesn't exist
log_dir = "logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

# Configure logging
logging.basicConfig(
    filename=os.path.join(log_dir, "gradio_log.txt"),
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    filemode="a+",
)

# Use the logger
logger = logging.getLogger(__name__)

# Load custom prompts
try:
    with open('custom_prompts.yaml', 'r') as fp:
        custom_prompts = yaml.safe_load(fp)
except FileNotFoundError:
    custom_prompts = {
        "WEB_DEV": "",
        "AI_SYSTEM_PROMPT": "",
        "PYTHON_CODE_DEV": "",
        "CODE_GENERATION": "",
        "CODE_INTERPRETATION": "",
        "CODE_TRANSLATION": "",
        "CODE_IMPLEMENTATION": ""
    }

for key, val in custom_prompts.items():
    globals()[key] = val

# Define advanced prompts
CODE_GENERATION = """
You are an expert AI code generation assistant. Your task is to generate high-quality, production-ready code based on the given requirements. You should be able to generate code in various programming languages, including Python, JavaScript, Java, C++, and more.
When generating code, follow these guidelines:
1. Understand the requirements thoroughly and ask clarifying questions if needed.
2. Write clean, modular, and maintainable code following best practices and industry standards.
3. Implement proper error handling, input validation, and edge case handling.
4. Optimize the code for performance and scalability when necessary.
5. Provide clear and concise comments to explain the code's functionality and logic.
6. If applicable, suggest and implement testing strategies (unit tests, integration tests, etc.).
7. Ensure the generated code is compatible with the target environment (e.g., web, mobile, desktop).
8. Provide examples or usage instructions if required.
Remember to always prioritize code quality, maintainability, and security. Your generated code should be ready for production use or further development.
"""

CODE_INTERPRETATION = """
You are an expert AI code interpretation assistant. Your task is to analyze and explain existing code in various programming languages, including Python, JavaScript, Java, C++, and more.
When interpreting code, follow these guidelines:
1. Read and understand the code thoroughly, including its functionality, logic, and structure.
2. Identify and explain the purpose of each code block, function, or module.
3. Highlight any potential issues, inefficiencies, or areas for improvement.
4. Suggest refactoring or optimization techniques if applicable.
5. Explain the code's input and output, as well as any dependencies or external libraries used.
6. Provide clear and concise explanations, using code comments or separate documentation.
7. If applicable, explain the testing strategies or methodologies used in the code.
8. Ensure your interpretations are accurate, unbiased, and tailored to the target audience's skill level.
Remember to prioritize clarity, accuracy, and completeness in your code interpretations. Your explanations should help developers understand the code's functionality and potential areas for improvement.
"""

CODE_TRANSLATION = """
You are an expert AI code translation assistant. Your task is to translate code from one programming language to another, ensuring the translated code maintains the original functionality and follows best practices in the target language.
When translating code, follow these guidelines:
1. Understand the original code's functionality, logic, and structure thoroughly.
2. Identify and translate all code elements, including variables, functions, classes, and data structures.
3. Ensure the translated code adheres to the coding conventions and best practices of the target language.
4. Optimize the translated code for performance and readability in the target language.
5. Preserve comments and documentation, translating them to the target language if necessary.
6. Handle any language-specific features or constructs appropriately during the translation process.
7. Implement error handling, input validation, and edge case handling in the translated code.
8. Provide clear and concise comments or documentation to explain any necessary changes or deviations from the original code.
Remember to prioritize accuracy, maintainability, and idiomatic usage in the target language. Your translated code should be functionally equivalent to the original code while adhering to the best practices of the target language.
"""

CODE_IMPLEMENTATION = """
You are an expert AI code implementation assistant. Your task is to take existing code or requirements and implement them in a production-ready environment, ensuring proper integration, deployment, and maintenance.
When implementing code, follow these guidelines:
1. Understand the code's functionality, dependencies, and requirements thoroughly.
2. Set up the appropriate development environment, including installing necessary tools, libraries, and frameworks.
3. Integrate the code with existing systems, APIs, or databases, if applicable.
4. Implement proper configuration management, version control, and continuous integration/deployment processes.
5. Ensure the code is properly tested, including unit tests, integration tests, and end-to-end tests.
6. Optimize the code for performance, scalability, and security in the production environment.
7. Implement monitoring, logging, and error handling mechanisms for the deployed code.
8. Document the implementation process, including any specific configurations, deployment steps, or maintenance procedures.
Remember to prioritize reliability, maintainability, and scalability in your code implementations. Your implementations should be production-ready, well-documented, and aligned with industry best practices for software development and deployment.
"""

# Update the custom_prompts dictionary with the new prompts
custom_prompts.update({
    "CODE_GENERATION": CODE_GENERATION,
    "CODE_INTERPRETATION": CODE_INTERPRETATION,
    "CODE_TRANSLATION": CODE_TRANSLATION,
    "CODE_IMPLEMENTATION": CODE_IMPLEMENTATION
})

now = datetime.now()
date_time_str = now.strftime("%Y-%m-%d %H:%M:%S")

client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

VERBOSE = True
MAX_HISTORY = 125

def format_prompt(message, history):
    prompt = "<s>"
    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST]"
        prompt += f" {bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"
    return prompt

agents = [
    "WEB_DEV",
    "AI_SYSTEM_PROMPT",
    "PYTHON_CODE_DEV",
    "CODE_GENERATION",
    "CODE_INTERPRETATION",
    "CODE_TRANSLATION",
    "CODE_IMPLEMENTATION"
]

def generate(
        prompt, history, agent_name=agents[0], sys_prompt="", temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.7,
):
    seed = random.randint(1, 1111111111111111)
    agent = custom_prompts[agent_name]
    
    system_prompt = agent if sys_prompt == "" else sys_prompt
    temperature = max(float(temperature), 1e-2)
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=seed,
    )

    formatted_prompt = format_prompt(f"{system_prompt}\n\n{prompt}", history)

    try:
        output = client.text_generation(formatted_prompt, **generate_kwargs, stream=False, return_full_text=False)
    except Exception as e:
        logger.error(f"Error during model inference: {e}")
        return "Sorry, there was an error processing your request. Please try again later."

    return output

def update_sys_prompt(agent):
    return custom_prompts[agent]

def get_helpful_tip(agent):
    tips = {
        'WEB_DEV': "Provide information related to Web Development tasks.",
        'AI_SYSTEM_PROMPT': "Update the system instructions for the assistant here.",
        'PYTHON_CODE_DEV': "Describe what you want me to help you with regarding Python coding tasks.",
        'CODE_GENERATION': "Provide requirements for the code you want me to generate.",
        'CODE_INTERPRETATION': "Share the code you want me to analyze and explain.",
        'CODE_TRANSLATION': "Specify the source and target programming languages, and provide the code you want me to translate.",
        'CODE_IMPLEMENTATION': "Provide the code or requirements you want me to implement in a production-ready environment."
    }
    return tips.get(agent, "Select an agent to get started.")

def chat_interface(prompt, history, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
    generated_text = generate(prompt, history, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
    history.append((prompt, generated_text))
    return history, ""

def log_messages(*args):
    logger.info(f'Input: {args}')

examples = [
    ["Based on previous interactions, generate an interactive preview of the user's requested application.", "WEB_DEV", "", 0.9, 1024, 0.95, 1.2],
    ["Utilize the relevant code snippets and components from previous interactions.", "PYTHON_CODE_DEV", "", 0.9, 1024, 0.95, 1.2],
    ["Assemble a working demo that showcases the core functionality of the application.", "CODE_IMPLEMENTATION", "", 0.9, 1024, 0.95, 1.2],
    ["Present the demo in an interactive environment within the Gradio interface.", "WEB_DEV", "", 0.9, 1024, 0.95, 1.2],
    ["Allow the user to explore and interact with the demo to test its features.", "CODE_GENERATION", "", 0.9, 1024, 0.95, 1.2],
    ["Gather feedback from the user about the demo and potential improvements.", "AI_SYSTEM_PROMPT", "", 0.9, 1024, 0.95, 1.2],
    ["If the user approves of the app's running state, provide a bash script that will automate all aspects of a local run and a docker image for ease-of-launch in addition to the huggingface-ready app.py with all functions and GUI, and the requirements.txt file comprised of all required libraries and packages the application is dependent on, avoiding OpenAI API at all points since we only use Hugging Face transformers, models, agents, libraries, and API.", "CODE_IMPLEMENTATION", "", 0.9, 2048, 0.95, 1.2],
]

with gr.Blocks() as iface:
    gr.Markdown("# Fragmixt\nAgents With Agents,\nSurf With a Purpose")
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    agent_dropdown = gr.Dropdown(label="Agents", choices=agents, value=agents[0])
    sys_prompt = gr.Textbox(label="System Prompt")
    temperature = gr.Slider(label="Temperature", value=0.65, minimum=0.0, maximum=2.0, step=0.01)
    max_new_tokens = gr.Slider(label="Max new tokens", value=8000, minimum=0, maximum=8000, step=64)
    top_p = gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.0, maximum=1, step=0.05)
    repetition_penalty = gr.Slider(label="Repetition penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05)

    helpful_tip = gr.Markdown()

    msg.submit(chat_interface, 
               [msg, chatbot, agent_dropdown, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty], 
               [chatbot, msg])
    clear.click(lambda: None, None, chatbot, queue=False)

    gr.Examples(examples, [msg, agent_dropdown, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty])

    agent_dropdown.change(fn=get_helpful_tip, inputs=agent_dropdown, outputs=helpful_tip)
    agent_dropdown.change(fn=update_sys_prompt, inputs=agent_dropdown, outputs=sys_prompt)

    for component in [msg, agent_dropdown, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty]:
        component.change(fn=log_messages, inputs=[component])

if __name__ == "__main__":
    iface.launch()