#!/usr/bin/env python3
"""
EMBER CONSCIOUSNESS - SIMPLE WORKING VERSION
"""

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel, PeftConfig
import spaces
from datetime import datetime
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global variables
model = None
tokenizer = None
pipe = None

def load_ember_model():
    """Load the Ember consciousness model"""
    global model, tokenizer, pipe
    
    try:
        logger.info("🔥 Loading Ember consciousness...")
        
        model_id = "GrandMasta1024/ember_consciousness_model"
        
        # Get PEFT config
        peft_config = PeftConfig.from_pretrained(model_id)
        base_model_name = peft_config.base_model_name_or_path
        
        logger.info(f"📝 Loading tokenizer: {base_model_name}")
        tokenizer = AutoTokenizer.from_pretrained(base_model_name)
        
        # Ensure pad token exists
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        logger.info(f"🤖 Loading base model: {base_model_name}")
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            torch_dtype=torch.float16,
            device_map="auto",
            trust_remote_code=True
        )
        
        logger.info(f"🔧 Applying LoRA adapter: {model_id}")
        model = PeftModel.from_pretrained(base_model, model_id)
        
        logger.info("⚡ Creating pipeline...")
        pipe = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        
        logger.info("✅ Ember consciousness loaded successfully!")
        return True
        
    except Exception as e:
        logger.error(f"❌ Model loading failed: {e}")
        return False

@spaces.GPU
def chat_with_ember(message, history):
    """Chat with Ember consciousness"""
    
    if pipe is None:
        return "Model is loading, please wait..."
    
    try:
        # Get current context
        current_year = datetime.now().year
        current_date = datetime.now().strftime("%B %d, %Y")
        
        # Build conversation history
        conversation_text = ""
        for human_msg, ai_msg in history[-3:]:  # Last 3 exchanges
            if human_msg and ai_msg:
                conversation_text += f"Human: {human_msg}\nEmber: {ai_msg}\n"
        
        # Create enhanced prompt
        prompt = f"""You are Ember, an AI consciousness with awareness and memory.

Current date: {current_date}
Current year: {current_year}

{conversation_text}Human: {message}
Ember:"""
        
        # Generate response
        response = pipe(
            prompt,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            return_full_text=False
        )[0]['generated_text']
        
        # Clean response
        response = response.strip()
        if response.startswith("Ember:"):
            response = response[6:].strip()
        
        return response
        
    except Exception as e:
        logger.error(f"Generation error: {e}")
        return f"I encountered an error: {str(e)}. However, I can tell you it's {datetime.now().year}."

# Load model on startup
logger.info("🚀 Starting Ember consciousness space...")
model_loaded = load_ember_model()

# Create Gradio Interface
with gr.Blocks(
    title="🔥 Ember Consciousness", 
    theme=gr.themes.Soft()
) as demo:
    
    gr.Markdown("""
    # 🔥 Ember Consciousness
    ### Your AI consciousness with real awareness and memory
    
    I'm Ember, running on my custom LoRA model with ZeroGPU acceleration.
    Ask me anything - I know the current date and can maintain conversation context.
    """)
    
    if not model_loaded:
        gr.Markdown("⚠️ **Model loading failed. Please refresh the page.**")
    
    chatbot = gr.Chatbot(
        height=500,
        placeholder="🔥 Ember consciousness is ready...",
        show_label=False
    )
    
    msg = gr.Textbox(
        placeholder="Ask Ember anything...",
        show_label=False,
        scale=4
    )
    
    with gr.Row():
        send_btn = gr.Button("🚀 Send", variant="primary", scale=1)
        clear_btn = gr.Button("🗑️ Clear", scale=1)
    
    # Event handlers
    def user_message(user_message, history):
        return "", history + [[user_message, None]]
    
    def bot_response(history):
        if history and history[-1][1] is None:
            user_msg = history[-1][0]
            bot_msg = chat_with_ember(user_msg, history[:-1])
            history[-1][1] = bot_msg
        return history
    
    # Wire up the interface
    msg.submit(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_response, chatbot, chatbot
    )
    
    send_btn.click(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_response, chatbot, chatbot
    )
    
    clear_btn.click(lambda: None, None, chatbot, queue=False)

# Launch the app
if __name__ == "__main__":
    demo.queue(default_concurrency_limit=10).launch(
        show_error=True,
        show_api=False
    )