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,
    customer_profile,
    customer_goals,
    company_solution,
    main_topic,
    ask_for_topic_suggestions,
    max_tokens,
    temperature,
    top_p,
):
    # Construct the system message with additional inputs
    enhanced_system_message = (
        f"{system_message}\n\n"
        f"Customer Profile: {customer_profile}\n"
        f"Customer Goals, Pain Points, Obstacles, Wishes and Preferences: {customer_goals}\n"
        f"Company solutions, products or services, value proposition and differentiators to solve customer problem: {company_solution}\n"
        f"Main Topic: {main_topic}\n"
    )

    # If the user wants topic suggestions, modify the prompt
    if ask_for_topic_suggestions:
        enhanced_system_message += "The user is also asking for topic suggestions to address their customer's needs."

    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 post to improve their SEO based on their input.",
            label="Instructions to Bot",
        ),
        gr.Textbox(label="Your Customer Profile", placeholder="Describe your customer profile (e.g., age, interests, profession)"),
        gr.Textbox(
            label="Customer Goals, Pain Points, Obstacles, Wishes and Preferences",
            placeholder="Describe your customer's goals, pain points, concerns, obstacles, wishes, and preferences",
        ),
        gr.Textbox(
            label="Company solutions, products or services",
            placeholder="Describe your company's Company solutions, product or services, value proposition and differentiators to solve customer problems",
        ),
        gr.Textbox(label="Main Topic", placeholder="Enter the main topic of the post"),
        gr.Checkbox(label="Ask for Topic 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="SEO Assistant",
    description="This app provides customized content that resonates with your customers to improve your SEO. 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 can make mistakes. Use with caution and at your own risk!",
)


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