"""
Vish AI - Virtual Intelligent System Hub
Lightweight multimodal AI assistant optimized for Hugging Face Spaces
Production-ready version
"""

import gradio as gr
import os
from datetime import datetime
import time
import importlib

# Supabase imports
try:
    from supabase import create_client, Client
    SUPABASE_AVAILABLE = True
except ImportError:
    SUPABASE_AVAILABLE = False
    print("⚠️ Supabase not available - running in demo mode")

# AI model imports
torch = None

try:
    torch = importlib.import_module("torch")
    AI_AVAILABLE = True
except ImportError:
    AI_AVAILABLE = False
    torch = None

if not AI_AVAILABLE:
    print("⚠️ AI models not available - using fallback mode")

# Supabase configuration
SUPABASE_URL = os.getenv("NEXT_PUBLIC_SUPABASE_URL", "https://lyebtceryednzafhyunq.supabase.co")
SUPABASE_KEY = os.getenv("NEXT_PUBLIC_SUPABASE_ANON_KEY", "")

# Initialize Supabase client
supabase = None
if SUPABASE_AVAILABLE and SUPABASE_KEY:
    try:
        supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
        print("✅ Supabase connected successfully")
    except Exception as e:
        print(f"⚠️ Supabase initialization error: {e}")
else:
    print("⚠️ Supabase credentials not configured")

# Global variables for unified model
phi3_model = None
phi3_tokenizer = None

def initialize_models():
    """Initialize Phi-3 unified model for all AI tasks"""
    global phi3_model, phi3_tokenizer
    
    if not AI_AVAILABLE:
        print("⚠️ AI libraries not available - using demo mode")
        return False
    
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
        import traceback
        
        # Load Phi-3 Mini - Unified model for all tasks (~7.4GB)
        print("📥 Loading Phi-3 Mini unified model...")
        print("   Model: microsoft/Phi-3-mini-4k-instruct")
        print("   Capabilities: Chat, Summarization, Sentiment Analysis")
        print("   This may take 5-15 minutes on first run (downloading ~7GB)...")
        
        # Load tokenizer
        print("   Loading tokenizer...")
        phi3_tokenizer = AutoTokenizer.from_pretrained(
            "microsoft/Phi-3-mini-4k-instruct",
            trust_remote_code=True
        )
        print("   ✅ Tokenizer loaded")
        
        # Load model with CPU optimization for Hugging Face Spaces
        print("   Loading model (this is the slow part)...")
        phi3_model = AutoModelForCausalLM.from_pretrained(
            "microsoft/Phi-3-mini-4k-instruct",
            device_map="cpu",
            torch_dtype=torch.float32,  # Use float32 for CPU
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
        
        print("✅ Phi-3 Mini model loaded successfully!")
        print("🎉 Unified model ready for all tasks!")
        print(f"   Model parameters: {phi3_model.num_parameters():,}")
        return True
    except Exception as e:
        print(f"❌ Error loading Phi-3 model: {e}")
        print("Detailed error:")
        import traceback
        traceback.print_exc()
        return False

def verify_user_token(token: str) -> dict:
    """Verify Supabase user authentication token"""
    if not supabase or not token:
        return {"authenticated": False, "user": None}
    
    try:
        user = supabase.auth.get_user(token)
        return {"authenticated": True, "user": user.user.email if user.user else None}
    except Exception as e:
        return {"authenticated": False, "error": str(e)}

def log_interaction(user_email: str, prompt: str, response: str, model_type: str):
    """Log user interactions to Supabase"""
    if not supabase:
        return
    
    try:
        data = {
            "user_email": user_email,
            "prompt": prompt,
            "response": response,
            "model_type": model_type,
            "timestamp": datetime.utcnow().isoformat()
        }
        supabase.table("vish_ai_logs").insert(data).execute()
    except Exception as e:
        print(f"Logging error: {e}")

def generate_phi3_response(prompt: str, max_new_tokens: int = 256, temperature: float = 0.7) -> str:
    """Generate response using Phi-3 model"""
    if not phi3_model or not phi3_tokenizer:
        return None
    
    try:
        # Format prompt for Phi-3 instruct format
        messages = [{"role": "user", "content": prompt}]
        
        # Apply chat template
        formatted_prompt = phi3_tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        # Tokenize
        inputs = phi3_tokenizer(formatted_prompt, return_tensors="pt")
        
        # Generate
        with torch.no_grad():
            outputs = phi3_model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=True,
                top_p=0.9,
                pad_token_id=phi3_tokenizer.eos_token_id
            )
        
        # Decode and extract response
        full_response = phi3_tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the assistant's response (after the prompt)
        if "<|assistant|>" in full_response:
            response = full_response.split("<|assistant|>")[-1].strip()
        else:
            response = full_response[len(formatted_prompt):].strip()
        
        return response
    except Exception as e:
        print(f"Error generating response: {e}")
        return None

