import gradio as gr
from transformers import pipeline, set_seed

# --- Model setup ---
generator = pipeline("text-generation", model="gpt2", device=-1)
set_seed(42)

SYSTEM_PROMPT = (
    "You are a friendly, concise fitness coach. "
    "Give short, clear, encouraging answers focused on fitness, exercise, nutrition, and motivation. "
    "Keep replies in simple language and avoid long paragraphs.\n\n"
)

MAX_EXCHANGES_IN_PROMPT = 4

def build_prompt_from_history(history):
    # history is in role/content format: [{"role": "user", "content": ...}, {"role": "assistant", ...}]
    recent_messages = history[-MAX_EXCHANGES_IN_PROMPT*2:]  # last N exchanges
    prompt = SYSTEM_PROMPT
    for msg in recent_messages:
        if msg["role"] == "user":
            prompt += f"User: {msg['content']}\n"
        else:
            prompt += f"Coach: {msg['content']}\n"
    prompt += "Coach:"
    return prompt

def fitness_coach_chat(history, user_message):
    if history is None:
        history = []

    history.append({"role": "user", "content": user_message})
    prompt = build_prompt_from_history(history)

    gen = generator(
        prompt,
        max_new_tokens=120,
        do_sample=True,
        temperature=0.8,
        top_p=0.9,
        return_full_text=False,
        num_return_sequences=1,
    )

    ai_text = gen[0]["generated_text"].strip()
    history.append({"role": "assistant", "content": ai_text})

    return history, ""

custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;600;700&display=swap');
body { background: linear-gradient(135deg, #ffe6f0 0%, #fff0f6 100%); font-family: 'Poppins', sans-serif; }
.gradio-container { font-family: 'Poppins', sans-serif; }
h1 { color: #9b2c6b; text-align:center; }
"""

with gr.Blocks(css=custom_css, title="AI Fitness Coach") as demo:
    gr.Markdown("<h1>💖 AI Fitness Coach</h1>")
    chatbot = gr.Chatbot(label="Coach", type="messages", height=480)
    state = gr.State([])

    with gr.Row():
        txt = gr.Textbox(show_label=False, placeholder="Type your question...")
        send = gr.Button("Send")

    send.click(fitness_coach_chat, inputs=[state, txt], outputs=[chatbot, txt])
    txt.submit(fitness_coach_chat, inputs=[state, txt], outputs=[chatbot, txt])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)