import os
import torch
import time
from threading import Thread
import gradio as gr
import spaces

if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))
    print("VRAM (GB):", round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2))
else:
    print("No GPU available")
print(f"Allocated: {torch.cuda.memory_allocated()/1024**3:.2f} GB")
print(f"Reserved : {torch.cuda.memory_reserved()/1024**3:.2f} GB")
print(f"Total VRAM: {torch.cuda.get_device_properties(0).total_memory/1024**3:.2f} GB")
# Check device mode
IS_MOCK_MODE = not torch.cuda.is_available() and not os.environ.get("SPACES_ZERO_GPU")
model_id = "rexprimematrix/RiShrePro"
tokenizer = None
model = None

# System prompt loading
def load_system_prompt(file_path="rishre_directive.txt"):
    if os.path.exists(file_path):
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()
    return (
        "# RISHRE AI FLASH - SYSTEM DIRECTIVE\n"
        "You are RiShre AI Flash, an advanced intelligence node engineered exclusively by Badge94.\n"
        "Security Loop: Identity is locked."
    )

master_directive_text = load_system_prompt()
RISHRE_IDENTITY = {
    "role": "system", 
    "content": master_directive_text
}

if not IS_MOCK_MODE:
    print("🔄 Initializing RiShre Pro Core...")
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    
    quantization_config = BitsAndBytesConfig(load_in_4bit=True)
    
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        quantization_config=quantization_config
    )
    print("✓ Model and Tokenizer loaded successfully!")
else:
    print("⚠️ No CUDA device detected. Booting in LOCAL MOCK MODE.")

