# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import partial
from collections import defaultdict
import threading
import time

import torch

from verl import DataProto
from verl.utils.reward_score import _default_compute_score


def compute_score_with_timeout(evaluation_func, data_source, solution_str, ground_truth, extra_info, timeout=300.0):
    """Compute score with timeout handling using threading."""
    result = [None]
    exception = [None]
    
    def target():
        try:
            result[0] = evaluation_func(data_source, solution_str, ground_truth, extra_info)
        except Exception as e:
            exception[0] = e
    
    thread = threading.Thread(target=target)
    thread.start()
    thread.join(timeout)
    
    if thread.is_alive():
        print(f"Timeout occurred for solution: {solution_str[:50]}...")
        # Note: We can't forcefully kill the thread, but we can return a default value
        return -2.0
    
    if exception[0] is not None:
        print(f"Error processing solution: {solution_str[:50]}..., Error: {exception[0]}")
        return -2.0
    
    return result[0] if result[0] is not None else -2.0


def parallel_compute_score_dapo(evaluation_func, data_sources, solution_strs, ground_truths, extra_infos, num_workers=64):
    """Parallel compute score for DAPO using ThreadPoolExecutor."""
    scores = []
    
    if extra_infos is None:
        extra_infos = [None] * len(data_sources)
    
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        # Submit all tasks
        future_to_index = {}
        for i, (data_source, solution_str, ground_truth, extra_info) in enumerate(zip(data_sources, solution_strs, ground_truths, extra_infos)):
            future = executor.submit(
                compute_score_with_timeout,
                evaluation_func,
                data_source,
                solution_str,
                ground_truth,
                extra_info,
                timeout=300.0
            )
            future_to_index[future] = i
        
        # Initialize results list
        results = [None] * len(data_sources)
        
        # Collect results as they complete
        for future in as_completed(future_to_index):
            index = future_to_index[future]
            try:
                result = future.result()
                results[index] = result
            except Exception as e:
                print(f"Error in thread execution: {e}")
                results[index] = -2.0
    
    # Process results
    for result in results:
        if result is None:
            scores.append(-2.0)
        elif isinstance(result, dict):
            scores.append(result)
        else:
            scores.append(float(result))
    
    return scores


class DAPORewardManager:
    """The reward manager."""

    def __init__(
        self,
        tokenizer,
        num_examine,
        compute_score=None,
        reward_fn_key="data_source",
        max_resp_len=None,
        overlong_buffer_cfg=None,
    ) -> None:
        self.tokenizer = tokenizer
        self.num_examine = num_examine  # the number of batches of decoded responses to print to the console
        self.compute_score = compute_score or _default_compute_score
        self.reward_fn_key = reward_fn_key
        self.overlong_buffer_cfg = overlong_buffer_cfg
        self.max_resp_len = max_resp_len

        if self.overlong_buffer_cfg is not None:
            assert self.max_resp_len is not None, f"max_resp_len must be provided if {overlong_buffer_cfg=}, but got None"

    def __call__(self, data: DataProto, return_dict: bool = False):
        """We will expand this function gradually based on the available datasets"""

        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn
        if "rm_scores" in data.batch.keys():
            if return_dict:
                return {"reward_tensor": data.batch["rm_scores"]}
            else:
                return data.batch["rm_scores"]

        reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
        reward_extra_info = defaultdict(list)

        already_print_data_sources = {}

        # Collect all data for parallel processing
        data_sources = []
        solution_strs = []
        ground_truths = []
        extra_infos = []
        valid_response_lengths = []
        
        # Preprocessing: collect all data items
        for i in range(len(data)):
            data_item = data[i]  # DataProtoItem

            prompt_ids = data_item.batch["prompts"]
            prompt_length = prompt_ids.shape[-1]

            valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum()
            valid_prompt_ids = prompt_ids[-valid_prompt_length:]

            response_ids = data_item.batch["responses"]
            valid_response_length = data_item.batch["attention_mask"][prompt_length:].sum()
            valid_response_ids = response_ids[:valid_response_length]

            # decode
            prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True)
            response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)
            eos_token = self.tokenizer.eos_token
            if response_str.endswith(eos_token):
                response_str = response_str[: -len(eos_token)]

            ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"]
            data_source = data_item.non_tensor_batch[self.reward_fn_key]
            extra_info = data_item.non_tensor_batch.get("extra_info", None)

            # Collect data for parallel processing
            data_sources.append(data_source)
            solution_strs.append(response_str)
            ground_truths.append(ground_truth)
            extra_infos.append(extra_info)
            valid_response_lengths.append(valid_response_length)

        # Parallel compute scores using ThreadPoolExecutor (fixed 64 workers)
        try:
            results = parallel_compute_score_dapo(
                self.compute_score,
                data_sources,
                solution_strs,
                ground_truths,
                extra_infos,
                num_workers=64,
            )
        except Exception as e:
            print(f"Unexpected error in parallel reward computing. Setting all as -2.0: {e}")
            results = [-2.0 for _ in range(len(data_sources))]

        # Post-processing: apply results and handle overlong buffer
        for i in range(len(data)):
            result = results[i]
            
            # Handle different result types
            score: float
            if isinstance(result, dict):
                score = result["score"]
                # Store the information including original reward
                for key, value in result.items():
                    reward_extra_info[key].append(value)
            else:
                score = result

            reward = score

            # Apply overlong buffer penalty if enabled
            if self.overlong_buffer_cfg.enable:
                overlong_buffer_len = self.overlong_buffer_cfg.len
                expected_len = self.max_resp_len - overlong_buffer_len
                exceed_len = valid_response_lengths[i] - expected_len
                overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor
                overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0)
                reward += overlong_reward
                if self.overlong_buffer_cfg.log:
                    reward_extra_info["overlong_reward"].append(overlong_reward)
                    reward_extra_info["overlong"].append(overlong_reward < 0)

            reward_tensor[i, valid_response_lengths[i] - 1] = reward

            # Print examples for debugging
            data_source = data_sources[i]
            if data_source not in already_print_data_sources:
                already_print_data_sources[data_source] = 0

            if already_print_data_sources[data_source] < self.num_examine:
                already_print_data_sources[data_source] += 1
                
                # Reconstruct prompt for printing
                data_item = data[i]
                prompt_ids = data_item.batch["prompts"]
                prompt_length = prompt_ids.shape[-1]
                valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum()
                valid_prompt_ids = prompt_ids[-valid_prompt_length:]
                prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True)
                
                print("[prompt]", prompt_str)
                print("[response]", solution_strs[i])
                print("[ground_truth]", ground_truths[i])
                if isinstance(result, dict):
                    for key, value in result.items():
                        print(f"[{key}]", value)
                else:
                    print("[score]", score)

        if return_dict:
            return {
                "reward_tensor": reward_tensor,
                "reward_extra_info": reward_extra_info,
            }
        else:
            return reward_tensor
