import asyncio
import json
import random
from concurrent.futures import ALL_COMPLETED, ThreadPoolExecutor, wait
from copy import deepcopy
from tree_rollout import (DataSampleTree, DivergenceStrategyMapping, FinishedReason, SampleStatus,
                          _increment_tree_idx_depth, _repeat_list_interleave, extract_last_boxed)
from typing import Any, Dict, List, Optional, Union

from swift.infer_engine import RequestConfig
from swift.infer_engine.protocol import ChatCompletionResponse, RolloutInferRequest, RolloutOutput
from swift.rewards import MultiTurnScheduler, multi_turns


class TreeRolloutScheduler(MultiTurnScheduler):
    """
    Base class for multi-turn tree-rollout scheduling.

    Provides default implementation for multi-turn conversation management.

    CUSTOMIZATION:
        Implement the required `step()` method and optionally override `check_finished()`
        - Uses TreeRolloutScheduler's run() method infrastructure
        - Only need to implement turn transition logic in step()
        - Optionally customize termination conditions

    Attributes:
        max_tree_width (int):
            For GRPO, it must be equal to num_generations.
        max_tree_depth (int):
            Controls the maximum number of reasoning turns for a single prompt.
        root_divergence (int):
            Number of branches generated in the first-round inference at the root node.
        max_divergence (int):
            Maximum number of branches allowed for each node.
        divergence_strategy (str):
            Strategy for selecting branch nodes; defaults to logprobs.
    """

    def __init__(self, infer_engine=None, max_turns=None, *args, **kwargs):
        super().__init__(infer_engine, max_turns, *args, **kwargs)
        self.max_tree_width = 8
        self.max_tree_depth = max_turns | 6
        self.max_divergence = 2
        self.divergence_strategy = 'logprobs'
        self.root_divergence = 1

        self.executor = ThreadPoolExecutor(max_workers=self.max_tree_width)

    async def async_infer(self,
                          infer_requests: List[Union['RolloutInferRequest', Dict[str, Any]]],
                          request_config: 'RequestConfig',
                          *,
                          use_tqdm: Optional[bool] = None,
                          **kwargs) -> List['RolloutOutput']:
        # dedup_requests_by_messages
        processed_request = []
        seen = set()
        uuids = []

        for item in infer_requests:
            if isinstance(item, dict):
                req = RolloutInferRequest(**item)
            else:
                req = item

            msg_key = json.dumps(req.messages, sort_keys=True)
            uuids.append(req.uuid)

            if msg_key not in seen:
                seen.add(msg_key)
                processed_request.append(req)

        request_config.logprobs = True

        outputs = await super().async_infer(processed_request, request_config, use_tqdm=use_tqdm, **kwargs)

        assert len(outputs) == len(uuids), '[Tree Rollout] Please check the max_tree_width is equal to num_generations.'

        for idx, output in enumerate(outputs):
            output.response.id = uuids[idx]

        return outputs

    async def run(self, infer_request: Union[List[RolloutInferRequest], RolloutInferRequest],
                  request_config: 'RequestConfig', **kwargs) -> List['RolloutOutput']:
        if isinstance(infer_request, RolloutInferRequest):
            infer_request = [infer_request]
        else:
            infer_request = list(infer_request)

        request_config.logprobs = True

        finished_rollout_by_root: Dict[int, List[RolloutOutput]] = {i: [] for i in range(len(infer_request))}
        finished_samples: Dict[int, List[DataSampleTree]] = {i: [] for i in range(len(infer_request))}

        samples_to_infer = []

        for root_idx in range(len(infer_request)):
            samples_to_infer.append(
                DataSampleTree(
                    tree_idx=str(root_idx),
                    request_id=infer_request[root_idx].uuid,
                    messages=infer_request[root_idx].messages,
                    status=SampleStatus.TO_INFER))

        # first step
        next_infer_step = 1
        samples_to_infer = _repeat_list_interleave(samples_to_infer, self.root_divergence)
        samples_to_infer = _increment_tree_idx_depth(samples_to_infer, next_infer_step)

        while len(samples_to_infer) > 0:
            # resolve the error: Request id xxx already running
            vllm_inputs = [
                RolloutInferRequest(messages=sample.messages, uuid=f'{sample.request_id}-{sample.tree_idx}')
                for sample in samples_to_infer
            ]

            # Get model response
            tasks = [self.infer_engine.infer_async(request, request_config, **kwargs) for request in vllm_inputs]
            outputs: List[ChatCompletionResponse] = await asyncio.gather(*tasks)

            assert len(vllm_inputs) == len(
                outputs), f'outputs length {len(outputs)} != inputs length {len(vllm_inputs)}'

            samples_last_step = deepcopy(samples_to_infer)
            samples_to_infer = []

            for idx, (sample, output) in enumerate(zip(samples_last_step, outputs)):
                assert len(output.choices) == 1, 'vllm should only generate one output'
                self.check_finished(sample, output)

                # bind the output and request
                output.id = sample.request_id
                choice = output.choices[0]
                child_sample = deepcopy(sample)
                child_sample.extend_response(choice)

                if child_sample.status == SampleStatus.FINISHED:
                    finished_samples[child_sample.root_node].append(child_sample)
                    finished_rollout_by_root[child_sample.root_node].append(
                        RolloutOutput(
                            response=output,
                            messages=deepcopy(child_sample.messages),
                            response_token_ids=deepcopy(child_sample.all_response_ids),
                            # If we use intermediate reasoning results when computing the reward,
                            # but loss_mask is not explicitly set,
                            # only the loss of the final round of reasoning will be computed.
                            response_loss_mask=[[1] * len(response_ids)
                                                for response_ids in child_sample.all_response_ids],
                            rollout_infos={'num_turns': next_infer_step},
                        ))
                else:
                    samples_to_infer.append(child_sample)

            # if we have budget, do divergence
            if len(samples_to_infer) > 0 and self.max_divergence > 1:
                for root_idx in finished_samples.keys():
                    root_to_infer_samples = [sample for sample in samples_to_infer if sample.root_node == root_idx]
                    root_finished_samples = finished_samples[root_idx]

                    budget = self.max_tree_width - len(root_finished_samples) - len(root_to_infer_samples)

                    if budget > 0 and len(root_to_infer_samples) > 0:
                        divergence_executor = DivergenceStrategyMapping[self.divergence_strategy]
                        if not divergence_executor:
                            raise ValueError(
                                f"[Tree Rollout] The divergence strategy: {self.divergence_strategy} doesn't exist.")

                        divergence_samples = divergence_executor.apply(root_idx, root_to_infer_samples, budget,
                                                                       self.max_divergence - 1)
                        samples_to_infer.extend(divergence_samples)

            # before end loop, if finished_count < max_tree_width, rollback
            if len(samples_to_infer) == 0 and any(count < self.max_tree_width
                                                  for count in [len(value) for value in finished_samples.values()]):
                samples_to_infer = self.roll_back_to_divergence(finished_samples)

            # tools call etc
            futures = [self.executor.submit(self.step, sample) for sample in samples_to_infer]
            wait(futures, return_when=ALL_COMPLETED)

            next_infer_step += 1
            samples_to_infer = _increment_tree_idx_depth(samples_to_infer, next_infer_step)

        # flatten finished outputs
        return [traj for lst in finished_rollout_by_root.values() for traj in lst]

    def step(self, sample: DataSampleTree, **kwargs):
        """
        You need to rewrite or modify this method to customize the next round of prompts, such as tools call.
        """

        # Special handling has already been done in the rollback.
        if sample.status == SampleStatus.ROLLBACK:
            sample.status = SampleStatus.TO_INFER
            return
        elif sample.status == SampleStatus.FINISH_NEXT_INFER:
            prompt = 'In this round of responses, you must generate an answer.'
        else:
            prompt = 'The answer is not correct, It seems You made a mistake, you need to recheck very carefully.'

        sample.messages.append({'role': 'user', 'content': prompt})

    def check_finished(self, sample: DataSampleTree, output: ChatCompletionResponse, **kwargs) -> bool:
        """
        Rewrite this method to add custom check logic
        """

        boxed_answer = extract_last_boxed(output.choices[0].message.content)

        if boxed_answer is not None:
            sample.status = SampleStatus.FINISHED
            sample.finished_reason = FinishedReason.ANSWER

        elif sample.status == SampleStatus.FINISH_NEXT_INFER:
            sample.status = SampleStatus.FINISHED
            sample.finished_reason = FinishedReason.MAX_INFER_STEP

        elif sample.depth >= self.max_tree_depth - 1:
            sample.status = SampleStatus.FINISH_NEXT_INFER

        return sample.status == SampleStatus.FINISHED

    def roll_back_to_divergence(
        self,
        finished_samples: Dict[int, List[DataSampleTree]],
    ) -> List[DataSampleTree]:
        """
        All nodes have completed inference, but there is still budget available, rollback.
        """

        sample_to_infer = []
        for root_idx, sample_list in finished_samples.items():
            if len(sample_list) >= self.max_tree_width:
                continue

            diff_count = self.max_tree_width - len(sample_list)
            result = random.sample(sample_list, min(diff_count, len(sample_list)))

            result_copy = deepcopy(result)

            # Randomly rollback several inference iterations; The rollback strategy can be optimized subsequently.
            for sample in result_copy:
                sample.status = SampleStatus.ROLLBACK
                truncate_len = sample.response_num
                sample.response_truncate(random.randint(1, truncate_len))

            sample_to_infer.extend(result_copy)

        return sample_to_infer


multi_turns['tree_rollout_scheduler'] = TreeRolloutScheduler
