import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the inference client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def respond(
    message, 
    history: list[tuple[str, str]], 
    system_message, 
    max_tokens, 
    temperature, 
    top_p,
):
    # Prepare messages for the chat completion
    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})
    
    # Stream 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

# Custom CSS for a more elegant design
css = """
.gradio-container {
    background: linear-gradient(135deg, #f6f8f9 0%, #e5ebee 100%);
    font-family: 'Inter', 'Helvetica Neue', Arial, sans-serif;
}

.chatbot {
    max-width: 800px;
    margin: 0 auto;
    background: white;
    border-radius: 15px;
    box-shadow: 0 10px 25px rgba(0,0,0,0.1);
    overflow: hidden;
}

.chatbot-header {
    background: linear-gradient(90deg, #6a11cb 0%, #2575fc 100%);
    color: white;
    padding: 20px;
    text-align: center;
    font-weight: bold;
}

.input-area {
    background: #f0f4f8;
    padding: 15px;
    border-top: 1px solid #e1e8f0;
}

.message-list {
    max-height: 500px;
    overflow-y: auto;
    padding: 20px;
}

.user-message {
    background-color: #e6f2ff;
    border-radius: 15px;
    padding: 10px 15px;
    margin-bottom: 10px;
    max-width: 80%;
    align-self: flex-end;
}

.assistant-message {
    background-color: #f0f0f0;
    border-radius: 15px;
    padding: 10px 15px;
    margin-bottom: 10px;
    max-width: 80%;
    align-self: flex-start;
}
"""

# Create the Gradio interface with enhanced design
demo = gr.ChatInterface(
    respond,
    title="I am your personal Photography AI Chatbot",
    description="Powered by Zephyr 7B Beta",
    css=css,
    additional_inputs=[
        gr.Textbox(
            value="You are a friendly, helpful AI assistant who provides clear and concise answers.", 
            label="System Message", 
            info="Define the personality and guidelines for the AI"
        ),
        gr.Slider(
            minimum=1, 
            maximum=2048, 
            value=512, 
            step=1, 
            label="Max New Tokens", 
            info="Maximum number of tokens to generate"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=4.0, 
            value=0.7, 
            step=0.1, 
            label="Temperature", 
            info="Controls randomness - lower is more focused, higher is more creative"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=1.0, 
            value=0.95, 
            step=0.05, 
            label="Top-p (Nucleus Sampling)", 
            info="Controls diversity of token selection"
        ),
    ],
    theme=gr.themes.Soft()
)

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