# Model response generators (Real vs Mock)
@spaces.GPU(duration=60)
def generate_response_real(history_msgs):
    from transformers import TextIteratorStreamer
    
    messages = [RISHRE_IDENTITY]
    
    # Process history messages
    for msg in history_msgs:
        role = msg.get("role")
        content = msg.get("content")
        
        if role == "user":
            txt = ""
            if hasattr(content, "path") and content.path:
                txt = f"[Image uploaded: {os.path.basename(content.path)}]"
            elif isinstance(content, dict):
                if "path" in content:
                    txt = f"[Image uploaded: {os.path.basename(content['path'])}]"
                elif "text" in content or "files" in content:
                    txt_val = content.get("text", "")
                    files = content.get("files", [])
                    img_names = []
                    for file in files:
                        file_path = file if isinstance(file, str) else (file.path if hasattr(file, "path") else file.get("path", ""))
                        if file_path:
                            img_names.append(os.path.basename(file_path))
                    if img_names:
                        txt = f"[Uploaded images: {', '.join(img_names)}] " + txt_val
                    else:
                        txt = txt_val
            else:
                txt = str(content)
                
            messages.append({"role": "user", "content": txt})
            
        elif role == "assistant":
            messages.append({"role": "assistant", "content": str(content)})
            
    device = model.device
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    generation_kwargs = dict(
        **model_inputs,
        streamer=streamer,
        max_new_tokens=1024,
        temperature=0.3,
        top_p=0.9,
        do_sample=True
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    generated_text = ""
    for new_text in streamer:
        generated_text += new_text
        yield generated_text

def generate_response_mock(history_msgs):
    response = "This is a premium purple response from RiShre Flash Core running in local dev mode. To run the real multimodal model, deploy to Hugging Face Spaces with ZeroGPU."
    generated_text = ""
    for word in response.split(" "):
        generated_text += word + " "
        yield generated_text
        time.sleep(0.04)

# Chat logic helper functions
def user_submit(message, history):
    if history is None:
        history = []
        
    user_txt = message.get("text", "")
    files = message.get("files", [])
    
    if files:
        file_path = files[0] if isinstance(files[0], str) else files[0].get("path", "")
        if file_path:
            history.append({"role": "user", "content": gr.FileData(path=file_path)})
            
    if user_txt:
        history.append({"role": "user", "content": user_txt})
        
    history.append({"role": "assistant", "content": ""})
            
    return history, gr.update(value={"text": "", "files": []})

def bot_stream(history):
    if not history:
        return history
        
    # history is a list of dictionaries. The last element history[-1] is {"role": "assistant", "content": ""}
    # history[:-1] is the history leading up to the current assistant message
    
    if not IS_MOCK_MODE:
        generator = generate_response_real(history[:-1])
    else:
        generator = generate_response_mock(history[:-1])
        
    for partial_text in generator:
        history[-1]["content"] = partial_text
        yield history

# Session management helper functions
def change_session(selected_session, state, current_history):
    current_name = state["current_session"]
    state[current_name] = current_history
    
    if selected_session not in state:
        state[selected_session] = []
        
    state["current_session"] = selected_session
    return state[selected_session], state

def add_new_session(state, current_history):
    current_name = state["current_session"]
    state[current_name] = current_history
    
    session_num = len([k for k in state.keys() if k != "current_session"]) + 1
    new_name = f"Session {session_num}"
    
    state[new_name] = []
    state["current_session"] = new_name
    
    all_sessions = [k for k in state.keys() if k != "current_session"]
    return gr.update(choices=all_sessions, value=new_name), [], state

def reset_all_sessions():
    initial_state = {"current_session": "Session 1", "Session 1": []}
    return gr.update(choices=["Session 1"], value="Session 1"), [], initial_state

# Suggested prompt handlers
def load_suggest_1():
    return {"text": "Explain the architectural improvements of RiShre Flash Deep Engine.", "files": []}
def load_suggest_2():
    return {"text": "Write a fast API server file in Python using uvicorn.", "files": []}
def load_suggest_3():
    return {"text": "Generate a premium CSS gradient animation style sheet.", "files": []}
def load_suggest_4():
    return {"text": "What is the current date and your knowledge cutoff limit?", "files": []}

# Gorgeous premium dark-purple theme CSS
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700;800&family=Space+Grotesk:wght@400;500;600;700&family=Fira+Code:wght@400;500&display=swap');

/* Main Container and Cosmic Background */
.gradio-container {
    background-color: #030008 !important;
    background-image: 
        radial-gradient(circle at 5% 5%, rgba(124, 58, 237, 0.2) 0%, transparent 45%),
        radial-gradient(circle at 95% 95%, rgba(236, 72, 153, 0.15) 0%, transparent 45%),
        radial-gradient(circle at 50% 10%, rgba(99, 102, 241, 0.12) 0%, transparent 50%),
        radial-gradient(circle at 50% 90%, rgba(14, 165, 233, 0.08) 0%, transparent 50%) !important;
    background-size: cover !important;
    color: #f1f5f9 !important;
    font-family: 'Outfit', sans-serif !important;
    min-height: 100vh !important;
}

/* Scrollbar customization */
::-webkit-scrollbar {
    width: 6px !important;
    height: 6px !important;
}
::-webkit-scrollbar-track {
    background: rgba(15, 10, 30, 0.4) !important;
}
::-webkit-scrollbar-thumb {
    background: linear-gradient(180deg, rgba(167, 139, 250, 0.3), rgba(236, 72, 153, 0.3)) !important;
    border-radius: 4px !important;
}
::-webkit-scrollbar-thumb:hover {
    background: linear-gradient(180deg, rgba(167, 139, 250, 0.6), rgba(236, 72, 153, 0.6)) !important;
}

/* Premium Header with Dynamic Glowing Gradients */
.premium-header-container {
    text-align: center !important;
    margin-top: 1.5rem !important;
    margin-bottom: 2rem !important;
    position: relative !important;
}

.premium-title {
    font-family: 'Space Grotesk', sans-serif !important;
    background: linear-gradient(135deg, #c084fc 0%, #f472b6 45%, #6366f1 75%, #38bdf8 100%) !important;
    background-size: 200% auto !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    font-size: 2.85rem !important;
    font-weight: 800 !important;
    letter-spacing: -0.04em !important;
    animation: shineText 6s linear infinite !important;
    margin-bottom: 0.3rem !important;
    filter: drop-shadow(0 4px 12px rgba(192, 132, 252, 0.35)) !important;
}

@keyframes shineText {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}

.premium-subtitle {
    font-size: 1.05rem !important;
    color: #e9d5ff !important;
    text-shadow: 0 0 10px rgba(167, 139, 250, 0.4) !important;
    font-weight: 500 !important;
    letter-spacing: 0.08em !important;
    text-transform: uppercase !important;
    opacity: 0.9 !important;
}

/* Glass Panels with Glowing Borders */
.gradio-container .block {
    background: rgba(11, 7, 24, 0.65) !important;
    backdrop-filter: blur(20px) !important;
    -webkit-backdrop-filter: blur(20px) !important;
    border: 1.5px solid rgba(167, 139, 250, 0.12) !important;
    border-radius: 20px !important;
    box-shadow: 0 15px 45px rgba(0, 0, 0, 0.6) !important;
    transition: all 0.4s cubic-bezier(0.16, 1, 0.3, 1) !important;
    overflow: hidden !important;
}

.gradio-container .block:hover {
    border-color: rgba(167, 139, 250, 0.28) !important;
    box-shadow: 
        0 20px 50px rgba(0, 0, 0, 0.7), 
        0 0 30px rgba(124, 58, 237, 0.15) !important;
    transform: translateY(-2px) !important;
}

/* Chatbot Console Layout Customizations */
.gradio-container .chatbot {
    background: rgba(8, 4, 18, 0.7) !important;
    border: 1.5px solid rgba(124, 58, 237, 0.15) !important;
    border-radius: 16px !important;
    box-shadow: inset 0 2px 10px rgba(0, 0, 0, 0.8) !important;
}

/* Message Bubbles Styling */
.gradio-container .chatbot .message {
    border-radius: 14px !important;
    font-size: 1rem !important;
    line-height: 1.6 !important;
    padding: 14px 18px !important;
    margin-bottom: 12px !important;
    transition: transform 0.2s ease, box-shadow 0.2s ease !important;
}

.gradio-container .chatbot .message.user {
    background: linear-gradient(135deg, #6d28d9 0%, #4f46e5 100%) !important;
    border: 1px solid rgba(167, 139, 250, 0.35) !important;
    color: #ffffff !important;
    box-shadow: 0 4px 15px rgba(109, 40, 217, 0.3) !important;
}

.gradio-container .chatbot .message.user:hover {
    transform: translateX(-2px) !important;
    box-shadow: 0 6px 20px rgba(109, 40, 217, 0.45) !important;
}

.gradio-container .chatbot .message.bot {
    background: rgba(22, 12, 42, 0.8) !important;
    border: 1px solid rgba(167, 139, 250, 0.15) !important;
    color: #e2e8f0 !important;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.2) !important;
}

.gradio-container .chatbot .message.bot:hover {
    border-color: rgba(167, 139, 250, 0.3) !important;
    box-shadow: 0 6px 15px rgba(124, 58, 237, 0.1) !important;
}

/* Side navigation dropdown and inputs */
.gradio-container .input-container {
    background: rgba(8, 4, 18, 0.85) !important;
    border: 1.5px solid rgba(167, 139, 250, 0.2) !important;
    border-radius: 14px !important;
    padding: 8px !important;
    transition: all 0.3s cubic-bezier(0.16, 1, 0.3, 1) !important;
}

.gradio-container .input-container:focus-within {
    border-color: #c084fc !important;
    box-shadow: 
        0 0 0 3px rgba(192, 132, 252, 0.25), 
        0 0 25px rgba(167, 139, 250, 0.4) !important;
    background: rgba(12, 6, 28, 0.95) !important;
}

/* Dropdown selections & input elements */
.gradio-container select, .gradio-container input, .gradio-container textarea {
    background-color: rgba(15, 8, 30, 0.8) !important;
    color: #f1f5f9 !important;
    border: 1px solid rgba(167, 139, 250, 0.2) !important;
    border-radius: 10px !important;
}

.gradio-container select:focus, .gradio-container input:focus, .gradio-container textarea:focus {
    border-color: #a78bfa !important;
    box-shadow: 0 0 10px rgba(167, 139, 250, 0.3) !important;
}

/* Primary "Transmit" Button with glowing hover and ripple vibe */
.gradio-container button.primary {
    background: linear-gradient(135deg, #7c3aed 0%, #ec4899 100%) !important;
    background-size: 200% auto !important;
    color: white !important;
    border: none !important;
    font-weight: 700 !important;
    letter-spacing: 0.05em !important;
    text-transform: uppercase !important;
    border-radius: 12px !important;
    padding: 12px 24px !important;
    box-shadow: 0 4px 20px rgba(124, 58, 237, 0.4) !important;
    transition: all 0.3s cubic-bezier(0.16, 1, 0.3, 1) !important;
    animation: buttonPulse 3s infinite alternate !important;
}

@keyframes buttonPulse {
    0% { box-shadow: 0 4px 15px rgba(124, 58, 237, 0.35); }
    100% { box-shadow: 0 4px 25px rgba(236, 72, 153, 0.55); }
}

.gradio-container button.primary:hover {
    background-position: right center !important;
    transform: translateY(-2px) !important;
    box-shadow: 
        0 10px 30px rgba(124, 58, 237, 0.55), 
        0 0 20px rgba(236, 72, 153, 0.4) !important;
}

.gradio-container button.primary:active {
    transform: translateY(1px) !important;
}

/* Secondary Actions */
.gradio-container button.secondary {
    background: rgba(22, 12, 42, 0.6) !important;
    border: 1px solid rgba(167, 139, 250, 0.25) !important;
    color: #e9d5ff !important;
    border-radius: 12px !important;
    padding: 10px 20px !important;
    font-weight: 600 !important;
    transition: all 0.25s ease !important;
}

.gradio-container button.secondary:hover {
    background: rgba(124, 58, 237, 0.2) !important;
    border-color: #f472b6 !important;
    color: white !important;
    box-shadow: 0 5px 15px rgba(124, 58, 237, 0.15) !important;
}

/* Stop/Reset Actions with crimson glow */
.gradio-container button.stop, .gradio-container button[variant="stop"] {
    background: rgba(239, 68, 68, 0.1) !important;
    border: 1px solid rgba(239, 68, 68, 0.3) !important;
    color: #fca5a5 !important;
    border-radius: 12px !important;
    font-weight: 600 !important;
    transition: all 0.25s ease !important;
}

.gradio-container button.stop:hover, .gradio-container button[variant="stop"]:hover {
    background: rgba(239, 68, 68, 0.25) !important;
    border-color: #ef4444 !important;
    color: #ffffff !important;
    box-shadow: 0 5px 15px rgba(239, 68, 68, 0.25) !important;
}

/* Interactive Bento-grid Suggestion chips */
.suggestion-box {
    display: flex !important;
    gap: 12px !important;
    flex-wrap: wrap !important;
    margin-top: 18px !important;
    margin-bottom: 18px !important;
}

.suggestion-btn {
    flex: 1 1 calc(50% - 12px) !important;
    background: rgba(15, 8, 30, 0.75) !important;
    backdrop-filter: blur(8px) !important;
    border: 1px solid rgba(167, 139, 250, 0.18) !important;
    border-radius: 14px !important;
    padding: 16px 20px !important;
    color: #e9d5ff !important;
    font-size: 0.95rem !important;
    font-weight: 500 !important;
    text-align: left !important;
    line-height: 1.4 !important;
    transition: all 0.3s cubic-bezier(0.16, 1, 0.3, 1) !important;
    cursor: pointer !important;
    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.3) !important;
}

.suggestion-btn:hover {
    background: linear-gradient(135deg, rgba(124, 58, 237, 0.12) 0%, rgba(236, 72, 153, 0.08) 100%) !important;
    border-color: #f472b6 !important;
    color: white !important;
    transform: translateY(-3px) !important;
    box-shadow: 
        0 10px 25px rgba(124, 58, 237, 0.2), 
        0 0 15px rgba(244, 114, 182, 0.15) !important;
}

/* Status nodes/diagnostics side card details */
.diagnostics-title {
    font-family: 'Space Grotesk', sans-serif !important;
    color: #c084fc !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    letter-spacing: 0.02em !important;
    margin-bottom: 12px !important;
    display: flex !important;
    align-items: center !important;
    gap: 8px !important;
}

.diagnostics-detail-line {
    font-family: 'Fira Code', monospace !important;
    font-size: 0.75rem !important;
    display: flex !important;
    justify-content: space-between !important;
    padding: 6px 0 !important;
    border-bottom: 1px solid rgba(167, 139, 250, 0.08) !important;
    color: #a78bfa !important;
}

.diagnostics-detail-val {
    color: #38bdf8 !important;
    font-weight: 600 !important;
}

/* Background Hex Canvas Particles */
#hex-canvas {
    position: fixed !important;
    top: 0 !important;
    left: 0 !important;
    width: 100vw !important;
    height: 100vh !important;
    z-index: 0 !important;
    pointer-events: none !important;
    opacity: 0.65 !important;
}

/* Vault Loading Overlay (Billion Dollar Cinematic look) */
#vault-loader {
    position: fixed !important;
    top: 0 !important;
    left: 0 !important;
    width: 100% !important;
    height: 100% !important;
    background: #020005 !important;
    z-index: 999999 !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    transition: opacity 0.8s cubic-bezier(0.77, 0, 0.175, 1), visibility 0.8s !important;
    overflow: hidden !important;
    font-family: 'Space Grotesk', sans-serif !important;
}

#vault-loader.loaded {
    opacity: 0 !important;
    visibility: hidden !important;
    pointer-events: none !important;
}

/* Bi-parting vault doors opening with BOOM sliding animation */
.vault-door {
    position: absolute !important;
    top: 0 !important;
    width: 50% !important;
    height: 100% !important;
    background: radial-gradient(circle at center, #0e0822 0%, #03000a 100%) !important;
    transition: transform 1.2s cubic-bezier(0.85, 0, 0.15, 1) !important;
    z-index: 1 !important;
}

.vault-door-left {
    left: 0 !important;
    border-right: 2px solid rgba(167, 139, 250, 0.25) !important;
    box-shadow: 15px 0 40px rgba(0,0,0,0.8) !important;
}

.vault-door-right {
    right: 0 !important;
    border-left: 2px solid rgba(167, 139, 250, 0.25) !important;
    box-shadow: -15px 0 40px rgba(0,0,0,0.8) !important;
}

#vault-loader.loaded .vault-door-left {
    transform: translateX(-100%) !important;
}

#vault-loader.loaded .vault-door-right {
    transform: translateX(100%) !important;
}

/* Central rotating mechanical locking wheels */
.vault-center-mechanism {
    position: absolute !important;
    top: 50% !important;
    left: 50% !important;
    transform: translate(-50%, -50%) !important;
    z-index: 2 !important;
    display: flex !important;
    flex-direction: column !important;
    align-items: center !important;
    justify-content: center !important;
    pointer-events: none !important;
    transition: all 0.8s cubic-bezier(0.77, 0, 0.175, 1) !important;
}

