import gradio as gr
from huggingface_hub import InferenceClient
import csv
import os

# Initialize the client with the open-source AI model
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Files for storing tested prompts and feedback
PROMPT_HISTORY_FILE = "prompt_history.csv"
FEEDBACK_FILE = "feedback.csv"

# Function to save prompts and responses in a CSV file
def log_prompt(prompt, response):
    file_exists = os.path.isfile(PROMPT_HISTORY_FILE)
    with open(PROMPT_HISTORY_FILE, mode="a", newline="", encoding="utf-8") as file:
        writer = csv.writer(file)
        if not file_exists:
            writer.writerow(["Prompt", "Response"])
        writer.writerow([prompt, response])

# Function to test prompts and store responses
def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]
    
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    
    messages.append({"role": "user", "content": message})
    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

    # Save the prompt and response
    log_prompt(message, response)

# Function to generate advanced prompts with different prompting techniques
def generate_prompt(prompt, technique):
    if technique == "Zero-shot":
        return prompt
    elif technique == "Few-shot":
        return f"Example 1: Input: 'Good morning', Output: 'Hello!'\nExample 2: Input: 'How are you?', Output: 'I'm fine, thank you!'\nNow generate a response for: '{prompt}'"
    elif technique == "Chain-of-Thought":
        return f"Step 1: Analyze the question '{prompt}'\nStep 2: Break down each element separately\nStep 3: Provide a detailed response."
    else:
        return prompt

# Function to collect feedback on AI responses
def rate_response(prompt, response, rating):
    with open(FEEDBACK_FILE, mode="a", newline="", encoding="utf-8") as file:
        writer = csv.writer(file)
        writer.writerow([prompt, response, rating])
    return "Thank you for your feedback! Your rating has been recorded."

# Main UI for testing Prompt Engineering
chat_interface = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
    title="Prompt Engineering Lab",
    description="Test your prompts with an open-source AI model and analyze the results.",
)

# UI for selecting advanced prompting techniques
prompt_interface = gr.Interface(
    fn=generate_prompt,
    inputs=[gr.Textbox(label="Prompt"), gr.Radio(["Zero-shot", "Few-shot", "Chain-of-Thought"], label="Technique")],
    outputs="text",
    title="Prompt Engineering Techniques",
    description="Select an advanced technique to generate more effective prompts.",
)

# UI for collecting feedback on AI responses
rating_interface = gr.Interface(
    fn=rate_response,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Textbox(label="AI Output"),
        gr.Slider(1, 5, step=1, label="Rating (1-5)"),
    ],
    outputs="text",
    title="Prompt Quality Evaluation",
    description="Rate the quality of the AI response to improve the prompts.",
)

# Creating a tabbed interface with all components
demo = gr.TabbedInterface(
    [chat_interface, prompt_interface, rating_interface],
    ["AI Chat", "Prompt Engineering", "Output Evaluation"]
)

# Launch the UI
if __name__ == "__main__":
    demo.launch()