"""
FastTRL Models - Model sarmalayıcılar ve yardımcıları
=====================================================
PPO için value head, referans model klonlama.
"""

from __future__ import annotations

import copy
import warnings
from typing import Optional, Tuple, Union, Dict, Any

import torch
import torch.nn as nn
from torch import Tensor

from fasttrl.utils import freeze_model, get_logger

logger = get_logger(__name__)


# ---------------------------------------------------------------------------
# Value Head
# ---------------------------------------------------------------------------

class ValueHead(nn.Module):
    """
    LM üstüne eklenen değer fonksiyonu başlığı.
    PPO'da kullanılır.

    Mimari:
        hidden → Dropout → Linear(hidden, hidden//2) → GELU → Linear(hidden//2, 1)
    """

    def __init__(
        self,
        hidden_size: int,
        dropout: float = 0.1,
        summary_last_kwargs: Optional[Dict] = None,
    ) -> None:
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.summary = nn.Linear(hidden_size, 1)

        # Küçük ağırlık başlatma (stabil eğitim için)
        nn.init.normal_(self.summary.weight, std=0.01)
        nn.init.zeros_(self.summary.bias)

    def forward(self, hidden_states: Tensor) -> Tensor:
        """
        Args:
            hidden_states: [batch, seq_len, hidden_size]
        Returns:
            values: [batch, seq_len]
        """
        x = self.dropout(hidden_states)
        return self.summary(x).squeeze(-1)


# ---------------------------------------------------------------------------
# AutoModelForCausalLMWithValueHead
# ---------------------------------------------------------------------------

class AutoModelForCausalLMWithValueHead(nn.Module):
    """
    Causal LM + Value Head birleşimi.
    PPO eğitimi için tasarlanmıştır.

    Kullanım::

        model = AutoModelForCausalLMWithValueHead.from_pretrained(
            "meta-llama/Llama-3-8B",
            value_head_dropout=0.1,
        )
        # Hem LM loss hem value tahmini için:
        outputs = model(input_ids, labels=labels)
        lm_loss = outputs.loss
        values  = outputs.value
    """

    SUPPORTED_MODULES = (
        "transformer",   # GPT-2 tarzı
        "model",         # LLaMA, Mistral tarzı
        "gpt_neox",      # NeoX tarzı
        "bloom",         # BLOOM tarzı
    )

    def __init__(
        self,
        pretrained_model: nn.Module,
        value_head_dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.pretrained_model = pretrained_model
        self.config = pretrained_model.config

        hidden_size = getattr(self.config, "hidden_size", None)
        if hidden_size is None:
            hidden_size = getattr(self.config, "n_embd", None)
        if hidden_size is None:
            raise ValueError(
                "Model config'inden hidden_size alınamadı. "
                "'hidden_size' veya 'n_embd' alanlarından biri olmalı."
            )

        self.value_head = ValueHead(hidden_size, dropout=value_head_dropout)
        self._is_peft_model = False

        # PEFT kontrolü
        try:
            from peft import PeftModel
            if isinstance(pretrained_model, PeftModel):
                self._is_peft_model = True
                logger.info("PEFT modeli tespit edildi.")
        except ImportError:
            pass

    @classmethod
    def from_pretrained(
        cls,
        model_name_or_path: str,
        value_head_dropout: float = 0.1,
        peft_config=None,
        **model_kwargs,
    ) -> "AutoModelForCausalLMWithValueHead":
        """
        Pretrained model yükler ve value head ekler.

        Args:
            model_name_or_path: HuggingFace model ismi veya lokal yol
            value_head_dropout: Value head dropout oranı
            peft_config:        PEFT konfigürasyonu (opsiyonel)
            **model_kwargs:     AutoModelForCausalLM.from_pretrained'a aktarılır
        """
        from transformers import AutoModelForCausalLM

        logger.info(f"Model yükleniyor: {model_name_or_path}")

        model_kwargs.setdefault("trust_remote_code", True)

        try:
            pretrained = AutoModelForCausalLM.from_pretrained(
                model_name_or_path, **model_kwargs
            )
        except Exception as e:
            raise RuntimeError(
                f"Model yüklenirken hata oluştu: {e}\n"
                f"Model: {model_name_or_path}\n"
                f"Kwargs: {model_kwargs}"
            ) from e

        if peft_config is not None:
            try:
                from peft import get_peft_model
                pretrained = get_peft_model(pretrained, peft_config)
                logger.info("PEFT adaptörü eklendi.")
            except ImportError as e:
                raise ImportError(
                    "peft kütüphanesi bulunamadı. "
                    "Kurmak için: pip install peft"
                ) from e

        return cls(pretrained, value_head_dropout=value_head_dropout)

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
        past_key_values=None,
        use_cache: bool = False,
        return_dict: bool = True,
        **kwargs,
    ) -> "ValueHeadModelOutput":
        """
        Returns:
            ValueHeadModelOutput ile:
            - loss:         LM loss (labels verilirse)
            - logits:       [batch, seq_len, vocab_size]
            - value:        [batch, seq_len]
            - past_key_values
            - hidden_states
        """
        outputs = self.pretrained_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_hidden_states=True,
            return_dict=True,
            **kwargs,
        )

        # Son hidden state'i al
        if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
            last_hidden = outputs.hidden_states[-1]
        else:
            # Bazı modeller hidden_states döndürmeyebilir
            last_hidden = outputs.logits  # fallback (hassas değil)
            warnings.warn(
                "Model hidden_states döndürmüyor. "
                "Lütfen output_hidden_states=True desteği olan model kullanın."
            )

        value = self.value_head(last_hidden)

        return ValueHeadModelOutput(
            loss=outputs.loss,
            logits=outputs.logits,
            value=value,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
        )

    def generate(self, *args, **kwargs):
        """Generasyonu pretrained modele devreder."""
        return self.pretrained_model.generate(*args, **kwargs)

    def get_lm_head(self) -> nn.Module:
        """LM başlığını döner."""
        for attr in ("lm_head", "embed_out", "cls"):
            if hasattr(self.pretrained_model, attr):
                return getattr(self.pretrained_model, attr)
        raise AttributeError("LM başlığı bulunamadı.")

    def state_dict_without_value_head(self) -> Dict[str, Tensor]:
        """Sadece temel model ağırlıklarını döner (value head hariç)."""
        return {
            k.replace("pretrained_model.", ""): v
            for k, v in self.state_dict().items()
            if not k.startswith("value_head.")
        }

    def push_to_hub(self, *args, **kwargs):
        """Sadece base modeli Hub'a gönderir (value head olmadan)."""
        self.pretrained_model.push_to_hub(*args, **kwargs)


