"""
LightningRL: A reinforcement learning component for preparing rollout datasets for training.

This module implements an agentic reinforcement learning algorithm where:
- Trajectories are decomposed into transitions
- Transitions for the same task are grouped for advantage estimation
- Rollout data is processed and prepared for the trainer

The component handles:
- Reconstructing sample information from completed rollouts
- Processing triplets (prompt/response pairs) with proper tokenization
- Padding and truncation of sequences
- Tensor creation and batching for training
- Reward assignment and token-level score computation
"""

import logging
from typing import Dict, Any, List, Tuple
import numpy as np
import torch
from tensordict import TensorDict
from verl import DataProto

from .logging import configure_logger

def get_left_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_token_id: int):
    """
    Left-pad (or truncate) a sequence of token IDs to a fixed length,
    and build the corresponding attention mask.

    Args:
        ids:             the original list of token IDs.
        max_length:      desired total length after padding/truncation.
        pad_token_id:    ID to use for padding.

    Returns:
        padded_ids (any):      list of length == max_length.
        attention_mask (any):  list of same length: 1 for non-pad tokens, 0 for pads.
    """
    seq_len = len(ids)

    if seq_len >= max_length:
        # too long → truncate from the left, keep the last max_length tokens
        trimmed = ids[-max_length:]
        attention_mask = [1] * max_length
        return trimmed, attention_mask

    # too short → pad on the left
    pad_len = max_length - seq_len
    padded_ids = [pad_token_id] * pad_len + ids
    attention_mask = [0] * pad_len + [1] * seq_len
    return padded_ids, attention_mask


def get_right_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_token_id: int):
    """
    Right-pad (or truncate) a sequence of token IDs to a fixed length,
    and build the corresponding attention mask.

    Args:
        ids:            the original list of token IDs.
        max_length:     desired total length after padding/truncation.
        pad_token_id:   ID to use for padding.

    Returns:
        padded_ids (any):     list of length == max_length.
        attention_mask (any): list of same length: 1 for non-pad tokens, 0 for pads.
    """
    seq_len = len(ids)

    if seq_len >= max_length:
        # too long → truncate to the first max_length tokens
        trimmed = ids[:max_length]
        attention_mask = [1] * max_length
        return trimmed, attention_mask

    # too short → pad on the right
    pad_len = max_length - seq_len
    padded_ids = ids + [pad_token_id] * pad_len
    attention_mask = [1] * seq_len + [0] * pad_len
    return padded_ids, attention_mask


logger = configure_logger(level=logging.INFO, name=__name__)


