# -*- coding: utf-8 -*-
"""Custom vLLM Worker."""
import ray
import torch
import torch.distributed

from trinity.common.models.vllm_patch.worker_patch import patch_vllm_prompt_logprobs
from trinity.manager.synchronizer import Synchronizer
from trinity.utils.distributed import init_process_group
from trinity.utils.log import get_logger


class WorkerExtension:
    def apply_patches(self):
        """Apply necessary patches to vLLM."""
        from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader

        patch_vllm_moe_model_weight_loader(self.model_runner.model)
        patch_vllm_prompt_logprobs(self.model_runner)

    def init_process_group(
        self,
        master_address: str,
        master_port: int,
        rank_offset: int,
        world_size: int,
        group_name: str,
        backend: str = "nccl",
        timeout: int = 1200,
        state_dict_meta: list = None,
        explorer_name: str = None,
        namespace: str = None,
    ):
        """Init torch process group for model weights update"""
        rank = torch.distributed.get_rank()
        self.logger = get_logger(f"vllm_worker_{rank}")

        assert torch.distributed.is_initialized(), "default torch process group must be initialized"
        assert group_name != "", "group name must not be empty"
        self._state_dict_meta = state_dict_meta
        self._weight_update_rank = rank + rank_offset
        self.logger.info(
            f"vLLM starting init_process_group:\n"
            f"  > address={master_address}:{master_port}\n"
            f"  > rank={rank}\n"
            f"  > rank_offset={rank_offset}\n"
            f"  > world_size={world_size}"
        )
        self._model_update_group = init_process_group(
            host=master_address,
            port=master_port,
            group_name=group_name,
            backend=backend,
            timeout=timeout,
            world_size=world_size,
            rank=self._weight_update_rank,
            device_id=self.device,
        )
        torch.distributed.barrier(group=self._model_update_group)
        self.logger.info("vLLM init_process_group finished.")
        self._explorer_name = explorer_name
        self._namespace = namespace
        self.synchronizer = Synchronizer.get_actor(namespace=self._namespace)
        self._checkpoint_converter = None

    def update_weight(self, source: str = "student"):
        """Broadcast weight to all vllm workers from source rank 0 (actor model)"""
        if not hasattr(self, "_weight_update_rank"):
            raise RuntimeError(
                "init_process_group must be called before update_weight. "
                "Ensure setup_weight_sync_group (or setup_model_level_weight_sync_group) "
                "is invoked for all models including auxiliary models (e.g. teacher)."
            )
        if self._weight_update_rank == 0:
            state_dict, model_version = ray.get(
                self.synchronizer.get_model_state_dict.remote(source)
            )
            if isinstance(state_dict, tuple):
                # currently only megatron return a tuple
                method, checkpoint_dir = state_dict
                if method == "megatron":
                    if self._checkpoint_converter is None:
                        from trinity.common.models.utils import get_megatron_converter

                        self._checkpoint_converter = get_megatron_converter(checkpoint_dir)
                    state_dict = self._checkpoint_converter.get_state_dict(checkpoint_dir)
                else:
                    raise NotImplementedError(f"{method} is not supported")
                ray.get(self.synchronizer.set_model_state_dict.remote(state_dict, model_version))
        # For NCCL student sync, metadata is provided during init_process_group.
        # Fall back to synchronizer metadata for other sources or legacy paths.
        state_dict_meta = self._state_dict_meta if source == "student" else None
        if state_dict_meta is None:
            state_dict_meta = ray.get(self.synchronizer.get_state_dict_meta.remote(source))
        if state_dict_meta is None:
            raise RuntimeError(f"No state dict metadata found for source `{source}`.")
        for name, dtype_str, shape in state_dict_meta:
            if self._weight_update_rank == 0:
                weight = state_dict[name]
                weight = weight.to(self.device)
            else:
                dtype = getattr(torch, dtype_str.split(".")[-1])
                weight = torch.empty(shape, dtype=dtype, device=self.device)
            torch.distributed.broadcast(weight, 0, group=self._model_update_group)
            weight = weight.type(self.model_config.dtype)
            self.model_runner.model.load_weights(weights=[(name, weight)])
            del weight
        torch.distributed.barrier(group=self._model_update_group)
        torch.cuda.synchronize()
        torch.cuda.empty_cache()