import os
import torch
from dataclasses import dataclass, field
from typing import Optional
from datasets import load_dataset, Features, Value, Sequence
from peft import LoraConfig, get_peft_model
from transformers.trainer_utils import get_last_checkpoint, set_seed
from src.datasets.images.utils import read_image
from src.logger import get_logger
from src.modeling.llava.trainer import LlavaNextTrainer, GpuStatsCallback
from src.modeling.llava.utils import get_masking_config, load_model, log_trainable_parameters
from transformers import TrainingArguments
from transformers.hf_argparser import HfArgumentParser
from peft import PeftModel
from transformers.modeling_utils import unwrap_model


logger = get_logger(__name__)


@dataclass
class LlavaTrainingArguments:
    # Basic arguments
    model_name_or_path: str = field(metadata={"help": "Path to the model directory"})
    data_path: str = field(metadata={"help": "Path to the dataset file"})
    images_path: str = field(metadata={"help": "Path to the directory containing images"})
    # Output arguments
    trainer_output_dir: str = field(default="models/checkpoints", metadata={"help": "Directory to save the trainer output"})
    model_output_dir: str = field(default="models/finetuned", metadata={"help": "Directory to save the final model"})
    # Run arguments
    run_name: str = field(default="llava_training", metadata={"help": "Name of the training run for logging"})
    # Training parts
    train_vision_tower: bool = field(default=False, metadata={"help": "Whether to train the vision tower."})
    train_multimodal_projector: bool = field(default=False, metadata={"help": "Whether to train the projector."})
    train_language_model: bool = field(default=True, metadata={"help": "Whether to train the language model."})
    train_lm_head: bool = field(default=False, metadata={"help": "Whether to train the language model head."})
    # Training arguments
    per_device_train_batch_size: int = field(default=1, metadata={"help": "Batch size per device for training"})
    per_device_eval_batch_size: int = field(default=1, metadata={"help": "Batch size per device for evaluation"})
    gradient_accumulation_steps: int = field(default=1, metadata={"help": "Number of gradient accumulation steps"})
    max_seq_length: int = field(default=4096, metadata={"help": "Maximum sequence length for training"})
    epochs: int = field(default=1, metadata={"help": "Number of training epochs"})
    lr: float = field(default=1e-5, metadata={"help": "Learning rate for training"})
    language_model_lr: Optional[float] = field(default=None, metadata={"help": "Learning rate for the language model"})
    vision_tower_lr: Optional[float] = field(default=None, metadata={"help": "Learning rate for the vision tower"})
    multimodal_projector_lr: Optional[float] = field(default=None, metadata={"help": "Learning rate for the multimodal projector"})
    weight_decay: float = field(default=0.01, metadata={"help": "Weight decay for optimizer"})
    warmup_ratio: float = field(default=0.03, metadata={"help": "Warmup ratio for learning rate scheduler"})
    lr_scheduler_type: str = field(default="cosine", metadata={"help": "Type of learning rate scheduler"})
    optim: str = field(default="adamw_torch_fused", metadata={"help": "Optimizer type"})
    # Optimization arguments
    deepspeed: Optional[str] = field(default=None, metadata={"help": "Path to the deepspeed configuration file"})
    gradient_checkpointing: bool = field(default=False, metadata={"help": "Enable gradient checkpointing"})
    use_peft: bool = field(default=False, metadata={"help": "Use LoRA for training"})
    lora_r : Optional[int] = field(default=16, metadata={"help": "LoRA rank"})
    lora_alpha : Optional[int] = field(default=32, metadata={"help": "LoRA alpha"})
    lora_dropout : Optional[float] = field(default=0.05, metadata={"help": "LoRA dropout"})
    lora_target_modules: Optional[str] = field(default=None, metadata={"help": "Comma-separated list of target modules for LoRA."})
    lora_modules_to_save: Optional[str] = field(default=None, metadata={"help": "Comma-separated list of modules to save with LoRA."})
    lora_layers_to_transform: Optional[str] = field(default=None, metadata={"help": "Comma-separated list of layers to apply LoRA."})
    # Other arguments
    seed: int = field(default=42, metadata={"help": "Random seed for initialization"})
    logging_steps: int = field(default=10, metadata={"help": "Number of steps between logging"})
    report_to: str = field(default="none", metadata={"help": "Reporting tool for logging"})
    save_steps: int = field(default=500, metadata={"help": "Number of steps between model saves"})
    eval_strategy: str = field(default="no", metadata={"help": "Evaluation strategy to use during training"})
    eval_steps: int = field(default=100, metadata={"help": "Number of steps between evaluations"})
    test_size: float = field(default=0.01, metadata={"help": "Proportion of the dataset to include in the evaluation split"})


