#!/usr/bin/env python3
"""
Enhanced Fine-tuning script for CodeLlama with optimized hyperparameters
Supports:
- Resume from checkpoint (automatic detection)
- Incremental fine-tuning (continue from existing adapter)
- Fresh training option
"""

import os
import sys
import torch
import json
from pathlib import Path
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    BitsAndBytesConfig,
    Trainer,
    DataCollatorForLanguageModeling,
    EarlyStoppingCallback,
)
from peft import (
    LoraConfig,
    PeftModel,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType,
)

def get_device_info():
    """Detect and return available compute device"""
    device_info = {
        "device": "cpu",
        "device_type": "cpu",
        "use_quantization": False,
        "dtype": torch.float32
    }
    
    if torch.cuda.is_available():
        device_info["device"] = "cuda"
        device_info["device_type"] = "cuda"
        device_info["use_quantization"] = True
        device_info["dtype"] = torch.float16
        device_info["device_count"] = torch.cuda.device_count()
        device_info["device_name"] = torch.cuda.get_device_name(0)
        print(f"✓ CUDA GPU detected: {device_info['device_name']} (Count: {device_info['device_count']})")
    else:
        print("⚠ No GPU detected, using CPU (training will be very slow)")
    
    return device_info

def get_bitsandbytes_config():
    """Get BitsAndBytes config if CUDA is available"""
    if torch.cuda.is_available():
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
    return None

def load_and_prepare_model(
    model_name: str,
    adapter_path: str | None = None,
    lora_r: int = 48,
    lora_alpha: int = 96,
    lora_dropout: float = 0.15
):
    """Load CodeLlama model with optimized LoRA configuration"""
    device_info = get_device_info()
    print(f"\nLoading model: {model_name}")
    
    # Tokenizer
    tokenizer_source = adapter_path if adapter_path and os.path.isdir(adapter_path) else model_name
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_source)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    
    # Quantization config
    bnb_config = get_bitsandbytes_config()
    
    # Model loading kwargs
    model_kwargs = {
        "trust_remote_code": True,
    }
    
    if bnb_config is not None:
        print("Using 4-bit quantization (CUDA)")
        model_kwargs["quantization_config"] = bnb_config
        model_kwargs["device_map"] = "auto"
    else:
        model_kwargs["torch_dtype"] = device_info["dtype"]
        model_kwargs["device_map"] = "auto"
    
    # Load base model
    base_model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
    
    # Prepare for k-bit training
    if bnb_config is not None:
        base_model = prepare_model_for_kbit_training(base_model)
    
    # LoRA configuration (optimized for CodeLlama)
    lora_config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=lora_dropout,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )
    
    # Load or create LoRA adapter
    if adapter_path and os.path.isdir(adapter_path):
        print(f"📂 Loading existing LoRA adapter from: {adapter_path}")
        print("   (Incremental fine-tuning mode - continuing from existing model)")
        model = PeftModel.from_pretrained(base_model, adapter_path, is_trainable=True)
    else:
        print("🆕 Creating new LoRA adapter (Fresh training mode)")
        model = get_peft_model(base_model, lora_config)
    
    # Enable gradient checkpointing
    model.gradient_checkpointing_enable()
    
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    trainable_ratio = (trainable_params / total_params) * 100
    
    print(f"\nModel loaded successfully!")
    print(f"  - Device: {device_info['device']}")
    print(f"  - Trainable parameters: {trainable_params:,}")
    print(f"  - Total parameters: {total_params:,}")
    print(f"  - Trainable ratio: {trainable_ratio:.2f}%")
    
    return model, tokenizer, device_info

def tokenize_function(examples, tokenizer, max_length=1536):
    """Tokenize function for dataset"""
    # Ensure pad_token is set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    
    # Combine instruction and response
    # For CodeLlama chat format: instruction already ends with [/INST]
    # So we just append: instruction + response + EOS
    texts = []
    for instruction, response in zip(examples["instruction"], examples["response"]):
        # Instruction already contains: <s>[INST]...[/INST]
        # We append response + EOS
        text = f"{instruction}{response}{tokenizer.eos_token}"
        texts.append(text)
    
    # Tokenize with padding to max_length for consistent batch sizes
    tokenized = tokenizer(
        texts,
        truncation=True,
        max_length=max_length,
        padding="max_length",
        return_tensors=None,  # Return lists, not tensors
    )
    
    # Labels are same as input_ids for causal LM
    labels = []
    pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
    
    # Set labels, masking padding tokens with -100 (ignored in loss)
    for input_ids_seq in tokenized["input_ids"]:
        label_seq = input_ids_seq.copy()
        # Mask padding tokens
        label_seq = [-100 if token_id == pad_token_id else token_id for token_id in label_seq]
        labels.append(label_seq)
    
    tokenized["labels"] = labels
    
    return tokenized