class ValueHeadModelOutput:
    """AutoModelForCausalLMWithValueHead çıkış nesnesi."""

    __slots__ = ("loss", "logits", "value", "past_key_values", "hidden_states")

    def __init__(
        self,
        loss=None,
        logits=None,
        value=None,
        past_key_values=None,
        hidden_states=None,
    ):
        self.loss = loss
        self.logits = logits
        self.value = value
        self.past_key_values = past_key_values
        self.hidden_states = hidden_states


# ---------------------------------------------------------------------------
# Referans model
# ---------------------------------------------------------------------------

def create_reference_model(
    model: nn.Module,
    num_shared_layers: Optional[int] = None,
) -> nn.Module:
    """
    PPO/DPO için referans model oluşturur.

    İki mod:
    1. num_shared_layers=None: Modelin tam kopyası, ağırlıklar dondurulur.
       Bellek tüketimi 2x artar.
    2. num_shared_layers=N: İlk N katman paylaşılır (bellek tasarrufu).
       Sadece son katmanlar bağımsız.

    Args:
        model:             Kaynak model
        num_shared_layers: Paylaşılan katman sayısı (None = tam kopya)
    """
    if num_shared_layers is None:
        ref_model = copy.deepcopy(model)
        freeze_model(ref_model)
        ref_model.eval()
        logger.info("Tam referans model kopyası oluşturuldu (ağırlıklar donduruldu).")
        return ref_model

    # Kısmi paylaşım: ilk N katman paylaşılır
    ref_model = copy.deepcopy(model)
    freeze_model(ref_model)

    # Paylaşılan katmanları orijinal modele bağla
    policy_layers = _get_transformer_layers(model)
    ref_layers = _get_transformer_layers(ref_model)

    if policy_layers is None or ref_layers is None:
        warnings.warn(
            "Katmanlar bulunamadı, tam kopya oluşturuluyor. "
            "Bellek optimizasyonu devre dışı."
        )
        return ref_model

    n = min(num_shared_layers, len(policy_layers))
    for i in range(n):
        ref_layers[i] = policy_layers[i]

    logger.info(
        f"Referans model: ilk {n} katman paylaşılıyor "
        f"(toplam {len(policy_layers)} katman)."
    )

    ref_model.eval()
    return ref_model


def _get_transformer_layers(model: nn.Module) -> Optional[nn.ModuleList]:
    """Transformer katmanlarını döner (çeşitli mimari isimlerini dener)."""
    layer_attrs = [
        "model.layers",       # LLaMA, Mistral
        "transformer.h",      # GPT-2
        "gpt_neox.layers",    # NeoX
        "model.decoder.layers", # OPT
        "transformer.blocks", # MPT
        "bloom.h",            # BLOOM
    ]
    for attr_path in layer_attrs:
        obj = model
        try:
            for part in attr_path.split("."):
                obj = getattr(obj, part)
            if isinstance(obj, nn.ModuleList):
                return obj
        except AttributeError:
            continue

    # İç içe arama
    for _, module in model.named_modules():
        if isinstance(module, nn.ModuleList) and len(module) > 0:
            first = module[0]
            if any(
                hasattr(first, attr)
                for attr in ("self_attn", "attention", "attn")
            ):
                return module

    return None
