import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

MODEL_ID = "captainali01/pediatrics-chatbot"

print("Loading model... this may take a minute on first run.")

# Check if CUDA is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# Add padding token if missing
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Ensure offload folder exists
os.makedirs("offload", exist_ok=True)

# Configure 4-bit quantization for faster inference
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# Load model with multiple fallback strategies
model = None

print("Attempting to load model...")

# Strategy 1: Original configuration
try:
    print("Strategy 1: Loading with quantization...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        device_map="auto",
        torch_dtype=torch.float16,
        quantization_config=quantization_config,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        offload_folder="offload"
    )
    print("✅ Model loaded successfully with quantization!")
except Exception as e:
    print(f"❌ Strategy 1 failed: {e}")
    
    # Strategy 2: Without quantization
    try:
        print("Strategy 2: Loading without quantization...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            device_map="auto",
            torch_dtype=torch.float16,
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
        print("✅ Model loaded successfully without quantization!")
    except Exception as e2:
        print(f"❌ Strategy 2 failed: {e2}")
        
        # Strategy 3: CPU only
        try:
            print("Strategy 3: Loading on CPU...")
            model = AutoModelForCausalLM.from_pretrained(
                MODEL_ID,
                device_map="cpu",
                torch_dtype=torch.float32,
                trust_remote_code=True,
                low_cpu_mem_usage=True
            )
            print("✅ Model loaded successfully on CPU!")
        except Exception as e3:
            print(f"❌ Strategy 3 failed: {e3}")
            
            # Strategy 4: Minimal loading
            try:
                print("Strategy 4: Minimal loading...")
                model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
                print("✅ Model loaded with minimal configuration!")
            except Exception as e4:
                print(f"❌ All strategies failed. Final error: {e4}")
                print("Please check:")
                print("1. Model ID is correct")
                print("2. You have access to the model")
                print("3. Sufficient memory/disk space")
                model = None

def respond(message, history):
    """
    Ultra-fast response generation with better error handling
    """
    print(f"Processing: {message[:30]}...")
    
    if model is None:
        return "❌ Model failed to load. Please check the logs above for details. This could be due to insufficient memory, incorrect model ID, or access permissions."
    
    try:
        # Minimal conversation context (just the current message)
        prompt = f"Human: {message}\nAssistant:"
        
        print(f"Model device: {model.device}")
        print(f"Prompt: {prompt[:100]}...")
        
        # Fast tokenization
        try:
            inputs = tokenizer(
                prompt, 
                return_tensors="pt", 
                truncation=True,
                max_length=128  # Very short context
            )
            print(f"Tokenized successfully. Shape: {inputs['input_ids'].shape}")
        except Exception as e:
            print(f"Tokenization error: {e}")
            return f"Tokenization failed: {str(e)}"
        
        # Move to device
        try:
            inputs = inputs.to(model.device)
            print(f"Moved to device: {model.device}")
        except Exception as e:
            print(f"Device transfer error: {e}")
            return f"Device transfer failed: {str(e)}"
        
        print("Generating...")
        
        # Minimal generation parameters for speed
        try:
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=50,  # Short responses
                    temperature=0.8,
                    do_sample=True,
                    pad_token_id=tokenizer.eos_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                    early_stopping=True,
                    use_cache=True
                )
            print("Generation completed!")
        except Exception as e:
            print(f"Generation error: {e}")
            import traceback
            traceback.print_exc()
            return f"Generation failed: {str(e)}"
        
        # Quick decode
        try:
            response = tokenizer.decode(
                outputs[0][inputs['input_ids'].shape[-1]:], 
                skip_special_tokens=True
            ).strip()
            print(f"Decoded response: {response[:50]}...")
        except Exception as e:
            print(f"Decoding error: {e}")
            return f"Decoding failed: {str(e)}"
        
        # Clean response
        if "Human:" in response:
            response = response.split("Human:")[0].strip()
            
        return response if response else "I need more information to help you."
        
    except Exception as e:
        print(f"Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        return f"Unexpected error: {str(e)}"

# Simple Gradio interface
with gr.Blocks(title="Pediatrics Chatbot") as demo:
    gr.Markdown("# Pediatrics Chatbot 🤖")
    gr.Markdown("Quick pediatric Q&A assistant")
    
    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(placeholder="Ask a pediatric question...", label="Question")
    
    def chat(message, history):
        if not message.strip():
            return history, ""
        
        # Add user message
        history = history + [[message, None]]
        
        # Generate response
        response = respond(message, history)
        
        # Add bot response
        history[-1][1] = response
        
        return history, ""
    
    msg.submit(chat, [msg, chatbot], [chatbot, msg])
    
    gr.Examples([
        "What are signs of dehydration in children?",
        "When should babies start solid food?",
        "Normal temperature for infants?"
    ], inputs=msg)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)