def chat_with_vish(message: str, history: list, auth_token: str = "") -> str:
    """Main chat function with authentication"""
    
    # Verify authentication (optional - remove if you want public access)
    user_info = verify_user_token(auth_token) if auth_token else {"authenticated": False}
    user_email = user_info.get("user", "anonymous")
    
    if not AI_AVAILABLE or not phi3_model:
        # Fallback response when AI is not available
        fallback = "🤖 **Vish AI (Demo Mode)**\n\nYou said: _{}_\n\n⚠️ AI models are not loaded. This happens when:\n- Running in Python 3.14 (PyTorch not supported)\n- First deployment (models downloading)\n\n✅ **This will work perfectly on Hugging Face Spaces!**\n\n_Response time: <0.1s_".format(message)
        history.append([message, fallback])
        return history
    
    try:
        start_time = time.time()
        
        # Build context from history
        context = ""
        if history:
            for h in history[-3:]:  # Last 3 exchanges for context
                context += f"User: {h[0]}\nAssistant: {h[1]}\n"
        
        # Create prompt with context
        prompt = f"{context}User: {message}\nAssistant:"
        if context:
            prompt = f"Previous conversation:\n{context}\nCurrent question: {message}\n\nProvide a helpful and concise response:"
        else:
            prompt = f"Question: {message}\n\nProvide a helpful and concise response:"
        
        # Generate response using Phi-3
        assistant_response = generate_phi3_response(prompt, max_new_tokens=200, temperature=0.7)
        
        if not assistant_response:
            assistant_response = "I apologize, but I encountered an error generating a response. Please try again."
        
        elapsed_time = time.time() - start_time
        
        # Log interaction
        log_interaction(user_email, message, assistant_response, "chat")
        
        final_response = f"{assistant_response}\n\n⚡ _Response time: {elapsed_time:.2f}s_"
        history.append([message, final_response])
        return history
        
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        history.append([message, error_msg])
        return history

def summarize_text(text: str, auth_token: str = "") -> str:
    """Summarize long text using Phi-3"""
    user_info = verify_user_token(auth_token) if auth_token else {"authenticated": False}
    user_email = user_info.get("user", "anonymous")
    
    if not AI_AVAILABLE or not phi3_model:
        # Fallback summary
        word_count = len(text.split())
        return f"📝 **Summary (Demo Mode)**\n\nReceived {word_count} words.\n\nFirst 150 characters:\n_{text[:150]}_...\n\n⚠️ Full AI summarization available on Hugging Face Spaces!\n\n_Processing time: <0.1s_"
    
    try:
        if len(text.split()) < 50:
            return "⚠️ Text is too short to summarize. Please provide at least 50 words."
        
        start_time = time.time()
        
        # Truncate if too long (model context limit)
        max_chars = 3000
        if len(text) > max_chars:
            text = text[:max_chars] + "..."
        
        # Create summarization prompt
        prompt = f"Summarize the following text concisely in 2-3 sentences:\n\n{text}\n\nSummary:"
        
        # Generate summary using Phi-3
        summary = generate_phi3_response(prompt, max_new_tokens=150, temperature=0.3)
        
        if not summary:
            return "❌ Error generating summary. Please try again."
        
        elapsed_time = time.time() - start_time
        
        log_interaction(user_email, text[:100], summary, "summarization")
        
        return f"{summary}\n\n⚡ _Processing time: {elapsed_time:.2f}s_"
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

