"""
AI Python Code Model Trainer
Hugging Face Space for continuous training with auto-resume
Username: himu1780 | Model: ai-python-model

FINAL VERSION - All optimizations applied
"""

import os
import gc
import gradio as gr
import threading
import time
from datetime import datetime
from huggingface_hub import HfApi, login
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset, Dataset

# Try to import torch for memory cleanup
try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

# ============ CONFIGURATION ============
HF_USERNAME = "himu1780"
MODEL_REPO = f"{HF_USERNAME}/ai-python-model"
DATASET_NAME = "jtatman/python-code-dataset-500k"
BASE_MODEL = "gpt2"

# Training hyperparameters (Memory optimized)
BATCH_SIZE = 1
GRADIENT_ACCUMULATION = 8
SAVE_STEPS = 500
LOGGING_STEPS = 50
MAX_LENGTH = 256
LEARNING_RATE = 5e-5
MAX_STEPS_PER_SESSION = 10000
EXAMPLES_PER_SESSION = 50000

# Continuous training settings
CONTINUOUS_TRAINING = True  # Set False to stop after one session
WAIT_BETWEEN_SESSIONS = 60  # Seconds to wait before next session

# ============ GLOBAL STATE ============
training_status = {
    "is_training": False,
    "current_step": 0,
    "total_loss": 0,
    "last_save": "Never",
    "start_time": None,
    "message": "Initializing...",
    "session_count": 0,
}

stop_requested = False

# ============ MEMORY CLEANUP ============
def cleanup_memory():
    """Free up memory after training"""
    gc.collect()
    if TORCH_AVAILABLE and torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("[INFO] Memory cleaned up")

# ============ AUTHENTICATION ============
def authenticate():
    """Login to Hugging Face Hub"""
    token = os.environ.get("HF_TOKEN")
    if token:
        login(token=token)
        training_status["message"] = "✅ Authenticated with Hugging Face"
        return True
    else:
        training_status["message"] = "❌ HF_TOKEN not found in secrets!"
        return False

# ============ MODEL LOADING ============
def load_model_and_tokenizer():
    """Load model from Hub (resume) or start fresh from base model"""
    global training_status
    
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    tokenizer.pad_token = tokenizer.eos_token
    
    try:
        training_status["message"] = f"🔄 Attempting to resume from {MODEL_REPO}..."
        model = AutoModelForCausalLM.from_pretrained(MODEL_REPO)
        training_status["message"] = f"✅ Resumed from {MODEL_REPO}"
        print(f"[INFO] Resumed training from {MODEL_REPO}")
    except Exception as e:
        training_status["message"] = f"🆕 Starting fresh from {BASE_MODEL}"
        model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
        print(f"[INFO] Starting fresh from {BASE_MODEL}: {e}")
    
    return model, tokenizer

# ============ DATASET PROCESSING ============
def prepare_dataset(tokenizer):
    """Load and prepare dataset"""
    global training_status
    training_status["message"] = "📥 Loading dataset (streaming mode)..."
    
    try:
        dataset = load_dataset(DATASET_NAME, split="train", streaming=True)
        dataset = dataset.take(EXAMPLES_PER_SESSION)
        
        def tokenize_function(examples):
            texts = []
            instructions = examples.get("instruction", [])
            outputs = examples.get("output", [])
            
            for instruction, output in zip(instructions, outputs):
                if instruction and output:
                    text = f"### Instruction:\n{instruction}\n\n### Response:\n{output}"
                    texts.append(text)
            
            if not texts:
                texts = [""]
            
            result = tokenizer(
                texts,
                truncation=True,
                max_length=MAX_LENGTH,
                padding="max_length",
                return_tensors=None,
            )
            result["labels"] = result["input_ids"].copy()
            return result
        
        tokenized_dataset = dataset.map(
            tokenize_function,
            batched=True,
            batch_size=100,
            remove_columns=["instruction", "output"],
        )
        
        training_status["message"] = "🔄 Converting dataset for Trainer..."
        
        all_examples = []
        for i, example in enumerate(tokenized_dataset):
            all_examples.append(example)
            # Progress every 5000 (IMPROVED)
            if i % 5000 == 0:
                training_status["message"] = f"📥 Loaded {i:,}/{EXAMPLES_PER_SESSION:,} examples..."
            if i >= EXAMPLES_PER_SESSION - 1:
                break
        
        train_dataset = Dataset.from_list(all_examples)
        
        training_status["message"] = f"✅ Dataset ready: {len(train_dataset):,} examples"
        return train_dataset
        
    except Exception as e:
        training_status["message"] = f"❌ Dataset error: {str(e)}"
        print(f"[ERROR] Dataset preparation failed: {e}")
        raise e