def load_data_files(args: LlavaTrainingArguments):
    if os.path.isdir(args.data_path):
        data_files = []
        for root, _, files in os.walk(args.data_path):
            for file_name in files:
                if file_name.endswith((".json", ".jsonl")):
                    data_files.append(os.path.join(root, file_name))
        if not data_files:
            raise ValueError(f"No JSON files found in the directory: {args.data_path}")
        logger.info(f"Found {len(data_files)} JSON files in the directory: {args.data_path}")
        return {"train": data_files}
    elif os.path.isfile(args.data_path):
        return {"train": args.data_path}
    else:
        raise ValueError(f"Invalid data path: {args.data_path}")


def train_eval_split(dataset, args: LlavaTrainingArguments):
    if args.eval_strategy == "no":
        return dataset["train"], None

    split = dataset["train"].train_test_split(test_size=args.test_size, seed=args.seed)
    return split["train"], split["test"]


def prepare_for_training(model, args: LlavaTrainingArguments):
    # Freeze not trained components
    def unfreeze(module):
        for param in module.parameters():
            param.requires_grad = True

    # By default, we freeze all components
    for param in model.parameters():
        param.requires_grad = False

    if args.train_vision_tower:
        logger.info("Enabling training for the vision tower.")
        unfreeze(model.vision_tower)

    if args.train_multimodal_projector:
        logger.info("Enabling training for the multimodal projector.")
        unfreeze(model.multi_modal_projector)

    if args.train_language_model and not args.use_peft:
        # When using PEFT, we do not unfreeze the language model,
        # as it will be trained through LoRA layers.
        logger.info("Enabling training for the language model.")
        unfreeze(model.language_model)

    if args.train_lm_head:
        logger.info("Enabling training for the language model head.")
        unfreeze(model.lm_head)

    # Enable gradient checkpointing
    def enable_gradient_checkpointing(module):
        if hasattr(module, "enable_input_require_grads"):
            module.enable_input_require_grads()
        if hasattr(module, "gradient_checkpointing_enable"):
            module.gradient_checkpointing_enable(dict(use_reentrant=True))

    if args.gradient_checkpointing:
        model.config.use_cache = False
        enable_gradient_checkpointing(model.vision_tower)
        enable_gradient_checkpointing(model.multi_modal_projector)
        enable_gradient_checkpointing(model.language_model)
        enable_gradient_checkpointing(model.lm_head)


def get_peft_config(args: LlavaTrainingArguments) -> Optional[LoraConfig]:
    if not args.use_peft:
        return None

    return LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        bias="none",
        target_modules=args.lora_target_modules.split(",") if args.lora_target_modules else None,
        modules_to_save=args.lora_modules_to_save.split(",") if args.lora_modules_to_save else None,
        layers_to_transform=([int(x) for x in args.lora_layers_to_transform.split(",")] if args.lora_layers_to_transform else None),
    )


def collate_fn(processor, args: LlavaTrainingArguments, examples):
    assistant_start_ids, response_end_token_id = get_masking_config(processor.tokenizer)

    texts, images = [], []
    for example in examples:
        # Build conversation
        conversation = []
        for msg in example["conversations"]:
            role = "user" if msg["from"] == "human" else "assistant"
            conversation.append({"role": role, "content": msg["value"]})
        
        # Apply Chat Template
        text = processor.apply_chat_template(conversation, tokenize=False)
        texts.append(text)

        # Load Image
        image = read_image(os.path.join(args.images_path, example["image_source"]), example["image"], example.get("image_shard", None))
        images.append(image)

    batch = processor(
        images=images, 
        text=texts, 
        return_tensors="pt", 
        padding=True, 
        truncation=True, 
        max_length=args.max_seq_length,
        return_token_type_ids=False
    )

    input_ids = batch["input_ids"]
    labels = input_ids.clone()
    
    labels[:] = -100
    
    start_len = len(assistant_start_ids)
    for i in range(input_ids.shape[0]):
        seq = input_ids[i].tolist()
        seq_len = len(seq)
        
        j = 0
        while j < seq_len - start_len:
            if seq[j : j + start_len] == assistant_start_ids:
                response_start_idx = j + start_len

                if response_end_token_id in seq[response_start_idx:]:
                    response_end_idx = seq.index(response_end_token_id, response_start_idx) + 1     # +1 to include the end token
                else:
                    response_end_idx = seq_len

                labels[i, response_start_idx : response_end_idx] = input_ids[i, response_start_idx : response_end_idx]
                j = response_end_idx
            else:
                j += 1
                
    batch["labels"] = labels
    return batch