def analyze_sentiment(text: str, auth_token: str = "") -> str:
    """Analyze sentiment of text using Phi-3"""
    user_info = verify_user_token(auth_token) if auth_token else {"authenticated": False}
    user_email = user_info.get("user", "anonymous")
    
    if not AI_AVAILABLE or not phi3_model:
        # Simple fallback sentiment
        positive_words = ['good', 'great', 'excellent', 'happy', 'love', 'wonderful', 'amazing', 'fantastic', 'brilliant']
        negative_words = ['bad', 'terrible', 'awful', 'hate', 'sad', 'horrible', 'worst', 'poor', 'disappointing']
        
        text_lower = text.lower()
        pos_count = sum(1 for word in positive_words if word in text_lower)
        neg_count = sum(1 for word in negative_words if word in text_lower)
        
        if pos_count > neg_count:
            emoji, label, score = "😊", "POSITIVE", 0.85
        elif neg_count > pos_count:
            emoji, label, score = "😞", "NEGATIVE", 0.85
        else:
            emoji, label, score = "😐", "NEUTRAL", 0.50
        
        return f"{emoji} **{label}** (Demo - Simple keyword detection)\n\nConfidence: ~{score:.0%}\n\n⚠️ Full AI sentiment analysis available on Hugging Face Spaces!\n\n_Analysis time: <0.1s_"
    
    try:
        start_time = time.time()
        
        # Create sentiment analysis prompt
        prompt = f"Analyze the sentiment of the following text. Respond with only one word: POSITIVE, NEGATIVE, or NEUTRAL.\n\nText: {text[:500]}\n\nSentiment:"
        
        # Generate sentiment using Phi-3
        result = generate_phi3_response(prompt, max_new_tokens=10, temperature=0.1)
        
        if not result:
            return "❌ Error analyzing sentiment. Please try again."
        
        # Parse result
        result_upper = result.upper().strip()
        if "POSITIVE" in result_upper:
            label = "POSITIVE"
            emoji = "😊"
        elif "NEGATIVE" in result_upper:
            label = "NEGATIVE"
            emoji = "😞"
        else:
            label = "NEUTRAL"
            emoji = "�"
        
        elapsed_time = time.time() - start_time
        
        log_interaction(user_email, text[:100], f"{label}", "sentiment")
        
        return f"{emoji} **{label}**\n\n⚡ _Analysis time: {elapsed_time:.2f}s_"
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

def get_model_info() -> str:
    """Get information about loaded models"""
    info = """
    ## 🤖 Vish AI - Unified AI Model
    
    **Powered by Microsoft Phi-3 Mini 4K Instruct:**
    - Model: microsoft/Phi-3-mini-4k-instruct
    - Size: ~7.4GB (optimized for CPU)
    - Context: 4K tokens
    - Capabilities: Chat, Summarization, Sentiment Analysis
    
    **Performance:**
    - Chat: ~1-3s per response
    - Summarization: ~2-4s per summary
    - Sentiment Analysis: ~0.5-2s per analysis
    
    **Features:**
    - Single unified model for all tasks
    - Fine-tunable for custom requirements
    - Optimized for CPU inference
    - Production-ready architecture
    
    **Advantages over previous setup:**
    - Better quality responses (3.8B parameters vs 82M-300M)
    - Consistent performance across all tasks
    - Single model to maintain and fine-tune
    - More context-aware understanding
    """
    return info

# Initialize models on startup
print("=" * 60)
print("🚀 Initializing Vish AI - Production Ready")
print("=" * 60)
print(f"Python Version: 3.x")
print(f"AI Available: {AI_AVAILABLE}")
print(f"Supabase Available: {SUPABASE_AVAILABLE}")
print("=" * 60)

