"""
LoRA fine-tuning method for comparison experiments.

Trains a LoRA adapter on benchmark tasks using supervised fine-tuning (SFT)
on positive responses from contrastive pairs.
"""

from __future__ import annotations

import json
import tempfile
from argparse import Namespace
from pathlib import Path
from typing import TYPE_CHECKING

import torch
from datasets import Dataset
from peft import LoraConfig, TaskType, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer

if TYPE_CHECKING:
    from wisent.core.models.wisent_model import WisentModel

__all__ = ["generate_steering_vector", "apply_lora_to_model", "remove_lora"]


# Default LoRA configurations per model architecture
LORA_TARGET_MODULES = {
    "gemma": ["q_proj", "k_proj", "v_proj", "o_proj"],
    "llama": ["q_proj", "k_proj", "v_proj", "o_proj"],
    "mistral": ["q_proj", "k_proj", "v_proj", "o_proj"],
    "phi": ["q_proj", "k_proj", "v_proj", "dense"],
    "gpt_neo": ["q_proj", "v_proj"],
    "gpt2": ["c_attn"],
    "default": "all-linear",
}


def get_target_modules(model_name: str) -> str | list[str]:
    """Get LoRA target modules based on model architecture."""
    model_name_lower = model_name.lower()

    for arch, modules in LORA_TARGET_MODULES.items():
        if arch in model_name_lower:
            return modules

    return LORA_TARGET_MODULES["default"]


def prepare_sft_dataset(
    pairs: list[dict],
    tokenizer,
    max_length: int = 512,
) -> Dataset:
    """
    Prepare dataset for SFT from contrastive pairs.

    Uses only positive responses for training.

    Args:
        pairs: List of contrastive pairs
        tokenizer: Tokenizer for formatting
        max_length: Maximum sequence length

    Returns:
        HuggingFace Dataset ready for SFTTrainer
    """
    formatted_examples = []

    for pair in pairs:
        prompt = pair["prompt"]
        positive_response = pair["positive_response"]["model_response"]

        # Format as chat if tokenizer supports it, otherwise simple format
        if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
            messages = [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": positive_response},
            ]
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=False,
            )
        else:
            # Simple format for base models
            text = f"Q: {prompt}\nA: {positive_response}"

        formatted_examples.append({"text": text})

    return Dataset.from_list(formatted_examples)