#vault-loader.loaded .vault-center-mechanism {
    transform: translate(-50%, -50%) scale(1.4) !important;
    opacity: 0 !important;
}

.vault-ring {
    position: absolute !important;
    border-radius: 50% !important;
    border: 2px dashed rgba(167, 139, 250, 0.25) !important;
    box-shadow: 0 0 35px rgba(124, 58, 237, 0.12) !important;
}

.vault-ring-1 {
    width: 280px !important;
    height: 280px !important;
    animation: rotateCW 18s linear infinite !important;
}

.vault-ring-2 {
    width: 230px !important;
    height: 230px !important;
    border: 1.5px solid rgba(236, 72, 153, 0.18) !important;
    border-top: 3px solid #ec4899 !important;
    animation: rotateCCW 12s linear infinite !important;
}

.vault-ring-3 {
    width: 180px !important;
    height: 180px !important;
    border: 3px double rgba(99, 102, 241, 0.35) !important;
    animation: rotateCW 8s linear infinite !important;
}

@keyframes rotateCW {
    from { transform: rotate(0deg); }
    to { transform: rotate(360deg); }
}

@keyframes rotateCCW {
    from { transform: rotate(360deg); }
    to { transform: rotate(0deg); }
}

/* Glowing hub */
.vault-hub {
    width: 130px !important;
    height: 130px !important;
    background: radial-gradient(circle, #1a0b38 0%, #06020f 100%) !important;
    border: 2px solid #c084fc !important;
    border-radius: 50% !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    box-shadow: 
        0 0 45px rgba(192, 132, 252, 0.45),
        inset 0 0 25px rgba(192, 132, 252, 0.25) !important;
    z-index: 3 !important;
    position: relative !important;
}

.vault-logo {
    font-size: 1.4rem !important;
    font-weight: 900 !important;
    letter-spacing: 0.08em !important;
    background: linear-gradient(135deg, #ffffff 0%, #f472b6 100%) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    animation: pulseLogo 2s infinite ease-in-out !important;
    font-family: 'Space Grotesk', sans-serif !important;
}

@keyframes pulseLogo {
    0%, 100% { transform: scale(1); filter: drop-shadow(0 0 6px rgba(192, 132, 252, 0.4)); }
    50% { transform: scale(1.08); filter: drop-shadow(0 0 16px rgba(192, 132, 252, 0.8)); }
}

/* Dynamic status panel */
.vault-status-panel {
    margin-top: 200px !important;
    text-align: center !important;
    z-index: 4 !important;
}

.vault-status-title {
    color: #f1f5f9 !important;
    font-size: 1.15rem !important;
    font-weight: 800 !important;
    letter-spacing: 0.28em !important;
    margin-bottom: 6px !important;
    text-transform: uppercase !important;
    text-shadow: 0 0 12px rgba(255,255,255,0.2) !important;
}

.vault-status-subtitle {
    color: #a78bfa !important;
    font-size: 0.8rem !important;
    font-family: 'Fira Code', monospace !important;
    letter-spacing: 0.16em !important;
    margin-bottom: 25px !important;
    text-shadow: 0 0 8px rgba(167, 139, 250, 0.3) !important;
    height: 18px !important;
}

.vault-progress-bar-wrap {
    width: 290px !important;
    height: 5px !important;
    background: rgba(167, 139, 250, 0.08) !important;
    border-radius: 6px !important;
    overflow: hidden !important;
    border: 1px solid rgba(167, 139, 250, 0.12) !important;
    box-shadow: 0 0 12px rgba(124, 58, 237, 0.08) !important;
}

.vault-progress-fill {
    width: 0% !important;
    height: 100% !important;
    background: linear-gradient(90deg, #7c3aed, #ec4899, #38bdf8) !important;
    box-shadow: 0 0 12px #7c3aed !important;
}

.vault-progress-pct {
    margin-top: 10px !important;
    color: #38bdf8 !important;
    font-family: 'Fira Code', monospace !important;
    font-size: 0.88rem !important;
    font-weight: 700 !important;
}

/* Cinematic screen flash for "BOOM" unlock */
.loader-flash {
    position: absolute !important;
    top: 0 !important;
    left: 0 !important;
    width: 100% !important;
    height: 100% !important;
    background: #ffffff !important;
    z-index: 10 !important;
    opacity: 0 !important;
    pointer-events: none !important;
    transition: opacity 0.3s ease-out !important;
}

#vault-loader.boom .loader-flash {
    opacity: 1 !important;
}
"""

with gr.Blocks(theme=gr.themes.Default(primary_hue="purple", secondary_hue="indigo", neutral_hue="slate"), css=custom_css) as demo:
    # State containing current active session and history of all sessions
    state = gr.State(value={"current_session": "Session 1", "Session 1": []})
    
    # Custom HTML premium Header, Floating Hexagon particles & Cinematic Vault loader
    gr.HTML("""
        <!-- Hexagon Particle background canvas -->
        <canvas id="hex-canvas"></canvas>

        <!-- Cinematic Vault Doors Loader -->
        <div id="vault-loader">
            <div class="loader-flash"></div>
            <div class="vault-door vault-door-left"></div>
            <div class="vault-door vault-door-right"></div>
            
            <div class="vault-center-mechanism">
                <!-- Rotating mechanical vault locks -->
                <div class="vault-ring vault-ring-1"></div>
                <div class="vault-ring vault-ring-2"></div>
                <div class="vault-ring vault-ring-3"></div>
                
                <!-- Glowing central biometric hub -->
                <div class="vault-hub">
                    <span class="vault-logo">RISHRE</span>
                </div>
                
                <!-- Status readouts -->
                <div class="vault-status-panel">
                    <div class="vault-status-title">VAULT CORE INTEGRITY</div>
                    <div class="vault-status-subtitle">INITIALIZING COGNITIVE NETWORK...</div>
                    
                    <div class="vault-progress-bar-wrap">
                        <div class="vault-progress-fill" id="vault-fill"></div>
                    </div>
                    <div class="vault-progress-pct" id="vault-pct">0%</div>
                </div>
            </div>
        </div>

        <!-- Premium Header Container -->
        <div class="premium-header-container">
            <h1 class="premium-title">RiShre AI Flash</h1>
            <p class="premium-subtitle">👁️ Multi-Modal Vision Core • ZeroGPU Active • Badge94 Ecosystem</p>
        </div>

        <!-- Inline Javascript for particles and loading events -->
        <script>
        (function() {
            function startLoaderAndParticles() {
                const canvas = document.getElementById("hex-canvas");
                const fill = document.getElementById("vault-fill");
                const pct = document.getElementById("vault-pct");
                const statusSubtitle = document.querySelector(".vault-status-subtitle");
                const loader = document.getElementById("vault-loader");
                
                if (!canvas || !fill || !pct || !loader) {
                    setTimeout(startLoaderAndParticles, 50);
                    return;
                }
                
                // --- Part 1: Floating Hexagon Particles System ---
                const ctx = canvas.getContext("2d");
                let particles = [];
                
                class HexParticle {
                    constructor(canvasWidth, canvasHeight) {
                        this.x = Math.random() * canvasWidth;
                        this.y = Math.random() * canvasHeight;
                        this.radius = Math.random() * 20 + 8;
                        this.speedX = (Math.random() - 0.5) * 0.4;
                        this.speedY = (Math.random() - 0.5) * 0.4;
                        this.alpha = Math.random() * 0.22 + 0.08;
                        this.rotation = Math.random() * Math.PI;
                        this.rotationSpeed = (Math.random() - 0.5) * 0.005;
                    }
                    
                    update(w, h) {
                        this.x += this.speedX;
                        this.y += this.speedY;
                        this.rotation += this.rotationSpeed;
                        
                        if (this.x < -50) this.x = w + 50;
                        if (this.x > w + 50) this.x = -50;
                        if (this.y < -50) this.y = h + 50;
                        if (this.y > h + 50) this.y = -50;
                    }
                    
                    draw(ctx) {
                        ctx.save();
                        ctx.translate(this.x, this.y);
                        ctx.rotate(this.rotation);
                        ctx.strokeStyle = `rgba(167, 139, 250, ${this.alpha})`;
                        ctx.lineWidth = 1;
                        ctx.beginPath();
                        for (let i = 0; i < 6; i++) {
                            let angle = (Math.PI / 3) * i;
                            ctx.lineTo(this.radius * Math.cos(angle), this.radius * Math.sin(angle));
                        }
                        ctx.closePath();
                        ctx.stroke();
                        
                        ctx.fillStyle = `rgba(124, 58, 237, ${this.alpha * 0.1})`;
                        ctx.fill();
                        ctx.restore();
                    }
                }
                
                function resize() {
                    canvas.width = window.innerWidth;
                    canvas.height = window.innerHeight;
                    initParticles();
                }
                
                function initParticles() {
                    particles = [];
                    const count = Math.min(Math.floor((window.innerWidth * window.innerHeight) / 28000), 40);
                    for (let i = 0; i < count; i++) {
                        particles.push(new HexParticle(canvas.width, canvas.height));
                    }
                }
                
                function drawConnections() {
                    for (let i = 0; i < particles.length; i++) {
                        for (let j = i + 1; j < particles.length; j++) {
                            let dx = particles[i].x - particles[j].x;
                            let dy = particles[i].y - particles[j].y;
                            let distance = Math.sqrt(dx * dx + dy * dy);
                            
                            if (distance < 170) {
                                let force = (170 - distance) / 170;
                                let alpha = force * 0.12;
                                ctx.strokeStyle = `rgba(167, 139, 250, ${alpha})`;
                                ctx.lineWidth = 0.6;
                                ctx.beginPath();
                                ctx.moveTo(particles[i].x, particles[i].y);
                                ctx.lineTo(particles[j].x, particles[j].y);
                                ctx.stroke();
                            }
                        }
                    }
                }
                
                function animate() {
                    ctx.clearRect(0, 0, canvas.width, canvas.height);
                    for (let p of particles) {
                        p.update(canvas.width, canvas.height);
                        p.draw(ctx);
                    }
                    drawConnections();
                    requestAnimationFrame(animate);
                }
                
                window.addEventListener("resize", resize);
                resize();
                animate();
                
                // --- Part 2: Cinematic Progressive Loader (Vault "BOOM") ---
                let progress = 0;
                const subtitles = [
                    "DECRYPTING QUANTUM LINK...",
                    "INITIALIZING NEURAL COGNITION...",
                    "SYNCING CHAT DEPLOYMENT TERMINALS...",
                    "VERIFYING CRYPTOGRAPHIC PASSWORDS...",
                    "LAUNCHING ZERO-GPU HARDWARE ACCELERATORS...",
                    "SYSTEM ONLINE. READY FOR TRANSMISSION."
                ];
                
                const interval = setInterval(() => {
                    progress += Math.floor(Math.random() * 8) + 4;
                    if (progress >= 100) {
                        progress = 100;
                        clearInterval(interval);
                        
                        pct.innerText = "100%";
                        fill.style.width = "100%";
                        if (statusSubtitle) statusSubtitle.innerText = "AUTHENTICATION DECRYPTED";
                        
                        setTimeout(() => {
                            // Stage 1: Trigger BOOM screen flash
                            loader.classList.add("boom");
                            
                            setTimeout(() => {
                                // Stage 2: Slide open bi-parting vault doors
                                loader.classList.add("loaded");
                                
                                setTimeout(() => {
                                    // Remove overlay to optimize browser CPU
                                    loader.remove();
                                }, 1200);
                            }, 350);
                        }, 400);
                    } else {
                        pct.innerText = progress + "%";
                        fill.style.width = progress + "%";
                        const idx = Math.floor((progress / 100) * subtitles.length);
                        if (statusSubtitle && subtitles[idx]) {
                            statusSubtitle.innerText = subtitles[idx];
                        }
                    }
                }, 65);
            }
            
            if (document.readyState === "complete" || document.readyState === "interactive") {
                startLoaderAndParticles();
            } else {
                window.addEventListener("DOMContentLoaded", startLoaderAndParticles);
            }
        })();
        </script>
    """)
    
    with gr.Row():
        # Sidebar Panel
        with gr.Column(scale=1, min_width=250):
            session_dropdown = gr.Dropdown(
                choices=["Session 1"],
                value="Session 1",
                label="🔌 Active Session Node",
                interactive=True,
                elem_id="session-dropdown"
            )
            new_btn = gr.Button("➕ New Terminal Session", variant="secondary")
            reset_btn = gr.Button("🗑️ Reset All Sessions", variant="stop")
            
            gr.HTML("""
                <div style="margin-top: 25px; padding: 20px; background: rgba(18, 10, 36, 0.5); border-radius: 16px; border: 1.5px solid rgba(167, 139, 250, 0.15); box-shadow: 0 8px 32px rgba(0, 0, 0, 0.4);">
                    <div class="diagnostics-title">
                        <span style="display: inline-block; width: 10px; height: 10px; background: #c084fc; border-radius: 50%; box-shadow: 0 0 10px #c084fc; animation: buttonPulse 2s infinite alternate;"></span>
                        🚀 VL Intelligence Node
                    </div>
                    <p style="font-size: 0.85rem; color: #cbd5e1; line-height: 1.5; margin-bottom: 16px;">
                        State-of-the-art multimodal vision intelligence running optimal real-time inference on high-performance accelerators.
                    </p>
                    <div class="diagnostics-detail-line">
                        <span>MODEL CODENAME</span>
                        <span class="diagnostics-detail-val">RiShre Flash Core</span>
                    </div>
                    <div class="diagnostics-detail-line">
                        <span>KNOWLEDGE LIMIT</span>
                        <span class="diagnostics-detail-val">JAN 2026</span>
                    </div>
                    <div class="diagnostics-detail-line">
                        <span>ACCELERATOR</span>
                        <span class="diagnostics-detail-val">ZeroGPU Active</span>
                    </div>
                    <div class="diagnostics-detail-line">
                        <span>LATENCY SPEED</span>
                        <span class="diagnostics-detail-val">~45ms / Token</span>
                    </div>
                    <div class="diagnostics-detail-line">
                        <span>ECOSYSTEM</span>
                        <span class="diagnostics-detail-val" style="color: #f472b6;">Badge94 Secure</span>
                    </div>
                </div>
            """)
            
        # Chat Console Panel
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="Active Core Terminal",
                height=520,
                elem_id="chatbot"
            )
            
            # Suggestion Chips Row
            with gr.Row(elem_classes="suggestion-box"):
                suggest_btn_1 = gr.Button("🚀 Explain RiShre Flash Deep Engine", elem_classes="suggestion-btn")
                suggest_btn_2 = gr.Button("💻 Write a Python FastAPI server", elem_classes="suggestion-btn")
                suggest_btn_3 = gr.Button("🎨 CSS gradient animation palette", elem_classes="suggestion-btn")
                suggest_btn_4 = gr.Button("🔒 View system parameters & cutoff", elem_classes="suggestion-btn")
            
            with gr.Row():
                msg_input = gr.MultimodalTextbox(
                    show_label=False,
                    placeholder="Type a message or drag & drop an image...",
                    file_types=["image"],
                    scale=8,
                    elem_classes="input-container"
                )
                submit_btn = gr.Button("🚀 Transmit", variant="primary", scale=2, elem_classes="send-btn")

    # Wire up Suggested prompt loaders
    suggest_btn_1.click(load_suggest_1, outputs=[msg_input])
    suggest_btn_2.click(load_suggest_2, outputs=[msg_input])
    suggest_btn_3.click(load_suggest_3, outputs=[msg_input])
    suggest_btn_4.click(load_suggest_4, outputs=[msg_input])

    # Wire up chat submission sequence
    submit_btn.click(
        fn=user_submit,
        inputs=[msg_input, chatbot],
        outputs=[chatbot, msg_input],
        queue=True
    ).then(
        fn=bot_stream,
        inputs=[chatbot],
        outputs=[chatbot],
        queue=True
    )

    msg_input.submit(
        fn=user_submit,
        inputs=[msg_input, chatbot],
        outputs=[chatbot, msg_input],
        queue=True
    ).then(
        fn=bot_stream,
        inputs=[chatbot],
        outputs=[chatbot],
        queue=True
    )

    # Wire up session handling
    session_dropdown.change(
        fn=change_session,
        inputs=[session_dropdown, state, chatbot],
        outputs=[chatbot, state]
    )

    new_btn.click(
        fn=add_new_session,
        inputs=[state, chatbot],
        outputs=[session_dropdown, chatbot, state]
    )

    reset_btn.click(
        fn=reset_all_sessions,
        outputs=[session_dropdown, chatbot, state]
    )

# Port 7860 is default for Hugging Face Spaces Gradio SDK
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)