if AI_AVAILABLE:
    print("\n🔄 Starting Phi-3 model initialization...")
    models_loaded = initialize_models()
    if models_loaded:
        print("\n✅ All systems ready!")
    else:
        print("\n⚠️ Running in demo mode")
else:
    print("\n⚠️ AI libraries not available - running in demo mode")
    print("💡 This is normal for Python 3.14 - deploy to Hugging Face Spaces for full AI!")

print("=" * 60)

# Create Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="Vish AI") as demo:
    # Dynamic header based on AI availability
    if AI_AVAILABLE and phi3_model:
        status_badge = "🟢 **PRODUCTION** - Phi-3 AI Model Active"
    else:
        status_badge = "🟡 **DEMO MODE** - Deploy to Hugging Face for Full AI"
    
    gr.Markdown(f"""
    # 🌟 Vish AI - Virtual Intelligent System Hub
    ### Lightweight, Fast, Multimodal AI Assistant
    
    {status_badge}
    
    Optimized for Hugging Face Spaces | Powered by Supabase
    """)
    
    with gr.Tabs():
        # Chat Tab
        with gr.Tab("💬 Chat Assistant"):
            with gr.Row():
                with gr.Column(scale=4):
                    chatbot = gr.Chatbot(height=400, label="Vish AI Chat", type="tuples")
                    msg = gr.Textbox(
                        label="Your Message",
                        placeholder="Ask me anything...",
                        lines=2
                    )
                    with gr.Row():
                        submit = gr.Button("Send", variant="primary")
                        clear = gr.Button("Clear")
                
                with gr.Column(scale=1):
                    auth_token_chat = gr.Textbox(
                        label="🔐 Auth Token (Optional)",
                        type="password",
                        placeholder="Supabase JWT token",
                        lines=3
                    )
                    gr.Markdown("""
                    **Usage Tips:**
                    - Just type and chat!
                    - No token needed for demo
                    - Add token for logging
                    """)
            
            def respond(message, history, token):
                return chat_with_vish(message, history or [], token)
            
            submit.click(respond, inputs=[msg, chatbot, auth_token_chat], outputs=chatbot)
            msg.submit(respond, inputs=[msg, chatbot, auth_token_chat], outputs=chatbot)
            clear.click(lambda: [], None, chatbot, queue=False)
        
        # Summarization Tab
        with gr.Tab("📝 Text Summarizer"):
            with gr.Row():
                with gr.Column():
                    input_text = gr.Textbox(
                        label="Enter Text to Summarize",
                        placeholder="Paste your long text here (minimum 50 words)...",
                        lines=10
                    )
                    auth_token_sum = gr.Textbox(
                        label="Auth Token (Optional)",
                        type="password"
                    )
                    summarize_btn = gr.Button("Summarize", variant="primary")
                
                with gr.Column():
                    summary_output = gr.Textbox(
                        label="Summary",
                        lines=10
                    )
            
            summarize_btn.click(summarize_text, [input_text, auth_token_sum], summary_output)
        
        # Sentiment Analysis Tab
        with gr.Tab("😊 Sentiment Analysis"):
            with gr.Row():
                with gr.Column():
                    sentiment_input = gr.Textbox(
                        label="Enter Text to Analyze",
                        placeholder="How do you feel about this?",
                        lines=5
                    )
                    auth_token_sent = gr.Textbox(
                        label="Auth Token (Optional)",
                        type="password"
                    )
                    analyze_btn = gr.Button("Analyze Sentiment", variant="primary")
                
                with gr.Column():
                    sentiment_output = gr.Textbox(
                        label="Sentiment Result",
                        lines=5
                    )
            
            analyze_btn.click(analyze_sentiment, [sentiment_input, auth_token_sent], sentiment_output)
        
        # Model Info Tab
        with gr.Tab("ℹ️ Model Info"):
            gr.Markdown(get_model_info())
    
    gr.Markdown("""
    ---
    **VIJ Project** | Powered by Supabase & Hugging Face | Built with ❤️ by Vishwas
    """)

# Launch the app
if __name__ == "__main__":
    demo.queue()  # Enable queuing for better performance
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )