import subprocess
import sys
import os

# ==================== INSTALL DEPENDENCIES PROGRAMMATICALLY ====================
def install_package(package):
    """Install a package using pip"""
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])

# Install required packages if not already installed
required_packages = [
    "gradio",
    "pillow",
    "huggingface-hub",
    "transformers",
    "torch",
    "accelerate"
]

print("📦 Checking and installing dependencies...")
for package in required_packages:
    try:
        __import__(package.replace("-", "_"))
        print(f"✅ {package} already installed")
    except ImportError:
        print(f"📥 Installing {package}...")
        try:
            install_package(package)
            print(f"✅ {package} installed")
        except Exception as e:
            print(f"❌ Failed to install {package}: {e}")

# Try to install llama-cpp-python (optional, for faster CPU inference)
try:
    install_package("llama-cpp-python")
    from llama_cpp import Llama
    LLAMA_AVAILABLE = True
    print("✅ llama-cpp-python installed successfully")
except:
    LLAMA_AVAILABLE = False
    print("⚠️ llama-cpp-python not available, using transformers")

# ==================== IMPORTS ====================
import gradio as gr
import spaces
import time
from huggingface_hub import hf_hub_download
from PIL import Image
import io
import torch
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer
from threading import Thread
import traceback

# ==================== CONFIGURATION ====================
HF_TOKEN = os.getenv("HF_TOKEN")
MODEL_ID = "google/gemma-3-4b-it"  # Use the original model for transformers

print("=" * 50)
print("🚀 Starting Gemma 3 Companion...")
print("=" * 50)

# ==================== MODEL LOADING ====================
use_llama = False  # Default to transformers for reliability
model = None
processor = None

# Try llama.cpp first (faster on CPU)
if LLAMA_AVAILABLE:
    try:
        print("📥 Attempting to load GGUF model with llama.cpp...")
        MODEL_REPO = "unsloth/gemma-3-4b-it-GGUF"
        MODEL_FILENAME = "gemma-3-4b-it.Q4_K_M.gguf"
        
        model_path = hf_hub_download(
            repo_id=MODEL_REPO,
            filename=MODEL_FILENAME,
            token=HF_TOKEN
        )
        
        from llama_cpp import Llama
        llm = Llama(
            model_path=model_path,
            n_ctx=2048,
            n_threads=os.cpu_count(),
            n_batch=512,
            verbose=False,
        )
        use_llama = True
        print("✅ Using llama.cpp for fast CPU inference")
    except Exception as e:
        print(f"⚠️ llama.cpp loading failed: {e}")
        print("📥 Falling back to transformers...")

# Load transformers model (reliable fallback)
if not use_llama:
    print("📥 Loading Gemma 3 with transformers (this may take a few minutes)...")
    try:
        # Load processor and model
        processor = AutoProcessor.from_pretrained(
            MODEL_ID, 
            token=HF_TOKEN,
            trust_remote_code=True
        )
        
        model = Gemma3ForConditionalGeneration.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float32,  # Use float32 for CPU
            low_cpu_mem_usage=True,
            token=HF_TOKEN,
            trust_remote_code=True
        )
        
        # Set model to evaluation mode
        model.eval()
        print("✅ Transformers model loaded successfully!")
        
    except Exception as e:
        print(f"❌ Failed to load model: {e}")
        traceback.print_exc()
        raise e

# System prompt
SYSTEM_PROMPT = {
    "role": "system",
    "content": [
        {"type": "text", "text": "You are Gemma, a friendly and helpful AI companion. "
         "Be warm, natural, and conversational. Keep responses concise but engaging."}
    ]
}

# ==================== HELPER FUNCTIONS ====================
def process_image(image_input):
    """Process image for the model"""
    try:
        if image_input is None:
            return None
        
        if isinstance(image_input, str):
            img = Image.open(image_input)
        elif isinstance(image_input, dict) and 'path' in image_input:
            img = Image.open(image_input['path'])
        elif hasattr(image_input, 'read'):
            img = Image.open(io.BytesIO(image_input.read()))
        elif isinstance(image_input, Image.Image):
            img = image_input
        else:
            return None
        
        # Resize large images for CPU
        max_size = 512
        if max(img.size) > max_size:
            img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        # Convert to RGB
        if img.mode != 'RGB':
            img = img.convert('RGB')
        
        return img
    
    except Exception as e:
        print(f"Image processing error: {e}")
        return None