def save_model(model, tokenizer, processor, args: LlavaTrainingArguments):
    logger.info("Saving the final model...")

    output_dir = os.path.join(args.model_output_dir, args.run_name)
    os.makedirs(output_dir, exist_ok=True)

    raw_model = unwrap_model(model)

    if isinstance(raw_model, PeftModel):
        logger.info("Detected PeftModel, merging for final model.")
        raw_model = raw_model.merge_and_unload()                          

    raw_model.save_pretrained(output_dir, safe_serialization=True)

    tokenizer.save_pretrained(output_dir)
    processor.save_pretrained(output_dir)

    logger.info("Final model, tokenizer, and processor saved successfully.")


def train(args: LlavaTrainingArguments):
    set_seed(args.seed)
    logger.info(f"Training with model: {args.model_name_or_path}")

    trainer_output_dir = os.path.join(args.trainer_output_dir, args.run_name)
    training_args = TrainingArguments(
        output_dir=trainer_output_dir,
        overwrite_output_dir=True,
        save_strategy="steps",
        save_steps=args.save_steps,
        save_total_limit=3,
        num_train_epochs=args.epochs,
        learning_rate=args.lr,
        weight_decay=args.weight_decay,
        warmup_ratio=args.warmup_ratio,
        per_device_train_batch_size=args.per_device_train_batch_size,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        optim=args.optim,
        lr_scheduler_type=args.lr_scheduler_type,
        bf16=True,
        dataloader_num_workers=4,
        dataloader_pin_memory=True,
        dataloader_persistent_workers=True,
        remove_unused_columns=False,  # keep image columns -> pixel_values
        logging_steps=args.logging_steps,
        report_to=args.report_to,
        run_name=args.run_name,
        eval_strategy=args.eval_strategy,
        eval_steps=args.eval_steps,
        deepspeed=args.deepspeed,
    )

    with training_args.main_process_first(desc="loading and tokenizing dataset"):
        data_files = load_data_files(args)
        dataset = load_dataset("json", data_files=data_files)
        logger.info(f"Loaded dataset from {args.data_path}. Found {len(dataset['train'])} training examples.")

        train_ds, eval_ds = train_eval_split(dataset, args)
        logger.info(f"Training dataset size: {len(train_ds)}. Evaluation dataset size: {len(eval_ds) if eval_ds else 'N/A'}.")
        

    model, tokenizer, processor = load_model(args.model_name_or_path)
    logger.info(f"Loaded model and tokenizer from {args.model_name_or_path}")

    # Freeze/unfreeze selections first so base LM stays frozen for PEFT
    prepare_for_training(model, args)
    logger.info("Model prepared for training (freezing/unfreezing & checkpointing).")

    peft_config = get_peft_config(args)
    if peft_config is not None:
        logger.info("Wrapping language_model model with PEFT LoRA adapters.")
        model = get_peft_model(model, peft_config)
        logger.info(f"LoRA configuration: {peft_config}")

    trainer = LlavaNextTrainer(
        model=model,
        args=training_args,
        projector_lr=args.multimodal_projector_lr,
        vision_lr=args.vision_tower_lr,
        llm_lr=args.language_model_lr,
        tokenizer=tokenizer,
        data_collator=lambda x: collate_fn(processor, args, x),
        train_dataset=train_ds,
        eval_dataset=eval_ds,
    )
    logger.info("Initialized Transformers Trainer with the provided configuration.")

    trainer.add_callback(GpuStatsCallback())

    log_trainable_parameters(trainer.model)

    last_checkpoint = get_last_checkpoint(trainer_output_dir) if os.path.exists(trainer_output_dir) else None
    if last_checkpoint is not None:
        logger.info(f"Resuming training from checkpoint: {last_checkpoint}")
        trainer.train(resume_from_checkpoint=last_checkpoint)
    else:
        logger.info("Starting training from scratch.")
        trainer.train()

    logger.info("Training completed successfully.")

    # save the final checkpoint for resuming later
    logger.info(f"Saving trainer state to {trainer_output_dir}")
    trainer.save_model(trainer_output_dir)
   
    if trainer.is_world_process_zero():
        # save the final model
        # save_model(trainer.model, tokenizer, processor, args)
        # zero3_merge_and_save_inplace(trainer, tokenizer, processor, args)
        logger.info(f"Model saved to {args.model_output_dir}/{args.run_name}")



if __name__ == "__main__":
    parser = HfArgumentParser([LlavaTrainingArguments])
    (args,) = parser.parse_args_into_dataclasses()

    train(args)
