import gradio as gr
import torch
import os
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from peft import LoraConfig, prepare_model_for_kbit_training
from trl import SFTTrainer
import json

def fine_tune_model():
    """Fine-tune model for personal assistant with progress updates"""
    
    try:
        yield "🔄 Starting fine-tuning process...\n"
        
        # Use a manageable model for Spaces
        model_name = "microsoft/DialoGPT-medium"
        yield f"📥 Loading model: {model_name}\n"
        
        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
        tokenizer.pad_token = tokenizer.eos_token
        yield "✅ Tokenizer loaded successfully\n"
        
        # Load model
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_map="auto" if torch.cuda.is_available() else None,
        )
        yield "✅ Model loaded successfully\n"
        
        # Load dataset
        yield "📊 Loading training dataset...\n"
        dataset = load_dataset("databricks/databricks-dolly-15k", split="train[:1000]")
        yield f"✅ Dataset loaded: {len(dataset)} examples\n"
        
        # Format dataset
        def format_example(example):
            instruction = example["instruction"]
            response = example["response"]
            context = example.get("context", "")
            
            if context:
                text = f"Human: {instruction}\nContext: {context}\nAssistant: {response}<|endoftext|>"
            else:
                text = f"Human: {instruction}\nAssistant: {response}<|endoftext|>"
            return {"text": text}
        
        processed_dataset = dataset.map(format_example)
        yield "✅ Dataset formatted for training\n"
        
        # LoRA configuration
        peft_config = LoraConfig(
            lora_alpha=16,
            lora_dropout=0.1,
            r=8,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["c_attn", "c_proj"],
        )
        yield "✅ LoRA configuration set\n"
        
        # Training arguments
        training_arguments = TrainingArguments(
            output_dir="./results",
            per_device_train_batch_size=1,
            gradient_accumulation_steps=4,
            logging_steps=10,
            save_steps=100,
            learning_rate=5e-5,
            num_train_epochs=1,
            warmup_steps=50,
            remove_unused_columns=False,
            dataloader_pin_memory=False,
        )
        yield "✅ Training configuration set\n"
        
        # Create trainer
        trainer = SFTTrainer(
            model=model,
            train_dataset=processed_dataset,
            tokenizer=tokenizer,
            args=training_arguments,
            peft_config=peft_config,
            dataset_text_field="text",
        )
        yield "🚀 Starting model training...\n"
        
        # Start training
        trainer.train()
        yield "✅ Training completed successfully!\n"
        
        # Save model
        trainer.save_model("./fine_tuned_assistant")
        tokenizer.save_pretrained("./fine_tuned_assistant")
        yield "💾 Model saved successfully!\n"
        
        yield "🎉 Fine-tuning process completed! Your personal assistant is ready!\n"
        
    except Exception as e:
        yield f"❌ Error during fine-tuning: {str(e)}\n"

def chat_with_assistant(message, history):
    """Simple chat interface for testing"""
    if not message.strip():
        return history, ""
    
    # Placeholder response (you can implement actual model inference later)
    responses = [
        f"Thank you for your message: '{message}'. As your personal assistant, I'm here to help!",
        f"I appreciate you reaching out about '{message}'. How can I assist you further?",
        f"Regarding '{message}', I'm happy to help. What specific assistance do you need?",
        f"I understand you mentioned '{message}'. As your supportive assistant, I'm here for you!"
    ]
    
    import random
    response = random.choice(responses)
    
    history.append((message, response))
    return history, ""

# Create Gradio interface
with gr.Blocks(title="Personal Assistant Fine-Tuning", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🤖 Personal Assistant Fine-Tuning on Hugging Face Spaces
    
    Welcome to your personal AI assistant training platform! This space allows you to:
    - Fine-tune a language model to be your personal assistant
    - Test the assistant's responses
    - Create a kind, supportive, and intelligent AI companion
    """)
    
    with gr.Tab("🔧 Fine-Tune Model"):
        gr.Markdown("""
        ### Train Your Personal Assistant
        
        Click the button below to start fine-tuning your model. This process will:
        1. Load a pre-trained conversational model
        2. Train it on assistant-style conversations
        3. Optimize it to be supportive and helpful
        
        **Note:** Training may take 30-60 minutes depending on your hardware tier.
        """)
        
        tune_button = gr.Button("🚀 Start Fine-Tuning", variant="primary", size="lg")
        tune_output = gr.Textbox(
            label="Training Progress", 
            lines=15, 
            max_lines=20,
            show_copy_button=True,
            interactive=False
        )
        
        tune_button.click(
            fn=fine_tune_model,
            outputs=tune_output
        )
    
    with gr.Tab("💬 Test Assistant"):
        gr.Markdown("""
        ### Chat with Your Assistant
        
        Once fine-tuning is complete, use this interface to test your personal assistant.
        Your assistant is designed to be kind, supportive, and intelligent.
        """)
        
        chatbot = gr.Chatbot(
            height=500,
            show_copy_button=True,
            bubble_full_width=False
        )
        
        with gr.Row():
            msg = gr.Textbox(
                label="Your message",
                placeholder="Type your message here...",
                lines=2,
                scale=4
            )
            send_btn = gr.Button("Send", variant="primary", scale=1)
        
        clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
        
        # Event handlers
        msg.submit(chat_with_assistant, [msg, chatbot], [chatbot, msg])
        send_btn.click(chat_with_assistant, [msg, chatbot], [chatbot, msg])
        clear_btn.click(lambda: [], None, chatbot, queue=False)
    
    with gr.Tab("ℹ️ Info"):
        gr.Markdown("""
        ### About This Space
        
        This Hugging Face Space is designed to help you create your own personal AI assistant through fine-tuning.
        
        **Features:**
        - 🤖 Fine-tune conversational AI models
        - 💬 Interactive chat interface for testing
        - 🎯 Optimized for supportive, intelligent responses
        - 📊 Real-time training progress monitoring
        
        **Hardware Recommendations:**
        - **Free CPU:** Good for testing the interface
        - **T4 Small GPU ($0.60/hr):** Recommended for actual training
        - **A10G Small ($1.05/hr):** Faster training, better performance
        
        **Tips for Success:**
        1. Start with free tier to test everything works
        2. Upgrade to GPU when ready for actual training
        3. Monitor training progress in the Fine-Tune tab
        4. Test your assistant in the Chat tab after training
        
        Created with ❤️ for building personal AI assistants!
        """)

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