class LightningRL:
    """
    LightningRL component for preparing rollout datasets for reinforcement learning training.
    
    This component implements the core agentic RL algorithm where trajectories are decomposed
    into transitions and grouped by task for advantage estimation.
    """
    
    def __init__(self, pad_token_id: int, reward_fillna_value: float = 0.0):
        """
        Initialize the LightningRL component.
        
        Args:
            pad_token_id: Token ID used for padding sequences
            reward_fillna_value: Default value to use when reward is None/null
        """
        self.pad_token_id = pad_token_id
        self.reward_fillna_value = reward_fillna_value
    
    def _fillna_reward(self, rollout):
        """Fill null/None reward values with the configured fillna value."""
        # Handle both old CalcAgentResult objects and new Rollout objects
        reward = getattr(rollout, 'final_reward', None)
        if reward is None:
            reward = getattr(rollout, 'reward', None)
        if reward is None:
            return self.reward_fillna_value

        return reward
    
    def generate_training_batch(
        self, 
        completed_rollouts: Dict[str, Any], 
        task_id_to_original_sample: Dict[str, Any],
        max_prompt_length: int, 
        max_response_length: int, 
        device: torch.device
    ) -> Tuple[DataProto, Dict[str, Any]]:
        """
        Processes completed rollouts to generate a training data batch.

        This function implements the core agentic RL algorithm, reconstructing logic
        for decomposing trajectories into transitions and grouping them by task for
        advantage estimation. It handles padding, truncation, and tensor creation 
        for the training loop.
        
        Args:
            completed_rollouts: Dictionary mapping rollout_id to rollout objects
            task_id_to_original_sample: Dictionary mapping task_id to original sample data
            max_prompt_length: Maximum length for prompt sequences
            max_response_length: Maximum length for response sequences  
            device: PyTorch device to place tensors on
            
        Returns:
            Tuple of (DataProto containing training batch, metrics dictionary)
        """
        logger.info(f"Starting generate_training_batch with {len(completed_rollouts)} completed rollouts")

        # 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts
        finished_id_to_sample_info = {}
        rollouts_without_triplets = 0
        rollouts_with_empty_triplets = 0
        total_triplets_processed = 0
        
        for rollout_id, rollout in completed_rollouts.items():
            original_sample = task_id_to_original_sample[rollout_id]

            if not hasattr(rollout, 'triplets'):
                logger.warning(f"Rollout {rollout_id} has no 'triplets' attribute")

                rollouts_without_triplets += 1
                continue
                
            if not rollout.triplets:
                logger.warning(f"Rollout {rollout_id} has empty triplets list")
                rollouts_with_empty_triplets += 1
                continue

            # The client should report triplets that contain prompt_ids and response_ids.
            # Example triplet.prompt: {"token_ids": [...]}
            # Example triplet.response: {"token_ids": [...]}
            trace_list = []
            reward = None
            for i, t in enumerate(rollout.triplets):
                prompt_token_ids = t.prompt.get("token_ids", []) if isinstance(t.prompt, dict) else []
                response_token_ids = t.response.get("token_ids", []) if isinstance(t.response, dict) else []
                reward = t.reward
                trace_list.append({
                    "prompt_ids": prompt_token_ids, 
                    "response_ids": response_token_ids
                })
                total_triplets_processed += 1

            final_reward = self._fillna_reward(rollout) if reward is None else reward

            info = {
                "reward": final_reward,
                "trace_list": trace_list,
                "data_id": original_sample["data_id"],
            }
            finished_id_to_sample_info[rollout_id] = info
        
        logger.info(f"Rollout processing summary:")
        logger.info(f"  - Total rollouts: {len(completed_rollouts)}")
        logger.info(f"  - Rollouts without triplets attribute: {rollouts_without_triplets}")
        logger.info(f"  - Rollouts with empty triplets: {rollouts_with_empty_triplets}")
        logger.info(f"  - Valid rollouts processed: {len(finished_id_to_sample_info)}")
        logger.info(f"  - Total triplets processed: {total_triplets_processed}")
        
        # --- Data processing and tensor creation logic ---
        # Get all the reported data.
        # prompt_ids are left-padded.
        # response_ids are right-padded.
        # They are concatenated in the middle.
        # Discard handling:
        #   - Those exceeding max_prompt_length will be marked for discard, but not
        #     discarded here. They are only truncated and marked, to be discarded later.
        #     This is for the correctness of the advantage calculation.
        #   - The discard for the PPO mini-batch should also be handled this way.
        input_ids_list, input_attention_mask_list = [], []
        response_ids_list, response_attention_mask_list = [], []
        reward_list, data_id_list, rollout_id_list, turn_index_list, is_drop_list = [], [], [], [], []
        n_trunc_sample_because_of_response = 0
        empty_prompt_count = 0
        empty_response_count = 0
        
        for rollout_id, sample_info in finished_id_to_sample_info.items():
            for turn_index, trace in enumerate(sample_info["trace_list"]):
                reward_list.append(sample_info["reward"])
                prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]
                
                if not prompt_ids:
                    empty_prompt_count += 1
                    logger.warning(f"Empty prompt_ids in rollout {rollout_id}, turn {turn_index}")
                    
                if not response_ids:
                    empty_response_count += 1
                    logger.warning(f"Empty response_ids in rollout {rollout_id}, turn {turn_index}")

                # Mark samples with prompts exceeding max_prompt_length to be dropped later
                if len(prompt_ids) > max_prompt_length:
                    prompt_ids = prompt_ids[:max_prompt_length]
                    is_drop_list.append(True)
                else:
                    is_drop_list.append(False)

                # Truncate responses that exceed max_response_length
                if len(response_ids) > max_response_length:
                    response_ids = response_ids[:max_response_length]
                    n_trunc_sample_because_of_response += 1

                # Pad prompts to the left and responses to the right
                one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
                    prompt_ids, max_prompt_length, self.pad_token_id
                )
                one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
                    response_ids, max_response_length, self.pad_token_id
                )

                input_ids_list.append(one_input_ids)
                input_attention_mask_list.append(one_input_attention_mask)
                response_ids_list.append(one_response_ids)
                response_attention_mask_list.append(one_response_attention_mask)
                data_id_list.append(sample_info["data_id"])
                rollout_id_list.append(rollout_id)
                turn_index_list.append(turn_index)

        n_transition = len(input_ids_list)
        
        logger.info(f"Final processing summary:")
        logger.info(f"  - Empty prompts: {empty_prompt_count}")
        logger.info(f"  - Empty responses: {empty_response_count}")
        logger.info(f"  - Truncated responses: {n_trunc_sample_because_of_response}")
        logger.info(f"  - Total transitions: {n_transition}")
        
        if n_transition == 0:
            logger.warning("No valid transitions found! Returning empty batch.")
            # Return empty batch
            batch = TensorDict({
                "prompts": torch.empty((0, max_prompt_length), dtype=torch.long, device=device),
                "responses": torch.empty((0, max_response_length), dtype=torch.long, device=device),
                "input_ids": torch.empty((0, max_prompt_length + max_response_length), dtype=torch.long, device=device),
                "attention_mask": torch.empty((0, max_prompt_length + max_response_length), dtype=torch.long, device=device),
                "position_ids": torch.empty((0, max_prompt_length + max_response_length), dtype=torch.long, device=device),
                "is_drop_mask": torch.zeros((0,), dtype=torch.bool, device=device),
                "token_level_scores": torch.zeros((0, max_response_length), dtype=torch.float, device=device),
            }, batch_size=0)
            
            data_proto = DataProto(batch=batch)
            data_proto.non_tensor_batch["data_id_list"] = np.array([])
            data_proto.non_tensor_batch["rollout_id_list"] = np.array([])
            data_proto.non_tensor_batch["turn_index_list"] = np.array([])
            
            return data_proto, {"agent_mode/n_sample_to_train": 0}
        
        # Convert lists to tensors
        batch_input_ids = torch.LongTensor(input_ids_list).to(device)
        input_attention_mask = torch.LongTensor(input_attention_mask_list).to(device)
        batch_response_ids = torch.LongTensor(response_ids_list).to(device)
        response_attention_mask = torch.LongTensor(response_attention_mask_list).to(device)

        # Concatenate prompts and responses to form the full sequence
        batch_seq = torch.cat([batch_input_ids, batch_response_ids], dim=-1)
        attention_mask = torch.cat([input_attention_mask, response_attention_mask], dim=-1)
        position_ids = torch.clamp(torch.cumsum(attention_mask, dim=-1) - 1, min=0)
        is_drop_mask = torch.BoolTensor(is_drop_list).to(device)
        scores = torch.tensor(reward_list, dtype=torch.bfloat16).to(device)

        # Create token-level scores by placing the final reward at the last token position
        token_level_scores = torch.zeros_like(attention_mask, dtype=scores.dtype)
        # At the eos_mask_idx position of each sample, fill in the corresponding scores.
        # torch.arange(n_transition) generates [0,1,2,...,bsz-1] as indices for the batch dimension.
        eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1)  # (bsz,)
        token_level_scores[torch.arange(n_transition), eos_mask_idx] = scores
        # Only take the last response_length part of the sequence to get the token-level scores for the model's response part.
        token_level_scores = token_level_scores[:, -max_response_length:]

        # Form the final batch using TensorDict
        batch = TensorDict(
            {
                "prompts": batch_input_ids,
                "responses": batch_response_ids,
                "input_ids": batch_seq,  # here input_ids become the whole sentences
                "attention_mask": attention_mask,
                "position_ids": position_ids,
                "is_drop_mask": is_drop_mask,
                "token_level_scores": token_level_scores.contiguous(),
            },
            batch_size=n_transition,
        )
        data_proto = DataProto(batch=batch)

        data_metrics = {
            "agent_mode/n_trunc_sample_because_of_response": n_trunc_sample_because_of_response,
            "agent_mode/n_sample_to_train": n_transition,
        }

        # Add non-tensor data for advantage calculation and logging
        data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)
        data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)
        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)

        return data_proto, data_metrics
