"""
FastTRL SFTTrainer - Supervised Fine-Tuning
===========================================
Hızlı, güvenilir ve esnek SFT eğitimi.
Sequence packing, NEFTune, completion-only loss desteği.
"""

from __future__ import annotations

import os
import math
import warnings
from typing import Any, Callable, Dict, List, Optional, Union

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

from fasttrl.config import SFTConfig
from fasttrl.data import (
    DataCollatorForLanguageModeling,
    apply_chat_template,
    pack_sequences,
)
from fasttrl.utils import (
    get_logger,
    set_seed,
    print_trainable_parameters,
    get_dtype,
    compute_accuracy,
    compute_perplexity,
)

logger = get_logger(__name__)


class SFTTrainer:
    """
    Supervised Fine-Tuning Trainer.

    Özellikler:
    - Sequence packing (pad israfı yok, ~%30 hız artışı)
    - NEFTune gürültü artırımı
    - Completion-only loss (sadece yanıt token'larında loss)
    - Mixed precision (fp16/bf16)
    - Gradient checkpointing
    - Otomatik chat template uygulama
    - Checkpoint kayıt ve yükleme

    Kullanım::

        from fasttrl import SFTTrainer, SFTConfig
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from datasets import load_dataset

        model = AutoModelForCausalLM.from_pretrained("gpt2")
        tokenizer = AutoTokenizer.from_pretrained("gpt2")
        dataset = load_dataset("imdb", split="train")

        config = SFTConfig(
            output_dir="./output",
            max_seq_length=512,
            num_train_epochs=3,
            per_device_train_batch_size=8,
            learning_rate=2e-5,
        )

        trainer = SFTTrainer(
            model=model,
            tokenizer=tokenizer,
            args=config,
            train_dataset=dataset,
        )
        trainer.train()
    """

    def __init__(
        self,
        model: Union[nn.Module, str],
        tokenizer,
        args: SFTConfig,
        train_dataset,
        eval_dataset=None,
        data_collator: Optional[Any] = None,
        formatting_func: Optional[Callable] = None,
        compute_metrics: Optional[Callable] = None,
        callbacks: Optional[List] = None,
        peft_config=None,
    ) -> None:
        self.args = args
        self.tokenizer = tokenizer
        self.compute_metrics = compute_metrics
        self.callbacks = callbacks or []
        self._neftune_hook = None

        set_seed(args.seed)

        # Model yükle
        self.model = self._load_model(model, peft_config)

        # Tokenizer düzelt
        self._fix_tokenizer()

        # Veri setlerini hazırla
        self.train_dataset = self._prepare_dataset(
            train_dataset, formatting_func, packing=args.packing
        )
        self.eval_dataset = (
            self._prepare_dataset(eval_dataset, formatting_func, packing=args.eval_packing)
            if eval_dataset is not None
            else None
        )

        # Collator
        if data_collator is None:
            response_template_ids = None
            if args.completion_only and args.response_template:
                response_template_ids = tokenizer.encode(
                    args.response_template, add_special_tokens=False
                )
            self.data_collator = DataCollatorForLanguageModeling(
                tokenizer=self.tokenizer,
                max_length=args.max_seq_length,
                response_template_ids=response_template_ids,
            )
        else:
            self.data_collator = data_collator

        # Optimizer & Scheduler (tembel başlatma)
        self.optimizer = None
        self.lr_scheduler = None
        self._global_step = 0
        self._current_epoch = 0

        print_trainable_parameters(self.model)
        logger.info(
            f"SFTTrainer hazır. "
            f"Eğitim örnekleri: {len(self.train_dataset)}, "
            f"Packing: {args.packing}, "
            f"Max uzunluk: {args.max_seq_length}"
        )

    # ------------------------------------------------------------------
    # Hazırlık
    # ------------------------------------------------------------------

    def _load_model(self, model: Union[nn.Module, str], peft_config) -> nn.Module:
        """Model yükler veya varolan modeli alır."""
        if isinstance(model, str):
            from transformers import AutoModelForCausalLM

            dtype = get_dtype(self.args.fp16, self.args.bf16)
            logger.info(f"Model yükleniyor: {model} (dtype={dtype})")
            model = AutoModelForCausalLM.from_pretrained(
                model,
                torch_dtype=dtype,
                device_map=self.args.device_map,
                trust_remote_code=True,
            )

        if peft_config is not None:
            try:
                from peft import get_peft_model, prepare_model_for_kbit_training

                if getattr(model, "is_loaded_in_8bit", False) or getattr(
                    model, "is_loaded_in_4bit", False
                ):
                    model = prepare_model_for_kbit_training(model)

                model = get_peft_model(model, peft_config)
                logger.info("PEFT adaptörü eklendi.")
            except ImportError:
                raise ImportError(
                    "peft paketi bulunamadı. pip install peft ile kurun."
                )

        if self.args.gradient_checkpointing:
            if hasattr(model, "enable_input_require_grads"):
                model.enable_input_require_grads()
            model.gradient_checkpointing_enable()
            logger.info("Gradient checkpointing etkinleştirildi.")

        return model

    def _fix_tokenizer(self) -> None:
        """Tokenizer'ın gerekli token'larını düzeltir."""
        if self.tokenizer.pad_token is None:
            if self.tokenizer.eos_token is not None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
                logger.info("pad_token = eos_token olarak ayarlandı.")
            else:
                self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
                self.model.resize_token_embeddings(len(self.tokenizer))
                logger.info("[PAD] token eklendi ve embedding boyutu güncellendi.")

        self.tokenizer.padding_side = "right"

    def _prepare_dataset(self, dataset, formatting_func, packing: bool = True):
        """Veri setini tokenize eder ve isteğe bağlı olarak paketler."""
        if dataset is None:
            return None

        # Chat template veya formatting func uygula
        if formatting_func is not None:
            dataset = dataset.map(
                lambda ex: {"text": formatting_func(ex)},
                desc="Formatting",
            )
        elif self.args.chat_template or (
            hasattr(self.tokenizer, "chat_template")
            and self.tokenizer.chat_template is not None
            and "messages" in dataset.column_names
        ):
            dataset = dataset.map(
                lambda ex: apply_chat_template(
                    ex, self.tokenizer, template=self.args.chat_template
                ),
                desc="Chat template uygulama",
            )

        # Text alanını kontrol et
        text_field = self.args.dataset_text_field
        if text_field not in dataset.column_names:
            # Eğer tek bir metin sütunu varsa otomatik seç
            text_cols = [
                c for c in dataset.column_names
                if dataset.features[c].dtype == "string"
                or str(dataset.features[c]) == "Value(dtype='string', id=None)"
            ]
            if len(text_cols) == 1:
                text_field = text_cols[0]
                logger.info(f"Metin alanı otomatik belirlendi: '{text_field}'")
            elif not text_cols:
                # Zaten tokenize edilmiş olabilir
                if "input_ids" in dataset.column_names:
                    return dataset
                raise ValueError(
                    f"Veri setinde metin alanı bulunamadı. "
                    f"Mevcut sütunlar: {dataset.column_names}"
                )

        if packing:
            dataset = pack_sequences(
                dataset,
                tokenizer=self.tokenizer,
                max_seq_length=self.args.max_seq_length,
                text_field=text_field,
                num_proc=self.args.dataloader_num_workers,
            )
        else:
            def tokenize(batch):
                return self.tokenizer(
                    batch[text_field],
                    truncation=True,
                    max_length=self.args.max_seq_length,
                    padding=False,
                )

            dataset = dataset.map(
                tokenize,
                batched=True,
                remove_columns=[
                    c for c in dataset.column_names
                    if c not in ("input_ids", "attention_mask")
                ],
                desc="Tokenizasyon",
            )

        return dataset

    # ------------------------------------------------------------------
    # Optimizer & Scheduler
    # ------------------------------------------------------------------

    def _create_optimizer(self) -> torch.optim.Optimizer:
        """Optimizer oluşturur."""
        decay_params = []
        no_decay_params = []

        for name, param in self.model.named_parameters():
            if not param.requires_grad:
                continue
            if any(nd in name for nd in ("bias", "layer_norm", "layernorm")):
                no_decay_params.append(param)
            else:
                decay_params.append(param)

        param_groups = [
            {"params": decay_params, "weight_decay": self.args.weight_decay},
            {"params": no_decay_params, "weight_decay": 0.0},
        ]

        optim_name = self.args.optim

        # Fused optimizer deneme (CUDA gerektirir)
        if "fused" in optim_name:
            try:
                optimizer = torch.optim.AdamW(
                    param_groups,
                    lr=self.args.learning_rate,
                    fused=True,
                )
                logger.info("Fused AdamW optimizer kullanılıyor.")
                return optimizer
            except (TypeError, RuntimeError) as e:
                logger.warning(
                    f"Fused optimizer başlatılamadı ({e}), "
                    "standart AdamW'ya geçildi."
                )

        optimizer = torch.optim.AdamW(
            param_groups,
            lr=self.args.learning_rate,
        )
        logger.info("AdamW optimizer kullanılıyor.")
        return optimizer

    def _create_scheduler(
        self, optimizer, num_training_steps: int
    ):
        """Öğrenme oranı zamanlayıcısı oluşturur."""
        from transformers import get_scheduler

        num_warmup_steps = math.ceil(
            num_training_steps * self.args.warmup_ratio
        )

        scheduler = get_scheduler(
            name=self.args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=num_warmup_steps,
            num_training_steps=num_training_steps,
        )
        logger.info(
            f"LR scheduler: {self.args.lr_scheduler_type}, "
            f"warmup: {num_warmup_steps}/{num_training_steps} adım"
        )
        return scheduler

    # ------------------------------------------------------------------
    # NEFTune
    # ------------------------------------------------------------------

    def _attach_neftune(self) -> None:
        """NEFTune gürültüsünü etkinleştirir."""
        alpha = self.args.neftune_noise_alpha
        if alpha is None:
            return

        embeddings = self._get_embedding_layer()
        if embeddings is None:
            warnings.warn("Embedding katmanı bulunamadı, NEFTune devre dışı.")
            return

        def neftune_hook(module, input, output):
            if module.training:
                dim = output.size(-1)
                seq_len = output.size(1)
                mag_norm = alpha / math.sqrt(seq_len * dim)
                noise = torch.zeros_like(output).uniform_(-1, 1) * mag_norm
                return output + noise
            return output

        self._neftune_hook = embeddings.register_forward_hook(neftune_hook)
        logger.info(f"NEFTune etkinleştirildi (alpha={alpha}).")

    def _detach_neftune(self) -> None:
        """NEFTune kancasını kaldırır."""
        if self._neftune_hook is not None:
            self._neftune_hook.remove()
            self._neftune_hook = None

    def _get_embedding_layer(self) -> Optional[nn.Module]:
        for attr in ("model.embed_tokens", "transformer.wte", "gpt_neox.embed_in",
                     "embed_tokens", "word_embeddings"):
            obj = self.model
            try:
                for part in attr.split("."):
                    obj = getattr(obj, part)
                if isinstance(obj, nn.Embedding):
                    return obj
            except AttributeError:
                continue
        return None

    # ------------------------------------------------------------------
    # Eğitim döngüsü
    # ------------------------------------------------------------------

    def train(self) -> Dict[str, float]:
        """
        Ana eğitim döngüsü.

        Returns:
            {"train_loss": ..., "train_perplexity": ...}
        """
        args = self.args

        train_loader = DataLoader(
            self.train_dataset,
            batch_size=args.per_device_train_batch_size,
            shuffle=True,
            collate_fn=self.data_collator,
            num_workers=args.dataloader_num_workers,
            pin_memory=args.dataloader_pin_memory,
            drop_last=False,
        )

        steps_per_epoch = math.ceil(
            len(train_loader) / args.gradient_accumulation_steps
        )
        total_steps = steps_per_epoch * args.num_train_epochs

        self.optimizer = self._create_optimizer()
        self.lr_scheduler = self._create_scheduler(self.optimizer, total_steps)

        # Mixed precision scaler
        scaler = None
        if args.fp16:
            scaler = torch.cuda.amp.GradScaler()

        device = next(self.model.parameters()).device

        self._attach_neftune()
        self.model.train()

        logger.info(
            f"Eğitim başlıyor: {args.num_train_epochs} epoch, "
            f"{total_steps} toplam adım"
        )

        total_loss = 0.0
        log_loss = 0.0
        best_eval_loss = float("inf")

        for epoch in range(args.num_train_epochs):
            self._current_epoch = epoch
            epoch_loss = 0.0

            for step, batch in enumerate(train_loader):
                batch = {k: v.to(device) for k, v in batch.items()}

                # Forward pass
                with torch.cuda.amp.autocast(enabled=args.fp16 or args.bf16):
                    outputs = self.model(**batch)
                    loss = outputs.loss / args.gradient_accumulation_steps

                # Backward pass
                if scaler is not None:
                    scaler.scale(loss).backward()
                else:
                    loss.backward()

                epoch_loss += loss.item() * args.gradient_accumulation_steps

                # Gradient accumulation adımı tamamlandıysa güncelle
                is_update_step = (
                    (step + 1) % args.gradient_accumulation_steps == 0
                    or (step + 1) == len(train_loader)
                )

                if is_update_step:
                    # Gradient clipping
                    if args.max_grad_norm is not None:
                        if scaler is not None:
                            scaler.unscale_(self.optimizer)
                        torch.nn.utils.clip_grad_norm_(
                            self.model.parameters(), args.max_grad_norm
                        )

                    if scaler is not None:
                        scaler.step(self.optimizer)
                        scaler.update()
                    else:
                        self.optimizer.step()

                    self.lr_scheduler.step()
                    self.optimizer.zero_grad()
                    self._global_step += 1

                    # Loglama
                    if self._global_step % args.logging_steps == 0:
                        lr = self.lr_scheduler.get_last_lr()[0]
                        avg_loss = epoch_loss / (step + 1)
                        logger.info(
                            f"Epoch {epoch+1}/{args.num_train_epochs} | "
                            f"Adım {self._global_step}/{total_steps} | "
                            f"Loss: {avg_loss:.4f} | "
                            f"LR: {lr:.2e}"
                        )

                    # Kaydetme
                    if (
                        args.save_steps > 0
                        and self._global_step % args.save_steps == 0
                    ):
                        self._save_checkpoint(f"checkpoint-{self._global_step}")

                    # Değerlendirme
                    if (
                        self.eval_dataset is not None
                        and args.eval_steps > 0
                        and self._global_step % args.eval_steps == 0
                    ):
                        eval_results = self.evaluate()
                        eval_loss = eval_results.get("eval_loss", float("inf"))
                        if eval_loss < best_eval_loss:
                            best_eval_loss = eval_loss
                            self._save_checkpoint("best_model")
                        self.model.train()

            total_loss += epoch_loss / len(train_loader)
            logger.info(
                f"Epoch {epoch+1} tamamlandı. "
                f"Ortalama loss: {epoch_loss / len(train_loader):.4f}"
            )

        self._detach_neftune()

        # Son modeli kaydet
        self._save_checkpoint("final_model")

        avg_loss = total_loss / args.num_train_epochs
        result = {
            "train_loss": avg_loss,
            "train_perplexity": compute_perplexity(avg_loss),
        }
        logger.info(f"Eğitim tamamlandı: {result}")
        return result

    # ------------------------------------------------------------------
    # Değerlendirme
    # ------------------------------------------------------------------

    def evaluate(self, eval_dataset=None) -> Dict[str, float]:
        """
        Değerlendirme döngüsü.

        Returns:
            {"eval_loss": ..., "eval_perplexity": ..., "eval_accuracy": ...}
        """
        dataset = eval_dataset or self.eval_dataset
        if dataset is None:
            raise ValueError("Değerlendirme için eval_dataset gerekli.")

        eval_loader = DataLoader(
            dataset,
            batch_size=self.args.per_device_eval_batch_size,
            shuffle=False,
            collate_fn=self.data_collator,
            num_workers=self.args.dataloader_num_workers,
        )

        device = next(self.model.parameters()).device
        self.model.eval()

        total_loss = 0.0
        total_correct = 0
        total_tokens = 0
        n_batches = 0

        with torch.no_grad():
            for batch in eval_loader:
                batch = {k: v.to(device) for k, v in batch.items()}
                with torch.cuda.amp.autocast(
                    enabled=self.args.fp16 or self.args.bf16
                ):
                    outputs = self.model(**batch)

                total_loss += outputs.loss.item()
                n_batches += 1

                # Doğruluk
                if outputs.logits is not None:
                    labels = batch.get("labels")
                    if labels is not None:
                        acc = compute_accuracy(outputs.logits, labels)
                        valid_tokens = (labels != -100).sum().item()
                        total_correct += acc * valid_tokens
                        total_tokens += valid_tokens

        avg_loss = total_loss / n_batches if n_batches > 0 else float("inf")
        accuracy = total_correct / total_tokens if total_tokens > 0 else 0.0

        results = {
            "eval_loss": avg_loss,
            "eval_perplexity": compute_perplexity(avg_loss),
            "eval_accuracy": accuracy,
        }
        logger.info(f"Değerlendirme: {results}")
        return results

    # ------------------------------------------------------------------
    # Checkpoint
    # ------------------------------------------------------------------

    def _save_checkpoint(self, name: str) -> None:
        """Checkpoint kaydeder."""
        path = os.path.join(self.args.output_dir, name)
        os.makedirs(path, exist_ok=True)

        # PEFT modelini kaydet
        if hasattr(self.model, "save_pretrained"):
            self.model.save_pretrained(path)
        else:
            torch.save(self.model.state_dict(), os.path.join(path, "model.pt"))

        self.tokenizer.save_pretrained(path)
        self.args.save(os.path.join(path, "training_args.json"))
        logger.info(f"Checkpoint kaydedildi: {path}")

    def save_model(self, output_dir: Optional[str] = None) -> None:
        """Modeli kaydeder."""
        path = output_dir or self.args.output_dir
        self._save_checkpoint(path)

    def push_to_hub(self, repo_id: str, **kwargs) -> None:
        """Modeli HuggingFace Hub'a gönderir."""
        if hasattr(self.model, "push_to_hub"):
            self.model.push_to_hub(repo_id, **kwargs)
            self.tokenizer.push_to_hub(repo_id, **kwargs)
            logger.info(f"Model Hub'a gönderildi: {repo_id}")
        else:
            raise AttributeError(
                "Model 'push_to_hub' metoduna sahip değil."
            )
