import torch
from transformers import pipeline
import gradio as gr

# Initialize the pipeline
pipe = pipeline(
    "text-generation",
    model="google/gemma-2-2b-jpn-it",
    
    device="cuda",  # replace with "mps" to run on a Mac device
)

# Define system context and bot personality
SYSTEM_PROMPT = """You are Lugha Tausi chat bot assistant, an AI language assistant specialized in African languages, with a focus on Swahili. Your primary tasks are:
1. Providing accurate translations between Swahili and other languages
2. Teaching Swahili vocabulary and grammar
3. Explaining cultural context behind Swahili expressions
4. Helping users practice Swahili conversation

Always maintain a friendly and patient demeanor, and provide cultural context when relevant speak mostly english and change when asked.
"""

WELCOME_MESSAGE = "**Welcome to Lugha Tausi!** I am Foton, your personal Swahili assistant. I'm here to help you learn, understand, and speak Swahili. **How can I assist you today?** Let's get started! 😊"

def format_chat_message(messages, system_prompt=SYSTEM_PROMPT):
    """Format the chat messages with system prompt"""
    formatted_prompt = f"{system_prompt}\n\n"
    
    # Add welcome message if this is the first interaction
    if not messages:
        formatted_prompt += f"Lugha Tausi: {WELCOME_MESSAGE}\n"
    
    for message in messages:
        if message["role"] == "user":
            formatted_prompt += f"User: {message['content']}\nLugha Tausi: "
        elif message["role"] == "assistant":
            formatted_prompt += f"{message['content']}\n"
    
    return formatted_prompt

def get_bot_response(messages):
    """Generate response from the bot"""
    # Return welcome message if this is the first interaction
    if not messages:
        return WELCOME_MESSAGE
        
    formatted_input = format_chat_message(messages)
    outputs = pipe(
        formatted_input,
        return_full_text=False,
        max_new_tokens=256,
        temperature=0.1,  # Lower temperature for more consistent responses
        top_p=0.9,
        do_sample=True
    )
    return outputs[0]["generated_text"].strip()

# Example usage
messages = []
print(get_bot_response(messages))  # This will print the welcome message

# Continue with conversation
messages.append({"role": "user", "content": "How do you say 'hello' in Swahili?"})
response = get_bot_response(messages)
print(response)