from typing import Any, Dict, List, Optional, Union

from reward_kit.agent.models import StepData  # Internal StepData model

# Assuming models are structured as planned
from reward_kit.models import EvaluateResult  # Extended EvaluateResult

# Placeholder for actual Message type if needed for type hinting complex observation_data
# from reward_kit.models import Message


class RLDataAligner:
    """
    Component responsible for aligning outputs from user reward functions
    (EvaluateResult containing scores and/or base_rewards per step)
    with the system's internal StepData representation. This prepares
    the data for subsequent GiGPO (or other RL algorithm) advantage calculations.
    """

    def align_data_for_rl_processing(
        self,
        current_eval_result: EvaluateResult,
        current_step_data_list: List[StepData],
        rollout_id: str,  # For logging or if needed
    ) -> List[StepData]:
        """
        Aligns the EvaluateResult (from user's reward function) with the
        internal list of StepData for a single rollout.

        Populates `StepData.base_reward` from `EvaluateResult.step_outputs.base_reward`.
        Associates `EvaluateResult.score` with the rollout for GiGPO A_E calculation.
        (Association of final_score might happen by returning it alongside, or
         by the caller managing it). For now, this function focuses on base_reward.

        Args:
            current_eval_result: The EvaluateResult from the user's reward function for this rollout.
            current_step_data_list: The list of StepData objects collected by RLRolloutWorker.
            rollout_id: Identifier for the current rollout.

        Returns:
            The list of StepData objects, with `base_reward` populated.
            The `final_score` from current_eval_result should be handled by the caller
            for GiGPO A_E calculation.
        """

        # Store final_score (for GiGPO A_E) - The caller will handle this.
        # This function's primary job is to populate base_rewards in StepData.

        if current_eval_result.step_outputs:
            # Create a dictionary for quick lookup of user-defined step rewards
            user_step_rewards_map: Dict[Union[int, str], float] = {
                step_out.step_index: step_out.base_reward
                for step_out in current_eval_result.step_outputs
            }

            for s_data in current_step_data_list:
                # --- Critical Mapping Logic ---
                # Strategy: Use 'assistant_turn_index' stored in StepData.step_info
                # by RLRolloutWorker. User's StepOutput.step_index should match this.
                # This assumes RLRolloutWorker adds this info.
                user_defined_step_idx = s_data.step_info.get("assistant_turn_index")

                if user_defined_step_idx is not None:
                    if user_defined_step_idx in user_step_rewards_map:
                        s_data.base_reward = user_step_rewards_map[
                            user_defined_step_idx
                        ]
                    else:
                        # No base reward provided by user for this specific system step.
                        # s_data.base_reward remains None (or could be a default).
                        pass
                else:
                    # RLRolloutWorker did not provide 'assistant_turn_index' for this StepData,
                    # or the mapping key in step_info is different.
                    # This indicates a potential issue in RLRolloutWorker or mapping strategy.
                    pass
        else:
            # No step_outputs provided by the user. Base rewards will remain None.
            pass

        return current_step_data_list

    # TODO (Future): Consider a batch version if performance becomes an issue,
    # but the core logic per rollout remains the same.
    # def align_batch_data_for_rl_processing(...)