def find_checkpoint(output_dir):
    """Find the latest checkpoint in output directory"""
    checkpoint_dir = Path(output_dir)
    if not checkpoint_dir.exists():
        return None
    
    # Look for checkpoint directories
    checkpoints = []
    for item in checkpoint_dir.iterdir():
        if item.is_dir() and item.name.startswith("checkpoint-"):
            try:
                step_num = int(item.name.split("-")[1])
                trainer_state = item / "trainer_state.json"
                if trainer_state.exists():
                    checkpoints.append((step_num, str(item)))
            except (ValueError, IndexError):
                continue
    
    if checkpoints:
        # Sort by step number and return latest
        checkpoints.sort(key=lambda x: x[0], reverse=True)
        return checkpoints[0][1]
    
    return None

def load_training_data(file_path):
    """Load training data from JSONL file"""
    print(f"Loading training data from {file_path}")
    
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Training data file not found: {file_path}")
    
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    data.append(json.loads(line))
                except json.JSONDecodeError as e:
                    print(f"⚠️  Warning: Skipping invalid JSON line: {e}")
                    continue
    
    return data

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="Fine-tune CodeLlama with optimized hyperparameters")
    parser.add_argument("--base-model", required=True, help="Base model path or HuggingFace ID")
    parser.add_argument("--adapter-path", default=None, help="Path to existing LoRA adapter (for incremental fine-tuning)")
    parser.add_argument("--dataset", required=True, help="Path to training dataset JSONL")
    parser.add_argument("--output-dir", required=True, help="Output directory for fine-tuned model")
    parser.add_argument("--resume-from-checkpoint", default=None, help="Resume from specific checkpoint (or 'auto' to find latest)")
    parser.add_argument("--fresh", action="store_true", help="Force fresh training (ignore existing checkpoints)")
    
    # Hyperparameters (optimized for CodeLlama based on HYPERPARAMETER_ANALYSIS.md)
    parser.add_argument("--max-length", type=int, default=1536, help="Max sequence length (default: 1536)")
    parser.add_argument("--num-epochs", type=int, default=5, help="Number of epochs (default: 5)")
    parser.add_argument("--batch-size", type=int, default=2, help="Batch size per device (default: 2)")
    parser.add_argument("--gradient-accumulation", type=int, default=4, help="Gradient accumulation steps (default: 4)")
    parser.add_argument("--learning-rate", type=float, default=2e-5, help="Learning rate (default: 2e-5)")
    parser.add_argument("--lora-r", type=int, default=48, help="LoRA rank (default: 48)")
    parser.add_argument("--lora-alpha", type=int, default=96, help="LoRA alpha (default: 96)")
    parser.add_argument("--lora-dropout", type=float, default=0.15, help="LoRA dropout (default: 0.15)")
    parser.add_argument("--warmup-ratio", type=float, default=0.1, help="Warmup ratio (default: 0.1)")
    parser.add_argument("--eval-steps", type=int, default=25, help="Evaluation steps (default: 25)")
    parser.add_argument("--save-steps", type=int, default=25, help="Save steps (default: 25)")
    parser.add_argument("--early-stopping-patience", type=int, default=5, help="Early stopping patience (default: 5)")
    parser.add_argument("--logging-steps", type=int, default=5, help="Logging steps (default: 5)")
    
    args = parser.parse_args()
    
    print("=" * 70)
    print("🚀 CodeLlama Fine-Tuning with Optimized Hyperparameters")
    print("=" * 70)
    print(f"Base model: {args.base_model}")
    print(f"Dataset: {args.dataset}")
    print(f"Output dir: {args.output_dir}")
    if args.adapter_path:
        print(f"Adapter path: {args.adapter_path} (Incremental fine-tuning)")
    print("=" * 70)
    
    # Check for existing checkpoint
    resume_checkpoint = None
    if not args.fresh:
        if args.resume_from_checkpoint == "auto":
            resume_checkpoint = find_checkpoint(args.output_dir)
            if resume_checkpoint:
                print(f"\n✅ Found existing checkpoint: {resume_checkpoint}")
                print("   Training will resume from this checkpoint")
        elif args.resume_from_checkpoint:
            resume_checkpoint = args.resume_from_checkpoint
            if os.path.isdir(resume_checkpoint):
                print(f"\n📂 Resuming from specified checkpoint: {resume_checkpoint}")
            else:
                print(f"\n⚠️  Warning: Checkpoint path does not exist: {resume_checkpoint}")
                resume_checkpoint = None
    else:
        print("\n🆕 Fresh training mode - starting from scratch")
        # Clear any existing checkpoints if fresh mode
        if os.path.exists(args.output_dir):
            checkpoint_dir = Path(args.output_dir)
            for item in checkpoint_dir.iterdir():
                if item.is_dir() and item.name.startswith("checkpoint-"):
                    print(f"   Removing old checkpoint: {item.name}")
                    import shutil
                    shutil.rmtree(item)
    
    # Load model and tokenizer
    model, tokenizer, device_info = load_and_prepare_model(
        args.base_model,
        args.adapter_path,
        lora_r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout
    )
    
    # Check if using pre-split dataset (train.jsonl in split directory)
    dataset_path = Path(args.dataset)
    val_dataset_path = None
    use_presplit = False
    
    if dataset_path.name == "train.jsonl":
        # Check if val.jsonl exists in same directory
        val_path = dataset_path.parent / "val.jsonl"
        if val_path.exists():
            val_dataset_path = val_path
            use_presplit = True
            print(f"\n✅ Using pre-split dataset:")
            print(f"   Train: {dataset_path}")
            print(f"   Val: {val_dataset_path}")
    
    # Load training data
    training_data = load_training_data(args.dataset)
    
    # Convert to dataset format
    instructions = []
    responses = []
    
    for item in training_data:
        if "instruction" in item and "response" in item:
            instructions.append(item["instruction"])
            responses.append(item["response"])
        else:
            print(f"⚠️  Warning: Skipping invalid sample (missing instruction/response)")
    
    if not instructions:
        raise ValueError("No valid training samples found in dataset")
    
    print(f"\n✅ Loaded {len(instructions)} training samples")
    
    # Create training dataset
    train_dataset_dict = Dataset.from_dict({
        "instruction": instructions,
        "response": responses
    })
    
    # Tokenize training dataset
    print("Tokenizing training dataset...")
    tokenized_train = train_dataset_dict.map(
        lambda x: tokenize_function(x, tokenizer, max_length=args.max_length),
        batched=True,
        remove_columns=train_dataset_dict.column_names
    )
    
    # Load validation dataset if pre-split, otherwise split from training data
    if use_presplit and val_dataset_path:
        print(f"\n✅ Loading validation dataset from: {val_dataset_path}")
        val_data = load_training_data(str(val_dataset_path))
        val_instructions = []
        val_responses = []
        
        for item in val_data:
            if "instruction" in item and "response" in item:
                val_instructions.append(item["instruction"])
                val_responses.append(item["response"])
        
        val_dataset_dict = Dataset.from_dict({
            "instruction": val_instructions,
            "response": val_responses
        })
        
        print("Tokenizing validation dataset...")
        tokenized_val = val_dataset_dict.map(
            lambda x: tokenize_function(x, tokenizer, max_length=args.max_length),
            batched=True,
            remove_columns=val_dataset_dict.column_names
        )
        
        train_dataset = tokenized_train
        val_dataset = tokenized_val
        
        print(f"  - Training samples: {len(train_dataset)}")
        print(f"  - Validation samples: {len(val_dataset)}")
    else:
        # Split into train/validation (80/20)
        print("\nSplitting dataset into train/validation (80/20)...")
        train_val_split = tokenized_train.train_test_split(test_size=0.2, seed=42)
        train_dataset = train_val_split["train"]
        val_dataset = train_val_split["test"]
        
        print(f"  - Training samples: {len(train_dataset)}")
        print(f"  - Validation samples: {len(val_dataset)}")
    
    print(f"  - Training samples: {len(train_dataset)}")
    print(f"  - Validation samples: {len(val_dataset)}")
    
    # Calculate training steps
    use_fp16 = device_info["device_type"] == "cuda"
    effective_batch_size = args.batch_size * args.gradient_accumulation
    steps_per_epoch = max(1, len(train_dataset) // effective_batch_size)
    total_steps = steps_per_epoch * args.num_epochs
    warmup_steps = max(int(total_steps * args.warmup_ratio), 10)
    
    print(f"\n📊 Training Configuration:")
    print(f"  - Total training steps: {total_steps}")
    print(f"  - Steps per epoch: {steps_per_epoch}")
    print(f"  - Warmup steps: {warmup_steps} ({100*warmup_steps/total_steps:.1f}% of training)")
    
    # Training arguments (optimized for CodeLlama)
    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.num_epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation,
        warmup_steps=warmup_steps,
        learning_rate=args.learning_rate,
        weight_decay=0.01,
        fp16=use_fp16,
        logging_steps=args.logging_steps,
        save_steps=args.save_steps,
        eval_strategy="steps",
        eval_steps=args.eval_steps,
        save_total_limit=3,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        lr_scheduler_type="cosine",
        max_grad_norm=1.0,
        report_to="none",
        push_to_hub=False,
        dataloader_pin_memory=(device_info["device_type"] == "cuda"),
        remove_unused_columns=False,
        resume_from_checkpoint=resume_checkpoint,  # Resume support
    )
    
    print(f"\n⚙️  Hyperparameters (Optimized for CodeLlama):")
    print(f"  - Max length: {args.max_length}")
    print(f"  - Epochs: {args.num_epochs}")
    print(f"  - Batch size: {args.batch_size}")
    print(f"  - Gradient accumulation: {args.gradient_accumulation}")
    print(f"  - Learning rate: {args.learning_rate}")
    print(f"  - LoRA rank: {args.lora_r}")
    print(f"  - LoRA alpha: {args.lora_alpha}")
    print(f"  - LoRA dropout: {args.lora_dropout}")
    print(f"  - Device: {device_info['device']}")
    print(f"  - Mixed precision (fp16): {use_fp16}")
    print("=" * 70)
    
    # Data collator - since we pad during tokenization, collator mainly handles batching
    # Ensure pad_token_id is set
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
        tokenizer.pad_token = tokenizer.eos_token
    
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,  # Causal LM, not masked LM
    )
    
    # Create trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=args.early_stopping_patience)],
    )
    
    # Train
    print("\n🚀 Starting training...")
    if resume_checkpoint:
        print(f"   Resuming from: {resume_checkpoint}")
    print("=" * 70)
    
    trainer.train(resume_from_checkpoint=resume_checkpoint)
    
    # Save final model
    print(f"\n💾 Saving fine-tuned model to {args.output_dir}")
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    model.save_pretrained(args.output_dir)
    
    # Save training config
    config = {
        "base_model": args.base_model,
        "adapter_path": args.adapter_path if args.adapter_path else None,
        "dataset": args.dataset,
        "output_dir": args.output_dir,
        "hyperparameters": {
            "max_length": args.max_length,
            "num_epochs": args.num_epochs,
            "batch_size": args.batch_size,
            "gradient_accumulation": args.gradient_accumulation,
            "learning_rate": args.learning_rate,
            "lora_r": args.lora_r,
            "lora_alpha": args.lora_alpha,
            "lora_dropout": args.lora_dropout,
        },
        "training_mode": "incremental" if args.adapter_path else "fresh",
        "resumed_from_checkpoint": resume_checkpoint is not None
    }
    
    config_path = Path(args.output_dir) / "training_config.json"
    with open(config_path, 'w') as f:
        json.dump(config, f, indent=2)
    
    print("\n✅ Fine-tuning complete!")
    print(f"Model saved to: {args.output_dir}")
    print(f"Config saved to: {config_path}")
    print(f"\n💡 To continue training with new data (incremental fine-tuning):")
    print(f"   python finetune_codellama.py --base-model {args.base_model} \\")
    print(f"       --adapter-path {args.output_dir} \\")
    print(f"       --dataset <new_dataset.jsonl> \\")
    print(f"       --output-dir <new_output_dir>")
    print(f"\n💡 To resume from checkpoint if training is interrupted:")
    print(f"   python finetune_codellama.py --base-model {args.base_model} \\")
    print(f"       --dataset {args.dataset} \\")
    print(f"       --output-dir {args.output_dir} \\")
    print(f"       --resume-from-checkpoint auto")

if __name__ == "__main__":
    main()