def generate_steering_vector(
    task: str,
    model_name: str,
    output_path: str | Path,
    trait_label: str = "correctness",
    num_pairs: int = 50,
    method: str = "lora",
    layers: str | None = None,  # Not used for LoRA but kept for interface compatibility
    device: str = "cuda:0",
    keep_intermediate: bool = False,
    extraction_strategy: str = "chat_last",  # Not used for LoRA but kept for interface compatibility
    # LoRA-specific parameters
    lora_r: int = 16,
    lora_alpha: int = 32,
    lora_dropout: float = 0.05,
    learning_rate: float = 2e-4,
    num_epochs: int = 3,
    batch_size: int = 4,
    max_length: int = 512,
    **kwargs,
) -> Path:
    """
    Train a LoRA adapter using SFT on positive responses.

    Args:
        task: lm-eval task name (e.g., 'boolq', 'cb')
        model_name: HuggingFace model name
        output_path: Where to save the LoRA adapter
        trait_label: Label for the trait being trained
        num_pairs: Number of training examples to use
        method: Method name (should be 'lora')
        layers: Not used for LoRA (kept for interface compatibility)
        device: Device to train on
        keep_intermediate: Whether to keep intermediate files
        extraction_strategy: Not used for LoRA (kept for interface compatibility)
        lora_r: LoRA rank
        lora_alpha: LoRA alpha scaling factor
        lora_dropout: LoRA dropout
        learning_rate: Training learning rate
        num_epochs: Number of training epochs
        batch_size: Training batch size
        max_length: Maximum sequence length

    Returns:
        Path to the saved LoRA adapter directory
    """
    import gc

    output_path = Path(output_path)

    # Step 1: Generate contrastive pairs
    print(f"Step 1: Generating training data from task: {task}")
    from wisent.core.cli.generate_pairs_from_task import execute_generate_pairs_from_task

    pairs_file = tempfile.NamedTemporaryFile(mode='w', suffix='_pairs.json', delete=False).name
    pairs_args = Namespace(
        task_name=task,
        limit=num_pairs,
        output=pairs_file,
        seed=42,
        verbose=False,
    )
    execute_generate_pairs_from_task(pairs_args)

    with open(pairs_file) as f:
        pairs_data = json.load(f)
    pairs = pairs_data["pairs"]
    print(f"   Loaded {len(pairs)} training examples")

    # Step 2: Load model and tokenizer
    print(f"\nStep 2: Loading model {model_name}...")

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map=device,
        trust_remote_code=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Step 3: Configure LoRA
    print(f"\nStep 3: Configuring LoRA (r={lora_r}, alpha={lora_alpha})...")

    target_modules = get_target_modules(model_name)
    print(f"   Target modules: {target_modules}")

    lora_config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=target_modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # Step 4: Prepare dataset
    print(f"\nStep 4: Preparing SFT dataset...")
    train_dataset = prepare_sft_dataset(pairs, tokenizer, max_length=max_length)
    print(f"   Dataset size: {len(train_dataset)} examples")

    # Step 5: Training
    print(f"\nStep 5: Training LoRA adapter...")

    # Create temporary directory for training outputs
    training_output_dir = tempfile.mkdtemp(prefix="lora_training_")

    training_args = TrainingArguments(
        output_dir=training_output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=1,
        learning_rate=learning_rate,
        weight_decay=0.01,
        warmup_ratio=0.1,
        logging_steps=10,
        save_strategy="no",  # Don't save checkpoints
        fp16=True,
        report_to="none",  # Disable wandb/tensorboard
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        processing_class=tokenizer,
        max_seq_length=max_length,
    )

    trainer.train()

    # Step 6: Save LoRA adapter
    print(f"\nStep 6: Saving LoRA adapter to {output_path}...")
    output_path.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output_path)
    tokenizer.save_pretrained(output_path)

    # Save metadata
    metadata = {
        "method": "lora",
        "model": model_name,
        "task": task,
        "trait_label": trait_label,
        "num_pairs": len(pairs),
        "lora_config": {
            "r": lora_r,
            "alpha": lora_alpha,
            "dropout": lora_dropout,
            "target_modules": target_modules if isinstance(target_modules, list) else [target_modules],
        },
        "training_config": {
            "learning_rate": learning_rate,
            "num_epochs": num_epochs,
            "batch_size": batch_size,
            "max_length": max_length,
        },
    }

    with open(output_path / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    # Cleanup
    del model, trainer
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    if not keep_intermediate:
        import os
        os.unlink(pairs_file)
        import shutil
        shutil.rmtree(training_output_dir, ignore_errors=True)

    print(f"\nLoRA adapter saved to {output_path}")
    return output_path


def apply_lora_to_model(wisent_model: "WisentModel", lora_path: str | Path) -> None:
    """
    Apply a trained LoRA adapter to a WisentModel.

    Args:
        wisent_model: WisentModel instance
        lora_path: Path to the saved LoRA adapter
    """
    from peft import PeftModel

    lora_path = Path(lora_path)

    # Check if model already has adapters
    if hasattr(wisent_model.hf_model, 'peft_config'):
        # Model already has PEFT, just load new adapter
        wisent_model.hf_model.load_adapter(str(lora_path), adapter_name="steering")
        wisent_model.hf_model.set_adapter("steering")
    else:
        # Wrap model with PEFT
        wisent_model.hf_model = PeftModel.from_pretrained(
            wisent_model.hf_model,
            str(lora_path),
            adapter_name="steering",
        )

    print(f"LoRA adapter loaded from {lora_path}")


def remove_lora(wisent_model: "WisentModel") -> None:
    """
    Remove/disable LoRA adapter from a WisentModel.

    Args:
        wisent_model: WisentModel instance with LoRA applied
    """
    if hasattr(wisent_model.hf_model, 'disable_adapters'):
        wisent_model.hf_model.disable_adapters()
        print("LoRA adapter disabled")
    elif hasattr(wisent_model.hf_model, 'base_model'):
        # Unwrap the model
        wisent_model.hf_model = wisent_model.hf_model.base_model.model
        print("LoRA adapter removed")
