"""
FastTRL Data - Veri işleme yardımcıları
========================================
Sequence packing, chat template uygulama, data collator'lar.
"""

from __future__ import annotations

import warnings
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Union, Tuple

import numpy as np
import torch
from torch import Tensor


# ---------------------------------------------------------------------------
# Chat template
# ---------------------------------------------------------------------------

def apply_chat_template(
    example: Dict[str, Any],
    tokenizer,
    template: Optional[str] = None,
    tools: Optional[List[Dict]] = None,
    add_generation_prompt: bool = False,
) -> Dict[str, Any]:
    """
    Sohbet template'i uygular.

    Girdi formatları:
    - {"messages": [...]}          → çok turlu sohbet
    - {"prompt": "...", "completion": "..."} → tek turlu
    - {"text": "..."}              → ham metin (geçirilir)

    Args:
        example:              Veri örneği
        tokenizer:            HuggingFace tokenizer
        template:             Özel Jinja2 template (None = tokenizer'ın varsayılanı)
        tools:                Tool tanımları (function calling için)
        add_generation_prompt: Generasyon istemi ekle

    Returns:
        {"text": "<formatlanmış metin>"}
    """
    if template is not None:
        tokenizer = _clone_tokenizer_with_template(tokenizer, template)

    if "messages" in example:
        messages = example["messages"]
    elif "prompt" in example and "completion" in example:
        messages = [
            {"role": "user", "content": example["prompt"]},
            {"role": "assistant", "content": example["completion"]},
        ]
    elif "text" in example:
        return {"text": example["text"]}
    else:
        raise ValueError(
            "Örnek şu alanlardan birini içermeli: "
            "'messages', ('prompt' + 'completion'), veya 'text'. "
            f"Mevcut alanlar: {list(example.keys())}"
        )

    if not hasattr(tokenizer, "apply_chat_template"):
        raise AttributeError(
            "Tokenizer 'apply_chat_template' metoduna sahip değil. "
            "HuggingFace tokenizer kullandığınızdan emin olun."
        )

    kwargs: Dict[str, Any] = {
        "tokenize": False,
        "add_generation_prompt": add_generation_prompt,
    }
    if tools is not None:
        kwargs["tools"] = tools

    try:
        text = tokenizer.apply_chat_template(messages, **kwargs)
    except Exception as e:
        raise RuntimeError(
            f"Chat template uygulanırken hata oluştu: {e}\n"
            f"Mesajlar: {messages}"
        ) from e

    return {"text": text}


def _clone_tokenizer_with_template(tokenizer, template: str):
    """Tokenizer'ı özel template ile klonlar."""
    import copy
    tok = copy.copy(tokenizer)
    tok.chat_template = template
    return tok


# ---------------------------------------------------------------------------
# Sequence Packing
# ---------------------------------------------------------------------------

def pack_sequences(
    dataset,
    tokenizer,
    max_seq_length: int = 2048,
    text_field: str = "text",
    num_proc: int = 4,
    batch_size: int = 1000,
) -> Any:
    """
    Kısa dizileri birleştirerek dolu sekans paketleri oluşturur.
    Bu yöntem eğitim hızını önemli ölçüde artırır (padding israfı olmaz).

    Args:
        dataset:        HuggingFace Dataset
        tokenizer:      Tokenizer (EOS token bilgisi için)
        max_seq_length: Maksimum sekans uzunluğu
        text_field:     Metin alanı adı
        num_proc:       Paralel işlem sayısı
        batch_size:     Tokenizasyon batch boyutu

    Returns:
        Paketlenmiş HuggingFace Dataset
    """
    eos_id = tokenizer.eos_token_id
    if eos_id is None:
        warnings.warn(
            "Tokenizer'ın EOS token'ı yok. Paketleme sınır işaretleri olmayacak."
        )

    def tokenize_batch(batch):
        texts = batch[text_field]
        encoded = tokenizer(
            texts,
            add_special_tokens=True,
            truncation=False,
            return_attention_mask=False,
        )
        return {"input_ids": encoded["input_ids"]}

    def pack_batch(batch):
        all_ids: List[int] = []
        for ids in batch["input_ids"]:
            all_ids.extend(ids)
            if eos_id is not None:
                all_ids.append(eos_id)

        packed_ids: List[List[int]] = []
        packed_masks: List[List[int]] = []

        for i in range(0, len(all_ids), max_seq_length):
            chunk = all_ids[i : i + max_seq_length]
            if len(chunk) < max_seq_length:
                # Son chunk'ı padding ile doldur (ya da atla)
                pad_id = tokenizer.pad_token_id or eos_id or 0
                mask = [1] * len(chunk) + [0] * (max_seq_length - len(chunk))
                chunk = chunk + [pad_id] * (max_seq_length - len(chunk))
            else:
                mask = [1] * max_seq_length

            packed_ids.append(chunk)
            packed_masks.append(mask)

        return {
            "input_ids": packed_ids,
            "attention_mask": packed_masks,
        }

    # 1) Tokenize
    dataset = dataset.map(
        tokenize_batch,
        batched=True,
        batch_size=batch_size,
        num_proc=num_proc,
        remove_columns=dataset.column_names,
        desc="Tokenizasyon",
    )

    # 2) Pack
    dataset = dataset.map(
        pack_batch,
        batched=True,
        batch_size=batch_size,
        num_proc=num_proc,
        remove_columns=dataset.column_names,
        desc="Paketleme",
    )

    return dataset


# ---------------------------------------------------------------------------
# Data Collators
# ---------------------------------------------------------------------------

