# Copyright (c) ModelScope Contributors. All rights reserved.
from __future__ import annotations

import copy
import torch
from dacite import from_dict
from dataclasses import dataclass, field, fields
from typing import TYPE_CHECKING, Any, Dict, List, Optional

from swift.infer_engine.protocol import RolloutOutput
from swift.template import Messages
from swift.utils import get_logger, remove_response

if TYPE_CHECKING:
    from swift.infer_engine.protocol import RolloutInferRequest
    from swift.rlhf_trainers.gkd_loss import DataSource

from swift.dataset.preprocessor.core import _pair_keys as StandardKeys

logger = get_logger()

# Multimodal keys that a scheduler may override via ``rollout_infos``.
_MULTIMODAL_KEYS = ('images', 'videos', 'audios')


@dataclass
class OnPolicySample:
    """A single on-policy rollout trajectory (pre-collation, per-sample).

    Lifecycle of one sample::

        1. dataset row        -> messages + extra
        2. rollout            -> response_token_ids / rollout_logprobs / finish_reason
        3. rebuild messages   -> replace_assistant_response_with_ids
        4. encode             -> encoded = template.encode(self)   (per-sample dict)

    The collated model-forward inputs (input_ids[B,T], ...) are produced later
    at batch level from ``[s.encoded for s in samples]`` by
    ``collate_to_grpo_micro_batch`` (returns ``(model_inputs, grpo_batch)``), never on
    the sample.
    """
    # --- standard keys ---
    messages: List[Dict]
    images: List[Any] = field(default_factory=list)
    videos: List[Any] = field(default_factory=list)
    audios: List[Any] = field(default_factory=list)
    tools: Optional[List[Any]] = None
    objects: Dict[str, Any] = field(default_factory=dict)
    extra: Dict[str, Any] = field(default_factory=dict)  # dataset passthrough columns (flattened for reward)

    # --- id ---
    prompt_id: str = ''
    request_id: str = ''

    # --- rollout output ---
    response_token_ids: List[List[int]] = field(default_factory=list)
    response_loss_mask: List[List[int]] = field(default_factory=list)
    rollout_logprobs: List[List[float]] = field(default_factory=list)
    finish_reason: Optional[str] = None
    add_eos: bool = False
    rollout_infos: Dict[str, Any] = field(default_factory=dict)
    routed_experts: Optional[Any] = None  # R3 router replay (per-sample, pre-collation)

    # --- per-sample template.encode output, used for model forward kwargs ---
    encoded: Optional[Dict[str, Any]] = None

    # --- OPSD (On-Policy Self-Distillation), shared by GKD and OPD-RL ---
    teacher_prompt: Optional[Any] = None  # OPSD: dataset ``teacher_prompt`` column (pre-collation)
    teacher_messages: Optional[Messages] = None  # OPSD: messages with teacher_prompt replacing last user

    @property
    def is_truncated(self) -> bool:
        return self.finish_reason == 'length'

    def get_tag(self, tag_key: str = 'dataset') -> Optional[str]:
        """Return the multi-teacher routing tag from ``extra[tag_key]`` (``None`` if unset).

        Single source of truth for where a sample's routing tag lives (``--teacher_tag_key``,
        default ``dataset``); routing keys off exactly this, with no fallback to other columns.
        """
        val = self.extra.get(tag_key)
        return str(val) if val is not None else None

    def build_teacher_view(self) -> bool:
        """Populate the OPSD teacher view from ``teacher_prompt`` + the on-policy response.

        OPSD scores the teacher on its OWN (teacher_prompt + same on-policy response)
        sequence: replace the last user message with ``teacher_prompt`` and keep the
        assistant response. Teacher and student share ``response_token_ids`` (identical
        response tokens, only the prompt differs). Returns ``True`` when an OPSD view
        exists (idempotent) and ``False`` when ``teacher_prompt`` is unset (non-OPSD).
        """
        if self.teacher_messages is not None:
            return True

        if not self.teacher_prompt:
            return False

        messages = [dict(m) for m in self.messages]
        for msg in reversed(messages):
            if msg['role'] == 'user':
                msg['content'] = self.teacher_prompt
                break

        self.teacher_messages = messages
        return True

    def to_teacher_template_dict(self) -> Dict[str, Any]:
        """Reconstruct the teacher-side dict consumed by ``template.encode`` (OPSD).

        Uses ``teacher_messages`` (teacher_prompt-replaced) + the shared
        ``response_token_ids`` (same on-policy response as the student).
        """
        d = self._standard_fields()
        d['messages'] = self.teacher_messages
        chat_template_kwargs = self.extra.get('chat_template_kwargs')
        if chat_template_kwargs is not None:
            d['chat_template_kwargs'] = chat_template_kwargs
        if self.response_token_ids:
            d['response_token_ids'] = self.response_token_ids
        d['add_eos'] = False
        return d

    def _standard_fields(self) -> Dict[str, Any]:
        """Non-empty StandardKeys fields from this sample (replaces _multimodal_columns)."""
        return {k: getattr(self, k) for k in StandardKeys if getattr(self, k)}

    def to_reward_row(self) -> Dict[str, Any]:
        """Build the dict consumed by reward functions.

        Flattens ``extra`` to top level so dataset columns (``solution``,
        ``target``, ...) remain accessible as keyword args after
        ``RowPreprocessor.rows_to_batched``. ``encoded`` is excluded (heavy,
        model-internal). Keeps the legacy column contract intact so existing
        reward functions need no change.
        """
        row: Dict[str, Any] = {
            'messages': self.messages,
            'prompt_id': self.prompt_id,
            'request_id': self.request_id,
            'finish_reason': self.finish_reason,
            'is_truncated': self.is_truncated,
            'rollout_infos': self.rollout_infos,
            'response_token_ids': self.response_token_ids,
        }
        row.update(self._standard_fields())
        row.update(self.extra)
        return row

    @classmethod
    def from_row(cls, row: Dict[str, Any]) -> OnPolicySample:
        """Build a sample from a dataloader / rollout dict row.

        Known keys are mapped to explicit dataclass fields via field-name
        introspection (no hand-maintained key list); every other column
        (dataset passthrough like ``solution`` / ``chat_template_kwargs``) goes
        to ``extra``. ``is_truncated`` is a derived property and is dropped.

        ``row`` values are deep-copied so the sample owns its data: dataloaders
        (e.g. RepeatSampler with steps_per_generation) cache and re-yield the same
        row dict, and the rollout/encode pipeline mutates messages in place
        (remove_response / response injection). Sharing references would corrupt
        the cached dataset rows across micro-steps.
        """
        field_names = {f.name for f in fields(cls)} - {'extra', 'encoded'}
        standard: Dict[str, Any] = {}
        extra: Dict[str, Any] = {}
        for key, value in row.items():
            if key in field_names:
                standard[key] = copy.deepcopy(value)
            elif key == 'is_truncated':
                continue  # derived from finish_reason
            else:
                extra[key] = copy.deepcopy(value)
        return cls(extra=extra, **standard)

    def to_template_dict(self) -> Dict[str, Any]:
        """Reconstruct the dict consumed by ``template.encode``.

        messages + StandardKeys (images/videos/audios/tools/objects) +
        ``chat_template_kwargs`` (the only dataset-passthrough column encode
        consumes — drives enable_thinking / max_pixels / reasoning_effort) +
        add_eos. Other ``extra`` columns (solution/target/...) are reward-only
        and intentionally excluded from encode. Response tokens are already
        injected into ``messages`` via ``replace_assistant_response_with_ids``
        before encoding.
        """
        d = self._standard_fields()
        chat_template_kwargs = self.extra.get('chat_template_kwargs')
        if chat_template_kwargs is not None:
            d['chat_template_kwargs'] = chat_template_kwargs
        d['add_eos'] = self.add_eos
        return d

    def apply_rollout_output(self, *, rollout_output: RolloutOutput) -> None:
        """Merge one rollout output back onto this sample (used post-rollout).

        Single- vs multi-turn is inferred from ``rollout_output`` itself, not
        from external scheduler state: a multi-turn scheduler (colocate *or*
        server) always returns the full ``messages`` history, while single-turn
        leaves ``messages`` empty. ``choice.token_ids`` is the last single-shot
        generation, so it is only a safe fallback in the single-turn case.
        """
        choice = rollout_output.response.choices[0]
        return_by_scheduler = rollout_output.messages is not None

        # messages: multi-turn returns full history; single-turn appends response
        if return_by_scheduler:
            self.messages = rollout_output.messages
        else:
            remove_response(self.messages)
            self.messages.append({'role': 'assistant', 'content': choice.message.content})

        # response token ids: prefer explicit TITO / scheduler output; only fall
        # back to single-shot choice.token_ids in genuine single-turn rollout.
        if rollout_output.response_token_ids:
            self.response_token_ids = rollout_output.response_token_ids
            if rollout_output.response_loss_mask:
                self.response_loss_mask = rollout_output.response_loss_mask
        elif not return_by_scheduler and choice.token_ids:
            self.response_token_ids = choice.token_ids

        # rollout logprobs (for importance sampling); keep nested [turn][token]
        if rollout_output.rollout_logprobs:
            self.rollout_logprobs = rollout_output.rollout_logprobs
        elif choice.logprobs is not None and 'content' in choice.logprobs:
            self.rollout_logprobs = [[item['logprob'] for item in choice.logprobs['content']]]

        self.finish_reason = choice.finish_reason
        self.add_eos = False

        # R3 router replay: transfer routed_experts from the rollout choice
        routed_experts = getattr(choice, 'routed_experts', None)
        if routed_experts is not None:
            self.routed_experts = routed_experts

        # rollout_infos may carry scheduler-overridden multi-modal data
        if rollout_output.rollout_infos:
            self.rollout_infos = rollout_output.rollout_infos
            for key in _MULTIMODAL_KEYS:
                if key in rollout_output.rollout_infos:
                    setattr(self, key, rollout_output.rollout_infos[key])

    def to_infer_request(self, include_extra: bool = False) -> RolloutInferRequest:
        """Build the ``RolloutInferRequest`` consumed by the rollout engine.

        Maps messages + multimodal/standard columns (images/videos/audios/
        tools/objects) + ``uuid`` (defaults to ``request_id``). Images given as
        ``{'bytes': ...}`` / ``{'path': ...}`` dicts are normalized to base64 /
        path strings. ``tools`` given as a JSON string is parsed.

        When ``include_extra`` is True, dataset passthrough columns (``extra``)
        are forwarded via ``data_dict`` (used for server mode / multi-turn).
        """
        import base64
        import json

        from swift.infer_engine.protocol import RolloutInferRequest

        def _process_image_data(image_data):
            if isinstance(image_data, dict):
                if image_data.get('bytes'):
                    return base64.b64encode(image_data['bytes']).decode('utf-8')
                if image_data.get('path'):
                    return image_data['path']
            return image_data

        request_data: Dict[str, Any] = {'uuid': self.request_id}
        request_data.update(self._standard_fields())

        chat_template_kwargs = self.extra.get('chat_template_kwargs')
        if chat_template_kwargs:
            request_data['chat_template_kwargs'] = chat_template_kwargs

        if request_data.get('images'):
            imgs = request_data['images']
            if not isinstance(imgs, list):
                imgs = [imgs]
            request_data['images'] = [_process_image_data(img) for img in imgs]

        if isinstance(request_data.get('tools'), str):
            try:
                request_data['tools'] = json.loads(request_data['tools'])
            except json.JSONDecodeError:
                pass

        if include_extra and self.extra:
            base_data_dict = self.extra.get('data_dict')
            if base_data_dict is not None and not isinstance(base_data_dict, dict):
                raise ValueError('data_dict exists but is not a dictionary')
            extra_data = {k: v for k, v in self.extra.items() if k != 'data_dict' and v is not None}
            request_data['data_dict'] = {**extra_data, **(base_data_dict or {})}

        return from_dict(RolloutInferRequest, request_data)