# ============ CUSTOM TRAINER ============
class StatusTrainer(Trainer):
    """Custom trainer with status updates and stop support"""
    
    def training_step(self, model, inputs):
        global stop_requested
        if stop_requested:
            raise KeyboardInterrupt("Stop requested by user")
        return super().training_step(model, inputs)
    
    def log(self, logs):
        super().log(logs)
        if "loss" in logs:
            training_status["total_loss"] = logs["loss"]
        training_status["current_step"] = self.state.global_step
    
    def save_model(self, output_dir=None, _internal_call=False):
        super().save_model(output_dir, _internal_call)
        training_status["last_save"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# ============ SINGLE TRAINING SESSION ============
def run_training_session():
    """Run a single training session"""
    global training_status, stop_requested
    
    model = None
    trainer = None
    
    try:
        if not authenticate():
            return False
        
        model, tokenizer = load_model_and_tokenizer()
        train_dataset = prepare_dataset(tokenizer)
        
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=tokenizer,
            mlm=False,
        )
        
        training_args = TrainingArguments(
            output_dir="./temp_checkpoints",
            overwrite_output_dir=True,
            per_device_train_batch_size=BATCH_SIZE,
            gradient_accumulation_steps=GRADIENT_ACCUMULATION,
            learning_rate=LEARNING_RATE,
            warmup_steps=100,
            weight_decay=0.01,
            logging_steps=LOGGING_STEPS,
            save_steps=SAVE_STEPS,
            save_total_limit=1,
            push_to_hub=True,
            hub_model_id=MODEL_REPO,
            hub_strategy="every_save",
            report_to="none",
            max_steps=MAX_STEPS_PER_SESSION,
            fp16=False,
            dataloader_num_workers=0,
            remove_unused_columns=False,
        )
        
        trainer = StatusTrainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            data_collator=data_collator,
            tokenizer=tokenizer,
        )
        
        training_status["message"] = "🏃 Training in progress..."
        trainer.train()
        trainer.push_to_hub()
        
        training_status["session_count"] += 1
        training_status["message"] = f"✅ Session {training_status['session_count']} completed!"
        return True
        
    except KeyboardInterrupt:
        training_status["message"] = "⏹️ Training stopped by user"
        return False
    except Exception as e:
        training_status["message"] = f"❌ Error: {str(e)}"
        print(f"[ERROR] Training failed: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        # MEMORY CLEANUP (IMPROVED)
        del model, trainer
        cleanup_memory()

# ============ MAIN TRAINING LOOP ============
def start_training():
    """Main training function with continuous loop"""
    global training_status, stop_requested
    
    if training_status["is_training"]:
        return "Training already in progress!"
    
    training_status["is_training"] = True
    training_status["start_time"] = datetime.now()
    stop_requested = False
    
    # CONTINUOUS TRAINING LOOP (IMPROVED)
    while not stop_requested:
        training_status["message"] = f"🚀 Starting session {training_status['session_count'] + 1}..."
        
        success = run_training_session()
        
        if stop_requested:
            break
        
        if not CONTINUOUS_TRAINING:
            break
        
        if success:
            training_status["message"] = f"⏳ Waiting {WAIT_BETWEEN_SESSIONS}s before next session..."
            time.sleep(WAIT_BETWEEN_SESSIONS)
        else:
            training_status["message"] = "⚠️ Session failed, retrying in 60s..."
            time.sleep(60)
    
    training_status["is_training"] = False
    stop_requested = False
    training_status["message"] = f"✅ Training finished! Total sessions: {training_status['session_count']}"
    return training_status["message"]

# ============ GRADIO INTERFACE ============
def get_status():
    """Get current training status"""
    elapsed = ""
    if training_status["start_time"]:
        delta = datetime.now() - training_status["start_time"]
        hours, remainder = divmod(int(delta.total_seconds()), 3600)
        minutes, seconds = divmod(remainder, 60)
        elapsed = f"{hours}h {minutes}m {seconds}s"
    
    return f"""
## 🤖 AI Python Model Trainer

### Status
| Item | Value |
|------|-------|
| **State** | {"🟢 Training" if training_status["is_training"] else "🔴 Stopped"} |
| **Message** | {training_status["message"]} |
| **Sessions Completed** | {training_status["session_count"]} |

### Progress
| Metric | Value |
|--------|-------|
| **Current Step** | {training_status["current_step"]:,} / {MAX_STEPS_PER_SESSION:,} |
| **Current Loss** | {training_status["total_loss"]:.4f if training_status["total_loss"] else "N/A"} |
| **Last Checkpoint** | {training_status["last_save"]} |
| **Elapsed Time** | {elapsed if elapsed else "N/A"} |

### Configuration
| Setting | Value |
|---------|-------|
| **Model Repo** | [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO}) |
| **Continuous Mode** | {"✅ Enabled" if CONTINUOUS_TRAINING else "❌ Disabled"} |
| **Batch Size** | {BATCH_SIZE} (effective: {BATCH_SIZE * GRADIENT_ACCUMULATION}) |
| **Max Steps/Session** | {MAX_STEPS_PER_SESSION:,} |
"""

def start_training_async():
    """Start training in background"""
    if training_status["is_training"]:
        return "⚠️ Training already in progress!"
    thread = threading.Thread(target=start_training, daemon=True)
    thread.start()
    return "🚀 Training started in background!"

def stop_training():
    """Stop training"""
    global stop_requested
    if not training_status["is_training"]:
        return "⚠️ No training in progress"
    stop_requested = True
    training_status["message"] = "⏹️ Stopping after current step..."
    return "⏹️ Stop requested"

# ============ AUTO-START ============
def auto_start():
    """Auto-start continuous training on Space launch"""
    time.sleep(10)
    while True:
        if not training_status["is_training"] and not stop_requested:
            print("[INFO] Auto-starting training session...")
            start_training()
        time.sleep(WAIT_BETWEEN_SESSIONS)

auto_thread = threading.Thread(target=auto_start, daemon=True)
auto_thread.start()

# ============ GRADIO APP ============
with gr.Blocks(title="AI Python Trainer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🐍 AI Python Code Model Trainer")
    gr.Markdown(f"**Continuous training** on `{DATASET_NAME}` with auto-checkpoint")
    
    status_display = gr.Markdown(get_status)
    
    with gr.Row():
        start_btn = gr.Button("▶️ Start Training", variant="primary")
        stop_btn = gr.Button("⏹️ Stop Training", variant="stop")
        refresh_btn = gr.Button("🔄 Refresh Status")
    
    output = gr.Textbox(label="Output", interactive=False)
    
    start_btn.click(start_training_async, outputs=output)
    stop_btn.click(stop_training, outputs=output)
    refresh_btn.click(get_status, outputs=status_display)
    
    demo.load(get_status, outputs=status_display, every=30)

demo.launch()