@dataclass
class DataCollatorForLanguageModeling:
    """
    Causal LM için data collator.
    - Otomatik padding
    - Labels = input_ids (causal LM)
    - Padding token'ları labels'da -100 yapılır (loss'a dahil edilmez)
    - Opsiyonel: sadece completion token'larında loss (completion_only=True)
    """

    tokenizer: Any
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    return_tensors: str = "pt"
    # Completion-only eğitim
    response_template_ids: Optional[List[int]] = None  # Response başlangıcı

    def __call__(
        self, features: List[Dict[str, Any]]
    ) -> Dict[str, Tensor]:
        pad_id = (
            self.tokenizer.pad_token_id
            if self.tokenizer.pad_token_id is not None
            else self.tokenizer.eos_token_id
        )
        if pad_id is None:
            raise ValueError(
                "Tokenizer'ın pad_token_id veya eos_token_id'si olmalı."
            )

        input_ids = [f["input_ids"] for f in features]
        attention_mask = [f.get("attention_mask", [1] * len(f["input_ids"])) for f in features]

        # Max uzunluk
        max_len = max(len(x) for x in input_ids)
        if self.max_length is not None:
            max_len = min(max_len, self.max_length)
        if self.pad_to_multiple_of is not None:
            max_len = (
                (max_len + self.pad_to_multiple_of - 1)
                // self.pad_to_multiple_of
                * self.pad_to_multiple_of
            )

        # Padding (sağa padding)
        padded_ids = []
        padded_mask = []
        for ids, mask in zip(input_ids, attention_mask):
            ids = ids[:max_len]
            mask = mask[:max_len]
            pad_len = max_len - len(ids)
            padded_ids.append(ids + [pad_id] * pad_len)
            padded_mask.append(mask + [0] * pad_len)

        input_ids_t = torch.tensor(padded_ids, dtype=torch.long)
        attention_mask_t = torch.tensor(padded_mask, dtype=torch.long)
        labels = input_ids_t.clone()

        # Padding token'larını label'da maskele
        labels[attention_mask_t == 0] = -100

        # Completion-only: sadece response kısmında loss
        if self.response_template_ids is not None:
            labels = self._mask_prompt_tokens(labels, input_ids_t)

        return {
            "input_ids": input_ids_t,
            "attention_mask": attention_mask_t,
            "labels": labels,
        }

    def _mask_prompt_tokens(
        self, labels: Tensor, input_ids: Tensor
    ) -> Tensor:
        """Prompt token'larını labels'da -100 yapar."""
        tpl = self.response_template_ids
        tpl_len = len(tpl)
        batch_size, seq_len = labels.shape

        for i in range(batch_size):
            ids = input_ids[i].tolist()
            # Template'i bul
            response_start = None
            for j in range(seq_len - tpl_len + 1):
                if ids[j : j + tpl_len] == tpl:
                    response_start = j + tpl_len
                    break

            if response_start is None:
                # Template bulunamadı: tüm satırı maskele
                warnings.warn(
                    "Örnek içinde response template bulunamadı, "
                    "tüm token'lar maskelendi."
                )
                labels[i] = -100
            else:
                labels[i, :response_start] = -100

        return labels


@dataclass
class DataCollatorForPreferenceOptimization:
    """
    DPO / Reward eğitimi için veri toplama.
    Her örnek (chosen, rejected) çiftinden oluşur.
    """

    tokenizer: Any
    max_length: int = 1024
    max_prompt_length: int = 512
    label_pad_token_id: int = -100
    padding_value: int = 0
    truncation_mode: str = "keep_end"

    def __call__(
        self, features: List[Dict[str, Any]]
    ) -> Dict[str, Tensor]:
        batch = {}
        for split in ("chosen", "rejected"):
            ids_key = f"{split}_input_ids"
            mask_key = f"{split}_attention_mask"
            label_key = f"{split}_labels"

            ids_list = [f[ids_key] for f in features]
            mask_list = [f[mask_key] for f in features]
            label_list = [f[label_key] for f in features]

            padded = self._pad_sequence(
                ids_list, mask_list, label_list
            )
            batch[ids_key] = padded["input_ids"]
            batch[mask_key] = padded["attention_mask"]
            batch[label_key] = padded["labels"]

        return batch

    def _pad_sequence(
        self,
        ids_list: List[List[int]],
        mask_list: List[List[int]],
        label_list: List[List[int]],
    ) -> Dict[str, Tensor]:
        max_len = min(
            max(len(x) for x in ids_list),
            self.max_length,
        )
        pad_id = (
            self.tokenizer.pad_token_id
            if self.tokenizer.pad_token_id is not None
            else 0
        )

        padded_ids, padded_mask, padded_labels = [], [], []

        for ids, mask, labels in zip(ids_list, mask_list, label_list):
            if self.truncation_mode == "keep_end":
                ids = ids[-max_len:]
                mask = mask[-max_len:]
                labels = labels[-max_len:]
            else:
                ids = ids[:max_len]
                mask = mask[:max_len]
                labels = labels[:max_len]

            pad_len = max_len - len(ids)
            padded_ids.append(ids + [pad_id] * pad_len)
            padded_mask.append(mask + [0] * pad_len)
            padded_labels.append(labels + [self.label_pad_token_id] * pad_len)

        return {
            "input_ids": torch.tensor(padded_ids, dtype=torch.long),
            "attention_mask": torch.tensor(padded_mask, dtype=torch.long),
            "labels": torch.tensor(padded_labels, dtype=torch.long),
        }