@dataclass
class GRPOSample(OnPolicySample):
    """On-policy sample with GRPO reward/advantage signals."""
    rewards: Optional[List[Optional[float]]] = None  # optional mirror; main path uses rewards_per_func tensor
    advantages: Optional[torch.Tensor] = None  # filled after _compute_advantages (0-dim tensor per sample)


@dataclass
class GRPOBatch:
    """Batch data for GRPO loss computation (post-collation, batch-level).

    1. ``completion_mask``, ``truncated_mask``, ``seq_lengths`` — derived from
       collated ``labels`` right after ``data_collator``.
    2. ``old_per_token_logps``, ``ref_per_token_logps`` — computed via
       model/ref forward on the collated batch.
    3. ``advantages`` — computed from gathered rewards.
    4. ``rollout_per_token_logps``, ``num_items_in_batch`` — optional,
       filled when rollout IS / DAPO is enabled.
    """
    completion_mask: torch.Tensor  # [B, T]
    truncated_mask: torch.Tensor  # [B]
    seq_lengths: torch.Tensor  # [B] or [B+n] for padding_free

    old_per_token_logps: Optional[torch.Tensor] = None  # [B, T]
    ref_per_token_logps: Optional[torch.Tensor] = None  # [B, T]
    rollout_per_token_logps: Optional[torch.Tensor] = None  # [B, T]
    teacher_per_token_logps: Optional[torch.Tensor] = None  # [B, T], OPD-RL teacher logp on sampled tokens
    advantages: Optional[torch.Tensor] = None  # [B, T] per-token (base broadcast minus per-token teacher KL)
    num_items_in_batch: Optional[torch.Tensor] = None  # scalar
    logits_to_keep: Optional[int] = None

    def to_device(self, device) -> 'GRPOBatch':
        """Move all tensor fields to ``device`` in place (Ray: collated on the CPU
        driver, moved to the GPU worker before forward). Returns self."""
        for f in fields(self):
            v = getattr(self, f.name)
            if isinstance(v, torch.Tensor):
                setattr(self, f.name, v.to(device))
        return self


@dataclass
class GKDSample(OnPolicySample):
    pass


@dataclass
class GKDBatch:
    """Batch-level GKD signals (post-collation), symmetric to :class:`GRPOBatch`.

    - ``data_source``: STUDENT / TEACHER / DATASET for this micro-batch.
    - ``teacher_topk_logprobs`` / ``teacher_topk_indices``: assembled teacher
      top-k (teacher-API mode), batch tensors aligned to student tokens.
    """
    data_source: 'DataSource'
    teacher_topk_logprobs: Optional[torch.Tensor] = None
    teacher_topk_indices: Optional[torch.Tensor] = None
