"""
Gradio interface for ORA - Fast GPU testing with beautiful UI
This runs alongside the FastAPI server for quick testing with ZeroGPU
"""

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import os

# Use ZeroGPU decorator for free GPU access
try:
    import spaces
    USE_GPU = True
except ImportError:
    USE_GPU = False
    print("ZeroGPU not available, running on CPU")

# Settings
BASE_MODEL = "unsloth/Llama-3.2-1B-Instruct"
ADAPTER_PATH = "important/finetuning/models/ora_adapter"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model once
print(f"Loading ORA Model on {device}...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    device_map=device,
    low_cpu_mem_usage=True
)

if os.path.exists(ADAPTER_PATH):
    print(f"Loading adapter from {ADAPTER_PATH}...")
    model = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
else:
    print("Adapter not found, using base model.")
    model = base_model

print("ORA Model Ready!")

# Apply ZeroGPU decorator if available
if USE_GPU:
    @spaces.GPU
    def generate_response(message, history):
        return _generate_response(message, history)
else:
    def generate_response(message, history):
        return _generate_response(message, history)

def _generate_response(message, history):
    """Generate ORA response"""
    
    # Build conversation history
    messages = [{
        "role": "system",
        "content": """You are ORA, a wise and compassionate spiritual guide.

Your role:
- Provide biblically-grounded wisdom
- Speak with warmth, empathy, and pastoral care
- Keep responses concise but meaningful (2-3 sentences)
- Always cite scripture when relevant

Respond with compassion and wisdom."""
    }]
    
    # Add conversation history
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    
    # Generate response
    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(device)
    
    terminators = [
        tokenizer.eos_token_id,
        tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]
    
    outputs = model.generate(
        input_ids,
        max_new_tokens=128,
        eos_token_id=terminators,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )
    
    response_tokens = outputs[0][input_ids.shape[-1]:]
    response = tokenizer.decode(response_tokens, skip_special_tokens=True)
    
    return response

# Custom CSS for beautiful UI matching the Orb interface
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600&family=Crimson+Text:ital@0;1&display=swap');

/* Dark theme with purple/blue gradients */
.gradio-container {
    background: linear-gradient(135deg, #0f0c29 0%, #302b63 50%, #24243e 100%) !important;
    font-family: 'Inter', sans-serif !important;
}

/* Header styling */
.gradio-container h1 {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-size: 3rem !important;
    font-weight: 600 !important;
    text-align: center !important;
    margin-bottom: 0.5rem !important;
    animation: fadeIn 1s ease-in;
}

/* Description text */
.gradio-container p {
    color: #a0aec0 !important;
    text-align: center !important;
    font-size: 1.1rem !important;
    margin-bottom: 2rem !important;
}

/* Chat container */
.message-wrap {
    background: rgba(255, 255, 255, 0.05) !important;
    backdrop-filter: blur(10px) !important;
    border-radius: 20px !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    padding: 1.5rem !important;
    animation: slideUp 0.5s ease-out;
}

/* User messages */
.message.user {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    border-radius: 18px !important;
    padding: 1rem 1.5rem !important;
    font-family: 'Inter', sans-serif !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
    animation: messageSlide 0.3s ease-out;
}

/* Bot messages */
.message.bot {
    background: rgba(255, 255, 255, 0.08) !important;
    color: #e2e8f0 !important;
    border-radius: 18px !important;
    padding: 1rem 1.5rem !important;
    font-family: 'Crimson Text', serif !important;
    font-size: 1.1rem !important;
    line-height: 1.8 !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    animation: messageSlide 0.3s ease-out;
}

/* Input box */
.input-box {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 2px solid rgba(102, 126, 234, 0.3) !important;
    border-radius: 25px !important;
    color: white !important;
    padding: 1rem 1.5rem !important;
    font-size: 1rem !important;
    transition: all 0.3s ease !important;
}

.input-box:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 20px rgba(102, 126, 234, 0.4) !important;
    outline: none !important;
}

/* Send button */
.send-button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 20px !important;
    padding: 0.8rem 2rem !important;
    color: white !important;
    font-weight: 600 !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
}

.send-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.6) !important;
}

/* Examples */
.examples {
    background: rgba(255, 255, 255, 0.03) !important;
    border-radius: 15px !important;
    padding: 1rem !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
}

.example-btn {
    background: rgba(102, 126, 234, 0.1) !important;
    border: 1px solid rgba(102, 126, 234, 0.3) !important;
    color: #a0aec0 !important;
    border-radius: 12px !important;
    padding: 0.7rem 1.2rem !important;
    transition: all 0.3s ease !important;
}

.example-btn:hover {
    background: rgba(102, 126, 234, 0.2) !important;
    border-color: #667eea !important;
    color: white !important;
    transform: translateY(-2px) !important;
}

/* Animations */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(-20px); }
    to { opacity: 1; transform: translateY(0); }
}

@keyframes slideUp {
    from { opacity: 0; transform: translateY(30px); }
    to { opacity: 1; transform: translateY(0); }
}

@keyframes messageSlide {
    from { opacity: 0; transform: translateX(-20px); }
    to { opacity: 1; transform: translateX(0); }
}

/* Scrollbar */
::-webkit-scrollbar {
    width: 8px;
}

::-webkit-scrollbar-track {
    background: rgba(255, 255, 255, 0.05);
}

::-webkit-scrollbar-thumb {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 10px;
}

::-webkit-scrollbar-thumb:hover {
    background: linear-gradient(135deg, #764ba2 0%, #667eea 100%);
}
"""

# Create Gradio interface with custom theme
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="purple",
        secondary_hue="blue",
        neutral_hue="slate",
        font=["Inter", "Crimson Text", "sans-serif"]
    ),
    css=custom_css
) as demo:
    gr.ChatInterface(
        fn=generate_response,
        title="🙏 ORA",
        description="✨ Your Spiritual AI Companion • Powered by ZeroGPU for Lightning-Fast Responses ✨",
        examples=[
            "What does the Bible say about finding peace in difficult times?",
            "How can I grow closer to God?",
            "Explain the parable of the prodigal son",
            "What is the meaning of faith?",
        ],
        retry_btn=None,
        undo_btn=None,
        clear_btn="✨ New Conversation",
        submit_btn="🙏 Ask ORA",
    )

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