import gradio as gr
import torch
import os

# ===== Import sentence-transformers with error handling =====
try:
    from sentence_transformers import SentenceTransformer, util
    USE_SEMANTIC = True
    print("✓ Sentence transformers loaded successfully")
except ImportError as e:
    USE_SEMANTIC = False
    print(f"⚠ Warning: sentence-transformers not available: {e}")
    print("Using fallback keyword matching instead")

# Predefined Q&A pairs
qa_pairs = {
    "What is Augmented Reality?": "Hello user!! 👋\nAugmented Reality (AR) overlays digital content onto real-world environments.",
    "What is Virtual Reality?": "Hello user!! 👋\nVirtual Reality (VR) immerses users in a completely digital environment.",
    "What is SLAM in AR?": "Hello user!! 👋\nSLAM means Simultaneous Localization and Mapping, used by AR devices to track position.",
    "Explain foveated rendering in VR.": "Hello user!! 👋\nFoveated rendering optimizes performance by rendering only the focal area in high detail.",
    "What causes latency in VR systems?": "Hello user!! 👋\nLatency occurs due to delays between user movements and visual updates, mainly from hardware lag.",
    "How does spatial mapping work in AR?": "Hello user!! 👋\nSpatial mapping lets AR devices construct 3D meshes of real spaces for accurate content placement.",
    "What are haptics in VR?": "Hello user!! 👋\nHaptics simulate touch sensations in VR through vibrations or pressure feedback.",
    "What is occlusion in AR?": "Hello user!! 👋\nOcclusion ensures virtual objects appear correctly behind real-world ones for realism.",
    "What is the future of AR/VR?": "Hello user!! 👋\nThe future is mixed reality, AI-driven spatial computing, and haptic integration.",
    "What are digital twins in AR/VR?": "Hello user!! 👋\nDigital twins are virtual replicas of real-world entities used for monitoring and simulation."
}

questions = list(qa_pairs.keys())
model = None
question_embeddings = None

# Initialize model with local priority and timeout
if USE_SEMANTIC:
    try:
        # Priority 1: Local model (upload /models/all-MiniLM-L6-v2/ folder to repo)
        local_path = "./models/all-MiniLM-L6-v2"
        if os.path.exists(local_path):
            model = SentenceTransformer(local_path, local_files_only=True)
            print("✓ Local model loaded successfully")
        else:
            # Priority 2: Online with cache
            import torch.hub
            torch.hub.set_dir("./cache")  # Use local cache dir
            model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
            print("✓ Online model loaded successfully")
        
        question_embeddings = model.encode(questions, convert_to_tensor=True, normalize_embeddings=True)
        print("✓ Model embeddings computed successfully")
    except Exception as e:
        print(f"⚠ Model loading failed: {e}")
        USE_SEMANTIC = False

def keyword_fallback(user_input):
    """Simple keyword-based matching as fallback"""
    user_lower = user_input.lower()
    
    if "augmented reality" in user_lower or ("ar" in user_lower and "what is" in user_lower):
        return qa_pairs["What is Augmented Reality?"]
    elif "virtual reality" in user_lower or ("vr" in user_lower and "what is" in user_lower):
        return qa_pairs["What is Virtual Reality?"]
    elif "slam" in user_lower:
        return qa_pairs["What is SLAM in AR?"]
    elif "foveated" in user_lower or "rendering" in user_lower:
        return qa_pairs["Explain foveated rendering in VR."]
    elif "latency" in user_lower:
        return qa_pairs["What causes latency in VR systems?"]
    elif "spatial mapping" in user_lower or "mapping" in user_lower:
        return qa_pairs["How does spatial mapping work in AR?"]
    elif "haptic" in user_lower or "touch" in user_lower:
        return qa_pairs["What are haptics in VR?"]
    elif "occlusion" in user_lower:
        return qa_pairs["What is occlusion in AR?"]
    elif "future" in user_lower:
        return qa_pairs["What is the future of AR/VR?"]
    elif "digital twin" in user_lower:
        return qa_pairs["What are digital twins in AR/VR?"]
    
    return "Hello user!! 👋\nI'm not sure about that. Try asking about AR, VR, SLAM, haptics, or other VR/AR concepts!"

def semantic_chatbot(user_input):
    if not user_input.strip():
        return "Please ask a valid question!"
    
    if USE_SEMANTIC and model is not None:
        try:
            with torch.no_grad():
                input_emb = model.encode(user_input, convert_to_tensor=True, normalize_embeddings=True)
                scores = util.cos_sim(input_emb, question_embeddings)[0]
                best_match = torch.argmax(scores).item()
            return qa_pairs[questions[best_match]]
        except Exception as e:
            print(f"Semantic matching error: {e}, falling back to keywords")
            return keyword_fallback(user_input)
    else:
        return keyword_fallback(user_input)

# Cyberpunk CSS
css_dark = """
.gradio-container {background-color: #0b0b0b; font-family: 'Orbitron', sans-serif; color: #0ff;}
input, textarea {background-color: #0b0b0b; color: #0ff; border: 1px solid #0ff;}
button {background-color: #222; color: #fff; border: 1px solid #0ff;}
button:hover {box-shadow: 0 0 10px #ff00ff; transform: scale(1.05);}
"""

# Gradio UI
with gr.Blocks(css=css_dark, theme="gradio/soft") as demo:
    gr.Markdown("## ✨ Vbot")
    gr.Markdown("Explore AR/VR concepts through semantic answers powered by transformers.")
    
    user_input = gr.Textbox(label="Ask a Question 🌐", placeholder="e.g. What is SLAM in AR?")
    output = gr.Textbox(label="Answer", interactive=False)
    submit = gr.Button("Get Answer")
    
    submit.click(fn=semantic_chatbot, inputs=user_input, outputs=output)
    
    gr.Markdown("### 🔍 Try These:")
    with gr.Row():
        for q in list(qa_pairs.keys())[:5]:
            btn = gr.Button(q[:30] + "..." if len(q) > 30 else q, size="sm")
            btn.click(fn=lambda x=q: semantic_chatbot(x), inputs=None, outputs=output)

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