import gradio as gr
import torch
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer

model = None
tokenizer = None

alpaca_prompt = """පහත දැක්වෙන්නේ යම් කාර්යයක් පිළිබඳ විස්තර කරන උපදෙසක් සහ එයට අදාළ තොරතුරු ඇතුළත් ආදානයකි. ඉල්ලූ කාර්යය නිවැරදිව සම්පූර්ණ කළ හැකි ප්‍රතිචාරයක් සපයන්න.
### උපදෙස:
{}
### ආදානය:
{}
### ප්‍රතිචාරය:
{}"""

@spaces.GPU
def infer(message, history, enable_history=False, max_new_tokens=1024):
    global model, tokenizer
    
    if model is None:
        tokenizer = AutoTokenizer.from_pretrained("manthilaffs/Gamunu-4B-Instruct-Alpha")
        model = AutoModelForCausalLM.from_pretrained(
            "manthilaffs/Gamunu-4B-Instruct-Alpha",
            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
            device_map="auto",
        )
        model.eval()
    
    # Add history only if enabled
    if enable_history and history:
        prev = "\n".join(
            [f"User: {h[0]}\nGamunu: {h[1]}" for h in history if h[1] is not None]
        )
        context = f"{prev}\n\n{message}"
    else:
        context = message
    
    prompt = alpaca_prompt.format(
        "ඔබ ගැමුණු (Gamunu) නම් AI සහායකයායි. ඔබව නිර්මාණය කර ඇත්තේ මන්තිල විසිනි. "
        "ඔබේ කාර්යය වන්නේ පරිශීලකයන්ගේ උපදෙස් නිවැරදිව පිලිපැදීම හා අසා ඇති ප්‍රශ්නවලට නිවැරදිව පිළිතුරු සපයමින් ඔවුන්ට සහය වීමයි.",
        context.strip(),
        "",
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.inference_mode():
        outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
    
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if "### ප්‍රතිචාරය:" in text:
        text = text.split("### ප්‍රතිචාරය:")[-1].strip()
    
    return text

# Custom CSS for styling
custom_css = """
#splash-screen {
    position: fixed;
    top: 0;
    left: 0;
    width: 100vw;
    height: 100vh;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    display: flex;
    justify-content: center;
    align-items: center;
    z-index: 9999;
    animation: fadeOut 1s ease-in-out 3s forwards;
}

@keyframes fadeOut {
    to {
        opacity: 0;
        visibility: hidden;
    }
}

.sinhala-animation {
    font-size: 40px;
    font-weight: bold;
    color: white;
    animation: letterRotate 3s ease-in-out infinite;
    text-shadow: 0 0 10px rgba(255,255,255,0.5);
}

@keyframes letterRotate {
    0%, 100% { transform: rotateY(0deg) scale(1); opacity: 1; }
    25% { transform: rotateY(180deg) scale(1.2); opacity: 0.8; }
    50% { transform: rotateY(360deg) scale(1); opacity: 1; }
    75% { transform: rotateY(540deg) scale(1.2); opacity: 0.8; }
}

/* Smaller font sizes for chat */
.message-wrap .message {
    font-size: 0.9rem !important;
}

.message-wrap p {
    font-size: 0.9rem !important;
}

/* Avatar styling */
.message-wrap.user .avatar-container::before {
    content: "👤";
    font-size: 24px;
    display: flex;
    align-items: center;
    justify-content: center;
    width: 40px;
    height: 40px;
    border-radius: 50%;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
}

.message-wrap.bot .avatar-container::before {
    content: "🧠";
    font-size: 24px;
    display: flex;
    align-items: center;
    justify-content: center;
    width: 40px;
    height: 40px;
    border-radius: 50%;
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
}

/* Hide default avatars if they exist */
.avatar-container img {
    display: none;
}

.avatar-container {
    width: 40px !important;
    height: 40px !important;
    min-width: 40px !important;
}
"""

# Splash screen HTML
splash_html = """
<div id="splash-screen">
    <div class="sinhala-animation">ගැමුණු</div>
</div>
<script>
    setTimeout(() => {
        document.getElementById('splash-screen').style.display = 'none';
    }, 4000);
</script>
"""

# ---------------- UI ----------------
with gr.Blocks(css=custom_css) as demo:
    gr.HTML(splash_html)
    
    chat = gr.ChatInterface(
        fn=lambda message, history: infer(message, history, enable_history.value, max_new_tokens.value),
        title="🧠 Gamunu 4B Instruct - Demo",
        theme=gr.themes.Default(text_size="sm"),
        examples=[
            ["හෙලෝ ගැමුණු! මම සමන්, ඔයාට කොහොමද?"],
            ["ෆොටෝසින්තසිස් ක්‍රියාවලිය පැහැදිලි කරන්න."],
            ["මෙම වාක්‍යය සිංහලයට පරිවර්තනය කරන්න: 'The sun rises in the east.'"],
            ["'completed' තත්ත්වයේ ඇති වාර්තා ගණන ගණනය කිරීමට දත්ත සමුදා විමසුමක් (database query) ගොඩනඟන්න."],
            ["ඔබ ගුරුවරයෙකු ලෙස ක්‍රියාකරන්න. ශිෂ්‍යයාට ඉතිහාසය උගන්වන්න."],
            ["පහත ප්‍රකාශය ප්‍රංශ භාෂාවට පරිවර්තනය කරන්න. Laughter is the best medicine."],
            ["ඝන වස්තුවක හා ද්‍රවයක පරිමාවන්හි වෙනස පැහැදිලි කරන්න."],
            ["වෙස් මුහුණු කලාවේ ප්‍රධාන අංග මොනවාද? වර්තමානයේ මෙම කලාව ප්‍රචලිතව පවතින ප්‍රදේශ මොනවාද?"]
        ]
    )
    
    with gr.Accordion("⚙️ Advanced Settings", open=False):
        enable_history = gr.Checkbox(label="Enable chat history", value=False)
        max_new_tokens = gr.Slider(64, 1024, value=512, step=32, label="🔢 Max New Tokens")
    
    gr.Markdown("""
---
🪶 **Model:** [manthilaffs/Gamunu-4B-Instruct-Alpha](https://huggingface.co/manthilaffs/Gamunu-4B-Instruct-Alpha)  
© 2025 Gamunu Project | Experimental Release
""")

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