import torch
from transformers import pipeline
import gradio as gr

import os


from huggingface_hub import login, HfFolder
from getpass import getpass

def init_huggingface_auth():
    """
    Initialize Hugging Face authentication with multiple fallback options
    Returns:
        bool: True if authentication successful, False otherwise
    """
    # First try: Check if already logged in
    if HfFolder.get_token() is not None:
        print("Already logged in to Hugging Face!")
        return True
    
    # Second try: Environment variable
    token = os.getenv('HUGGINGFACE_TOKEN')
    
    # Third try: Manual input if no token found
    if not token:
        try:
            print("Please enter your Hugging Face token.")
            print("You can find it at: https://huggingface.co/settings/tokens")
            token = getpass("Hugging Face Token: ")
        except Exception as e:
            print(f"Error during token input: {e}")
            return False
    
    # Try to login with the token
    if token:
        try:
            login(token=token, add_to_git_credential=True)
            print("Successfully logged in to Hugging Face!")
            return True
        except Exception as e:
            print(f"Error during login: {e}")
            return False
    
    print("No Hugging Face token found or provided")
    return False
    
if not init_huggingface_auth():
        print("Warning: Hugging Face authentication failed")
        
print("Authentication successful - continuing with the application")

# Initialize the pipeline
pipe = pipeline(
    "text-generation",
    model="google/gemma-2-2b-jpn-it",
    
    device="cpu",  # 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)