"""
DocuMint Train - LoRA Training Pipeline
Base Model: Qwen2-0.5B-Instruct
"""

import os
import gc
import torch
from typing import Optional, Dict, Any
from datasets import load_dataset, Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
from huggingface_hub import login, HfApi


# ============ CONFIG ============

BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
OUTPUT_REPO = "himu1780/DocuMint-Models"
DATA_REPO = "himu1780/DocuMint-Data"
OUTPUT_DIR = "./lora_output"

# LoRA Configuration
LORA_R = 8
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]

# Training Configuration
MAX_LENGTH = 512
BATCH_SIZE = 1
GRADIENT_ACCUMULATION = 4
LEARNING_RATE = 2e-4
NUM_EPOCHS = 3
WARMUP_STEPS = 100
SAVE_STEPS = 500
LOGGING_STEPS = 50


# ============ GLOBAL STATE ============

training_status = {
    "is_training": False,
    "current_step": 0,
    "total_steps": 0,
    "loss": 0.0,
    "message": "Ready",
    "progress": 0
}


# ============ UTILS ============

def cleanup_memory():
    """Free memory."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def authenticate() -> bool:
    """Login to HuggingFace."""
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        login(token=hf_token)
        print("✅ Authenticated with HuggingFace")
        return True
    print("❌ No HF_TOKEN found!")
    return False


# ============ DATASET ============

def format_instruction(example: Dict) -> Dict:
    """Format dataset examples for instruction tuning."""
    # Adjust based on your dataset structure
    if "instruction" in example and "output" in example:
        # Alpaca format
        text = f"<|im_start|>user\n{example['instruction']}<|im_end|>\n<|im_start|>assistant\n{example['output']}<|im_end|>"
    elif "text" in example:
        # Plain text
        text = example["text"]
    elif "question" in example and "answer" in example:
        # Q&A format
        text = f"<|im_start|>user\n{example['question']}<|im_end|>\n<|im_start|>assistant\n{example['answer']}<|im_end|>"
    else:
        # Fallback - use all values
        text = str(example)
    
    return {"text": text}


def prepare_dataset(tokenizer, dataset_name: str = None, split: str = "train"):
    """Load and prepare dataset for training."""
    global training_status
    training_status["message"] = "Loading dataset..."
    
    try:
        if dataset_name:
            # Load specified dataset
            dataset = load_dataset(dataset_name, split=split)
        else:
            # Load from our private repo
            dataset = load_dataset(DATA_REPO, split=split)
        
        print(f"📊 Loaded {len(dataset)} examples")
        
        # Format for instruction tuning
        dataset = dataset.map(format_instruction, remove_columns=dataset.column_names)
        
        # Tokenize
        def tokenize(example):
            tokens = tokenizer(
                example["text"],
                truncation=True,
                max_length=MAX_LENGTH,
                padding="max_length"
            )
            tokens["labels"] = tokens["input_ids"].copy()
            return tokens
        
        dataset = dataset.map(tokenize, remove_columns=["text"])
        training_status["message"] = f"Dataset ready: {len(dataset)} examples"
        
        return dataset
        
    except Exception as e:
        training_status["message"] = f"Dataset error: {e}"
        print(f"❌ Failed to load dataset: {e}")
        return None


# ============ MODEL ============

def load_base_model():
    """Load Qwen2-0.5B base model."""
    global training_status
    training_status["message"] = "Loading base model..."
    
    print(f"🔄 Loading {BASE_MODEL}...")
    
    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL,
        trust_remote_code=True
    )
    
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float32,  # CPU
        device_map="cpu",
        trust_remote_code=True,
        low_cpu_mem_usage=True
    )
    
    print("✅ Base model loaded!")
    return model, tokenizer


def apply_lora(model):
    """Apply LoRA configuration to model."""
    global training_status
    training_status["message"] = "Applying LoRA..."
    
    lora_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        target_modules=TARGET_MODULES,
        task_type=TaskType.CAUSAL_LM,
        bias="none"
    )
    
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    
    print("✅ LoRA applied!")
    return model


# ============ TRAINING ============

class StatusCallback:
    """Callback to update training status."""
    
    def __init__(self, total_steps):
        self.total_steps = total_steps
    
    def on_step_end(self, args, state, control, **kwargs):
        global training_status
        training_status["current_step"] = state.global_step
        training_status["total_steps"] = self.total_steps
        training_status["progress"] = (state.global_step / self.total_steps) * 100
        if state.log_history:
            training_status["loss"] = state.log_history[-1].get("loss", 0)


def train(
    dataset_name: str = None,
    epochs: int = NUM_EPOCHS,
    batch_size: int = BATCH_SIZE,
    learning_rate: float = LEARNING_RATE
):
    """
    Main training function.
    
    Args:
        dataset_name: HuggingFace dataset to use (or None for DocuMint-Data)
        epochs: Number of training epochs
        batch_size: Training batch size
        learning_rate: Learning rate
    
    Returns:
        Success message or error
    """
    global training_status
    training_status["is_training"] = True
    training_status["message"] = "Starting training..."
    
    try:
        # Authenticate
        if not authenticate():
            return "❌ Authentication failed. Set HF_TOKEN environment variable."
        
        # Load model
        model, tokenizer = load_base_model()
        
        # Apply LoRA
        model = apply_lora(model)
        
        # Prepare dataset
        dataset = prepare_dataset(tokenizer, dataset_name)
        if dataset is None:
            return "❌ Failed to load dataset"
        
        # Calculate steps
        total_steps = (len(dataset) // (batch_size * GRADIENT_ACCUMULATION)) * epochs
        training_status["total_steps"] = total_steps
        
        # Training arguments
        training_args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            num_train_epochs=epochs,
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=GRADIENT_ACCUMULATION,
            learning_rate=learning_rate,
            warmup_steps=WARMUP_STEPS,
            logging_steps=LOGGING_STEPS,
            save_steps=SAVE_STEPS,
            save_total_limit=2,
            fp16=False,  # CPU
            bf16=False,
            optim="adamw_torch",
            lr_scheduler_type="cosine",
            report_to="none",
            remove_unused_columns=False
        )
        
        # Data collator
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=tokenizer,
            mlm=False
        )
        
        # Trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=dataset,
            data_collator=data_collator
        )
        
        training_status["message"] = "Training in progress..."
        
        # Train!
        trainer.train()
        
        training_status["message"] = "Saving model..."
        
        # Save locally
        model.save_pretrained(OUTPUT_DIR)
        tokenizer.save_pretrained(OUTPUT_DIR)
        
        # Push to Hub
        training_status["message"] = "Pushing to HuggingFace..."
        model.push_to_hub(OUTPUT_REPO)
        tokenizer.push_to_hub(OUTPUT_REPO)
        
        training_status["is_training"] = False
        training_status["message"] = "✅ Training complete! Model saved to " + OUTPUT_REPO
        training_status["progress"] = 100
        
        cleanup_memory()
        return f"✅ Training complete! LoRA adapters saved to {OUTPUT_REPO}"
        
    except Exception as e:
        training_status["is_training"] = False
        training_status["message"] = f"❌ Error: {str(e)}"
        return f"❌ Training failed: {str(e)}"


def get_status() -> Dict[str, Any]:
    """Get current training status."""
    return training_status.copy()


# ============ MAIN ============

if __name__ == "__main__":
    print("🏆 DocuMint Train - LoRA Training Pipeline")
    print("Run train() to start training.")