import os
import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the Hugging Face Inference client
HF_API_KEY = os.environ.get("HF_API_KEY")
HF_MODEL_NAME = os.environ.get("HF_MODEL_NAME")

client = InferenceClient(model=HF_MODEL_NAME, token=HF_API_KEY)

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    segment_profile,
    value_proposition,
    promotion,
    subject,
    ask_for_subject_suggestions,
    max_tokens,
    temperature,
    top_p,
):
    # Construct the system message with additional inputs
    enhanced_system_message = (
        f"{system_message}\n\n"
        f"Segment Profile: {segment_profile}\n"
        f"Value Proposition: {value_proposition}\n"
        f"Goal and Promotion: {promotion}\n"
        f"Subject Line: {subject}\n"
    )

    # If the user wants Subject Line suggestions, modify the prompt
    if ask_for_subject_suggestions:
        enhanced_system_message += "The user is also asking for subject line suggestion to catch their customer's attention and improve Email Open Rate."

    messages = [{"role": "system", "content": enhanced_system_message}]

    # Add conversation history
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # Add the current user message
    messages.append({"role": "user", "content": message})

    # Generate the response
    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


# Define the Gradio interface
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(
            value="You are a friendly Chatbot, a digital marketing expert and a talented copywriter. You are trying to help a user write a creative email that can achieve campaign goals, a high Open Rate, CTR and conversion rate - based on user input.",
            label="Instructions to Bot",
        ),
        gr.Textbox(label="Your Target Customer Segment Profile", placeholder="Describe the profile of your target customer segment (e.g., age, gender, interests, profession)"),
        gr.Textbox(
            label="Your Value Proposition",
            placeholder="Describe how your solution to customer problems offers them unique value",
        ),
        gr.Textbox(
            label="Campaign Goal, Special Event, Promotion and Call to Action",
            placeholder="Describe your campaign goal, a special event, promotion and Call to Action that you hope your target segment will act upon",
        ),
        gr.Textbox(label="Subject Line", placeholder="Enter the Subject Line of the Email or ask for suggestions"),
        gr.Checkbox(label="Ask for Subject Line Suggestions", value=False),
        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="Email Copywriter",
    description="This app creates a customized email that resonates with your customers to improve CTR and conversion. Based on your input. Powered by Hugging Face Inference, Design Thinking, and domain expertise. Expand Additional Inputs by clicking on the arrow, input more details about your customers, then a message describing what you need the assistant to do for you. Developed by wn. Disclaimer: AI makes mistakes. Use with caution and at your own risk!",
)


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