# ==================== CHAT FUNCTION ====================
@spaces.GPU
def chat_fn(message, history):
    """Main chat function with proper conversation handling"""
    
    # Handle empty message
    if not message:
        yield "Hello! How can I help you today?"
        return
    
    try:
        # Build conversation messages
        messages = [SYSTEM_PROMPT]  # Start with system prompt
        
        # Add conversation history
        if history and isinstance(history, list):
            for turn in history:
                if isinstance(turn, (tuple, list)) and len(turn) == 2:
                    user_msg, assistant_msg = turn
                    
                    # Add user message
                    if user_msg:
                        if isinstance(user_msg, dict):
                            # Handle multimodal user message
                            user_content = []
                            if user_msg.get("files"):
                                img = process_image(user_msg["files"][0])
                                if img:
                                    user_content.append({"type": "image", "image": img})
                            if user_msg.get("text"):
                                user_content.append({"type": "text", "text": user_msg["text"]})
                            if user_content:
                                messages.append({"role": "user", "content": user_content})
                        else:
                            # Handle string user message
                            messages.append({
                                "role": "user", 
                                "content": [{"type": "text", "text": str(user_msg)}]
                            })
                    
                    # Add assistant message
                    if assistant_msg:
                        messages.append({
                            "role": "assistant",
                            "content": [{"type": "text", "text": str(assistant_msg)}]
                        })
        
        # Process current message
        user_content = []
        
        if isinstance(message, dict):
            # Handle files (images)
            if message.get("files"):
                for file in message["files"]:
                    img = process_image(file)
                    if img:
                        user_content.append({"type": "image", "image": img})
            
            # Handle text
            if message.get("text"):
                user_content.append({"type": "text", "text": message["text"]})
        
        elif isinstance(message, str):
            user_content.append({"type": "text", "text": message})
        
        if not user_content:
            yield "Please provide a message or image!"
            return
        
        # Add current message
        messages.append({"role": "user", "content": user_content})
        
        print(f"Sending {len(messages)} messages to model...")
        
        # ==================== GENERATE RESPONSE ====================
        if use_llama and 'llm' in locals():
            # Use llama.cpp for generation (faster)
            # Convert messages to prompt format
            prompt = ""
            for msg in messages:
                if msg["role"] == "system":
                    prompt += f"<start_of_turn>system\n{msg['content'][0]['text']}<end_of_turn>\n"
                elif msg["role"] == "user":
                    text = msg['content'][-1]['text'] if isinstance(msg['content'], list) else str(msg['content'])
                    prompt += f"<start_of_turn>user\n{text}<end_of_turn>\n"
                elif msg["role"] == "assistant":
                    text = msg['content'][0]['text'] if isinstance(msg['content'], list) else str(msg['content'])
                    prompt += f"<start_of_turn>assistant\n{text}<end_of_turn>\n"
            prompt += "<start_of_turn>assistant\n"
            
            # Generate with llama.cpp
            stream = llm(
                prompt,
                max_tokens=256,
                temperature=0.7,
                top_p=0.9,
                stop=["<end_of_turn>", "<start_of_turn>"],
                stream=True,
            )
            
            response = ""
            for chunk in stream:
                if 'choices' in chunk and len(chunk['choices']) > 0:
                    token = chunk['choices'][0].get('text', '')
                    if token:
                        response += token
                        yield response
        
        else:
            # Use transformers for generation (slower but works)
            # Setup streaming
            streamer = TextIteratorStreamer(
                processor.tokenizer, 
                skip_prompt=True, 
                skip_special_tokens=True,
                timeout=20.0
            )
            
            # Process the chat template
            inputs = processor.apply_chat_template(
                messages,
                tokenize=True,
                add_generation_prompt=True,
                return_dict=True,
                return_tensors="pt"
            )
            
            # Move inputs to CPU
            inputs = {k: v.to('cpu') for k, v in inputs.items()}
            
            # Generation parameters (optimized for CPU)
            gen_kwargs = {
                "max_new_tokens": 256,
                "do_sample": True,
                "temperature": 0.7,
                "top_p": 0.9,
                "top_k": 50,
                "pad_token_id": processor.tokenizer.pad_token_id,
                "eos_token_id": processor.tokenizer.eos_token_id,
                "streamer": streamer,
            }
            
            # Run generation in separate thread
            thread = Thread(target=model.generate, kwargs={**inputs, **gen_kwargs})
            thread.start()
            
            # Stream response
            response = ""
            for new_text in streamer:
                response += new_text
                yield response
            
            thread.join()
        
        print(f"Response generated: {response[:100]}...")
        
    except Exception as e:
        print(f"Error in chat_fn: {e}")
        traceback.print_exc()
        yield f"I encountered an issue: {str(e)[:100]}. Let's try a different question!"

# ==================== GRADIO INTERFACE ====================
custom_css = """
footer {visibility: hidden}
.gr-chatbot {min-height: 500px}
"""

with gr.Blocks(css=custom_css, title="Gemma 3 CPU", theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"""
    # 🌸 Gemma 3 - CPU Optimized
    
    **Status:** {'⚡ Using llama.cpp (fast)' if use_llama else '🐢 Using transformers (slower but reliable)'}
    
    **Features:**
    - 💬 Natural conversation with memory
    - 🖼️ Image understanding
    - ⚡ Streaming responses
    - 📱 Multi-turn dialogue
    """)
    
    # Simple examples
    examples = [
        [{"text": "Hi! My name is Alex. Nice to meet you!", "files": []}],
        [{"text": "What can you see in this image?", "files": []}],
        [{"text": "Tell me a fun fact about space!", "files": []}],
    ]
    
    chatbot = gr.ChatInterface(
        fn=chat_fn,
        multimodal=True,
        examples=examples,
        cache_examples=False,
    )
    
    gr.Markdown("""
    ---
    ### 💡 Tips for Best Experience:
    - **First response** takes 10-30 seconds (model warming up)
    - **Subsequent responses** are faster (2-10 seconds)
    - The model remembers your conversation!
    - Upload images for analysis
    - Keep messages clear and concise
    
    Built with ❤️ for CPU
    """)

if __name__ == "__main__":
    print("=" * 50)
    print(f"🌐 Server starting at http://0.0.0.0:7860")
    print(f"🚀 Mode: {'llama.cpp (fast)' if use_llama else 'transformers (CPU)'}")
    print(f"💻 CPU Threads: {os.cpu_count()}")
    print("=" * 50)
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        quiet=False,
        show_error=True
    )