import logging
import os
import vllm
import torch 
import types
import gc

from .flash_quantization import get_quantize_fn

# Set up logger
logger = logging.getLogger(__name__)

def vllm_model_finder(vllm_llm):
    if not hasattr(vllm_llm.llm_engine, 'model_executor'):
        vllm_llm.llm_engine.model_executor = vllm_llm.llm_engine.engine_core.model_executor
    
    vllm_model = vllm_llm.llm_engine.model_executor.driver_worker.worker.model_runner.model
    return vllm_model

def bond_method_to_cls(func, obj):
    if hasattr(func, '__self__') or not callable(func):
        # If the function is already bound to an instance, return it as is
        return func
    else:
        return types.MethodType(func, obj)

recorded_loader_keys = [
    'weight_loader',
    'load_qkv_weight',
    'load_row_parallel_weight',
    'load_merged_column_weight',
    'output_dim',
    'input_dim',
    '_assert_and_load',
]
def patch_vllm_process_weights_after_loading():
    try:
        # Store the original process_weights_after_loading function
        from vllm.model_executor.model_loader import loader
        
        if not hasattr(loader, 'beforeflashrl_process_weights_after_loading'):
            
            original_process_weights_after_loading = loader._process_weights_after_loading
            loader.beforeflashrl_process_weights_after_loading = original_process_weights_after_loading
            
            def hacked_process_weights_after_loading(
                model, 
                model_config, 
                target_device, 
            ) -> None:
                if model_config is None and target_device is None:
                    model_config = getattr(model, 'hacked_model_config', None)
                    target_device = getattr(model, 'hacked_target_device', None)
                else:
                    setattr(model, 'hacked_model_config', model_config)
                    setattr(model, 'hacked_target_device', target_device)
            
                if getattr(model, 'hacked_not_need_process_weights_after_loading', False):
                    logger.debug("vllm process_weights_after_loading already processed")
                    return

                # print("Patched process_weights_after_loading function called")
                
                original_weights = dict(model.named_parameters())
                
                # this can be optimized for better memory usage, leave for future work...
                if not hasattr(model, 'hacked_original_weights_rebuild_keys'):
                    model.hacked_original_weights_rebuild_keys = {}
                    for name, p in original_weights.items():
                        model.hacked_original_weights_rebuild_keys[name] = (p.shape, p.stride(), p.dtype, p.untyped_storage().nbytes())
                
                # record weight_loader 
                recorded_loader = {k: dict() for k in recorded_loader_keys}
                for name, p in original_weights.items():
                    for k in recorded_loader.keys():
                        if hasattr(p, k):
                            attr = getattr(p, k)
                            if not callable(attr):
                                recorded_loader[k][name] = attr
                            elif p is attr.__self__:
                                recorded_loader[k][name] = attr.__func__
                            else:
                                recorded_loader[k][name] = attr
                
                original_process_weights_after_loading(model, model_config, target_device)
                        
                model.hacked_recorded_loader = recorded_loader

            # Patch the process_weights_after_loading function
            loader._process_weights_after_loading = hacked_process_weights_after_loading
            
            logger.debug("Successfully patched the process_weights_after_loading function of vllm")
        else:
            logger.debug("vllm process_weights_after_loading already patched")
        
        from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
        if not hasattr(BaseKVCacheMethod, 'beforeflashrl_process_weights_after_loading'):
            original_kvcache_process_weights_after_loading = BaseKVCacheMethod.process_weights_after_loading
            BaseKVCacheMethod.beforeflashrl_process_weights_after_loading = original_kvcache_process_weights_after_loading
            def hacked_kvcache_process_weights_after_loading(
                self,
                layer,
            ) -> None:
                # print("Patched kv_cache process_weights_after_loading function called")
                if not hasattr(layer, 'k_scale'):
                    layer.k_scale = -1.0
                if not hasattr(layer, 'v_scale'):
                    layer.v_scale = -1.0
                return original_kvcache_process_weights_after_loading(self, layer)
            BaseKVCacheMethod.process_weights_after_loading = hacked_kvcache_process_weights_after_loading
            
            logger.debug("Successfully patched kv_cache process_weights_after_loading")
        else:
            logger.debug("vllm kv_cache process_weights_after_loading already patched")
            
        return True
        
    except Exception as e:
        logger.debug(f"Error patching vllm process_weights_after_loading: {e}")
        return False

keys_to_overload = [
    'load_format', 
    'quantization', 
    'distributed_executor_backend',
]

def patch_vllm_llm():
    try:                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    
        if not hasattr(vllm.LLM, 'beforeflashrl__init__'):
            # Store the original LLM init function
            original_init = vllm.LLM.__init__
            vllm.LLM.beforeflashrl__init__ = original_init
        
            def hacked_init_(
                self, 
                model: str,
                **kwargs
            ) -> None:
                import vllm.envs as envs
                assert envs.VLLM_USE_V1, 'flash_rl only supports vllm v1 for now'

                config = os.environ.get("FLASHRL_CONFIG", None)
                
                if 'distributed_executor_backend' not in kwargs or kwargs['distributed_executor_backend'] != 'external_launcher':
                    logger.error("flash_rl only supports external_launcher for now")
                
                assert 'RANK' in os.environ and 'WORLD_SIZE' in os.environ, \
                    'flash_rl only supports external_launcher for now'
                    
                rank = int(os.environ.get("RANK", None))
                mp_size = kwargs.get('tensor_parallel_size', 1) * kwargs.get('pipeline_parallel_size', 1)
                dp_rank = rank // mp_size
                
                if config is not None:
                    # Load the config file and set the model
                    # Assuming config is a JSON file, you can use json.load() to read it
                    logger.info(f"flash_rl config detected.")
                    logger.info(f"Loading flash_rl config from: {config}")
                    
                    config_path = config.strip()
                    if not os.path.exists(config_path):
                        from huggingface_hub import hf_hub_download
                        config_path = config_path.split('/')
                        assert len(config_path) >= 3, f'Invalid flash_rl config path: {config_path}'
                        config_path = hf_hub_download(repo_id='/'.join(config_path[:2]), filename='/'.join(config_path[2:]))

                    import yaml
                    with open(config_path, 'r') as f:
                        config_data = yaml.safe_load(f)
                        
                    config_count = len(config_data['configs'])
                    config_index = dp_rank % config_count
                    logger.info(f"Using config {config_index} of {config_count}")
                    config_data = config_data['configs'][config_index]
                    
                    for k, v in config_data.items():
                        logger.info(f"rank {rank} flash_rl config: {k}: {v}")
                    
                    
                    for key in keys_to_overload:
                        if key in config_data:
                            logger.debug(f"Overloading {key} with {config_data[key]}")
                            kwargs[key] = config_data.get(key)
                    
                    if config_data.get('fn', 'int8') != 'bf16':
                        model = config_data.get('model')
                        
                        quant_profile = config_data.get('profile', os.path.join(model, 'profile.pt'))
                        logger.debug(f"Loading flash_rl profile from: {quant_profile}")
                        
                        quant_profile_path = quant_profile.strip()
                        if not os.path.exists(quant_profile_path):
                            from huggingface_hub import hf_hub_download
                            quant_profile_path = quant_profile_path.split('/')
                            assert len(quant_profile_path) >= 3, f'Invalid flash_rl profile path: {quant_profile_path}'
                            quant_profile_path = hf_hub_download(repo_id='/'.join(quant_profile_path[:2]), filename='/'.join(quant_profile_path[2:]))
                        
                        self.flash_rl_profile = torch.load(quant_profile_path)
                        
                        if 'module_attribute_to_preserve' in config_data:
                            logger.debug(f"flash_rl module_attribute_to_preserve: {config_data['module_attribute_to_preserve']}")
                            self.flash_rl_module_attribute_to_preserve = config_data.get('module_attribute_to_preserve')
                        else:
                            self.flash_rl_module_attribute_to_preserve = []
                        
                else:
                    logger.info(f"flash_rl config not detected.")
                    logger.info(f"Using the original model: {model}")
                
                # Call the parent's __init__ with the custom model
                init_return = original_init(
                    self, 
                    model, 
                    **kwargs,
                )
                
                if hasattr(self, 'flash_rl_profile') and \
                    (not hasattr(model, 'beforeflashrl_load_weights')) and \
                    (config_data.get('fn', 'int8') != 'bf16'):
                        
                    model = vllm_model_finder(self)
                    quant_fn = config_data.get('fn', 'int8')
                    logger.debug(f"flash_rl quantization function: {quant_fn}")
                    flash_quantize_fn = get_quantize_fn(quant_fn)
                     
                    # Store the original load_weights function
                    original_load_weights = model.load_weights
                    model.beforeflashrl_load_weights = original_load_weights
                    def hacked_load_weights(
                        weights,
                    ):
                        # print("flash_rl quant load_weights is called")
                        setattr(model, 'hacked_not_need_process_weights_after_loading', False)
                        
                        if len(self.flash_rl_module_attribute_to_preserve) > 0:
                            for _, module in model.named_modules():
                                for attr in self.flash_rl_module_attribute_to_preserve:
                                    if torch.is_tensor(getattr(module, attr, None)):
                                        # print(f"flash_rl reserving {attr} in module {module}")
                                        setattr(module, f'hacked_{attr}', getattr(module, attr))
                        
                        existing_params = dict(model.named_parameters())
                        
                        hacked_data_dict = {}
                        for name, p in existing_params.items():
                            hacked_data_dict[name] = p.data
                        
                        assert hasattr(model, "hacked_original_weights_rebuild_keys")
                        
                        for name, (shape, stride, dtype, nbytes) in model.hacked_original_weights_rebuild_keys.items():
                            if name in existing_params:
                                # existing_nbytes = existing_params[name].untyped_storage().nbytes()
                                # print(
                                #     f"flash_rl load_weights review {name} ({existing_nbytes}): "
                                #     f"{existing_params[name].dtype} x {existing_params[name].shape}"
                                #     f" --> {dtype} x {shape}"
                                # )
                                existing_params[name].data = torch.empty(shape, dtype=dtype) 
                        
                        for k, loader_k in model.hacked_recorded_loader.items():
                            for n, loader in loader_k.items():
                                if not hasattr(existing_params[n], k):
                                    setattr(existing_params[n], k, bond_method_to_cls(loader, existing_params[n]))

                        updated_params = original_load_weights(
                            flash_quantize_fn(weights, self.flash_rl_profile)
                        )
                        
                        if hasattr(model, 'hacked_model_config') and hasattr(model, 'hacked_target_device'):        
                            from vllm.model_executor.model_loader import loader
                            loader._process_weights_after_loading(model, None, None)
                            setattr(model, 'hacked_not_need_process_weights_after_loading', True)
                        else:
                            setattr(model, 'hacked_not_need_process_weights_after_loading', False)
                        
                        skipped_params = list()
                        for name, p in model.named_parameters():
                            assert name in hacked_data_dict, f'param {name} is not in hacked_data dict'
                            assert hacked_data_dict[name].dtype == p.data.dtype, f'param {name} dtype mismatch: {hacked_data_dict[name].dtype} vs {p.data.dtype}'
                            assert hacked_data_dict[name].numel() == p.data.numel(), f'param {name} numel() mismatch: {hacked_data_dict[name].numel()} vs {p.data.numel()}'
                            
                            if name in updated_params:
                                trided_data = torch.as_strided(p.data, hacked_data_dict[name].shape, hacked_data_dict[name].stride())
                                hacked_data_dict[name].copy_(trided_data)
                            else:
                                skipped_params.append(name)
                                
                            tmp_data = p.data
                            p.data = hacked_data_dict[name]
                            del tmp_data
                        
                        logger.debug(f"flash_rl load_weights skipped params: {skipped_params}")
                        del skipped_params
                        del hacked_data_dict
                        del existing_params
                        gc.collect()
                        torch.cuda.empty_cache()
                        
                        if len(self.flash_rl_module_attribute_to_preserve) > 0:
                            for _, module in model.named_modules():
                                for attr in self.flash_rl_module_attribute_to_preserve:
                                    if torch.is_tensor(getattr(module, attr, None)):
                                        # print(f"flash_rl recovering {attr} in module {module}")
                                        assert hasattr(module, f'hacked_{attr}'), f"module {module} does not have attribute hacked_{attr}"
                                        setattr(module, attr, getattr(module, f'hacked_{attr}'))
                                        delattr(module, f'hacked_{attr}')
                                        
                        # return []
                        return updated_params
                    
                    model.load_weights = hacked_load_weights
                    logger.debug("Successfully patched the load_weights function of vllm")
                else:
                    logger.debug("vllm load_weights patching skipped")
                
                return init_return
            
            # Patch the LLM init function
            vllm.LLM.__init__ = hacked_init_
            
            logger.debug("Successfully patched vllm")
        else:
            logger.debug("vllm LLM already patched")
        return True
        
    except Exception as e:
        logger.error(f"Error patching vllm LLM: {e}")
        return False

def patch_vllm_lmhead_to_fp32():
    try:
        if not hasattr(vllm.model_executor.layers.vocab_parallel_embedding.VocabParallelEmbedding, 'beforeflashrl_VocabParallelEmbedding'):
            # Store the original LLM init function
            original_init = vllm.model_executor.layers.vocab_parallel_embedding.VocabParallelEmbedding.__init__
            vllm.model_executor.layers.vocab_parallel_embedding.VocabParallelEmbedding.beforeflashrl_VocabParallelEmbedding = original_init

            def hacked_init_(
                self, 
                *args, 
                **kwargs
            ) -> None:
               
                if len(args) >= 3:
                    output_dtype = args[2]
                    args = list(args)
                    args[2] = torch.float32
                    args = tuple(args)
                else:
                    output_dtype = kwargs.get('params_dtype', torch.get_default_dtype())
                    kwargs['params_dtype'] = torch.float32

                if len(args) >= 6:
                    args = list(args)
                    args[5] = None
                    args = tuple(args)
                else:
                    kwargs['quant_config'] = None


                init_return = original_init(
                    self, 
                    *args,
                    **kwargs,
                )
                self.output_dtype = output_dtype
                
                if not hasattr(self, 'beforeflashrl_forward'):
                    original_forward = self.forward 
                    self.beforeflashrl_forward = original_forward

                    def hacked_forward(
                        self, 
                        *args, 
                        **kwargs,
                    ):
                        original_forward_return = original_forward(*args, **kwargs)
                        return original_forward_return.to(self.output_dtype)

                    setattr(self, 'forward', bond_method_to_cls(hacked_forward, self))
                    logger.debug("Successfully patched vllm VocabParallelEmbedding forward to fp32 weight")
                else: 
                    logger.debug("vllm VocabParallelEmbedding forward already patched to fp32 weight")
                return init_return
            
            # Patch the LLM init function
            vllm.model_executor.layers.vocab_parallel_embedding.VocabParallelEmbedding.__init__ = hacked_init_
            print("Successfully patched vllm VocabParallelEmbedding at init")
        else:
            print("vllm VocabParallelEmbedding init already patched")
        
        status = True
    except Exception as e:
        
        status = False
        logger.error(f"Error patching VocabParallelEmbedding init: {e}")

    try: 
        if not hasattr(vllm.model_executor.layers.vocab_parallel_embedding.UnquantizedEmbeddingMethod, 'beforeflashrl_apply'):
            # Store the original LLM init function
            original_apply = vllm.model_executor.layers.vocab_parallel_embedding.UnquantizedEmbeddingMethod.apply
            vllm.model_executor.layers.vocab_parallel_embedding.UnquantizedEmbeddingMethod.beforeflashrl_apply = original_apply

            def hacked_apply_(
                self,
                layer: torch.nn.Module,
                x: torch.Tensor,
                bias = None
            ) -> torch.Tensor:
                original_type = x.dtype
                x = x.to(dtype=torch.float32)
                return original_apply(self, layer, x, bias).to(original_type)
             
            vllm.model_executor.layers.vocab_parallel_embedding.UnquantizedEmbeddingMethod.apply = hacked_apply_
            print("Successfully patched vllm UnquantizedEmbeddingMethod at apply")
        else:
            print("vllm UnquantizedEmbeddingMethod apply already patched")
        status = status
    except Exception as e:
        status = False
        logger.error(f"Error patching UnquantizedEmbeddingMethod apply: {e}")
    
    return status 

def patch_vllm_logprob_compute():
    try:
        from vllm.v1.sample.sampler import Sampler
        if not hasattr(Sampler, 'beforeflashrl_forward'):
            # Store the original LLM init function
            original_forward = Sampler.forward
            Sampler.beforeflashrl_forward = original_forward

            def hacked_logprob_forward(
                self,
                logits: torch.Tensor,
                sampling_metadata,
            ):
                # NOTE(woosuk): Use the original logits (before any penalties or
                # temperature scaling) for the top-k logprobs.
                # This is different from the V0 sampler, which uses the logits that
                # is used for sampling (after penalties and temperature scaling).
                # TODO(rob): provide option for logprobs post sampling.
                # See https://vllm-dev.slack.com/archives/C07UUL8E61Z/p1735907856007919 # noqa: E501

                # Use float32 for the logits.
                logits = logits.to(torch.float32)
                # Apply allowed token ids.
                logits = self.apply_allowed_token_ids(logits, sampling_metadata)
                # Apply bad words exclusion.
                logits = self.apply_bad_words(logits, sampling_metadata)
                # Apply logits bias.
                logits = self.apply_logits_bias(logits, sampling_metadata)
                # Apply penalties (e.g., min_tokens, freq_penalties).
                logits = self.apply_penalties(logits, sampling_metadata)
                # Sample the next token.
                    
                sampled, processed_logits = self.sample(logits, sampling_metadata)
                # Convert sampled token ids to int64 (long) type to ensure compatibility
                # with subsequent operations that may use these values as indices.
                # This conversion is necessary because FlashInfer sampling operations
                # return int32 (while PyTorch argmax and topk return int64).
                sampled = sampled.long()

                # Gather the logprobs of the topk and sampled token (if requested).
                # Get logprobs and rank tensors (if requested)
                if sampling_metadata.max_num_logprobs is not None:
                    processed_logprobs = self.compute_logprobs(logits)
                    logprobs_tensors = self.gather_logprobs(processed_logprobs, 0, token_ids=sampled)
                else:
                    logprobs_tensors = None

                # Use int32 to reduce the tensor size.
                sampled = sampled.to(torch.int32)

                # These are GPU tensors.
                sampler_output = vllm.v1.outputs.SamplerOutput(
                    # The sampled tokens are expanded to 2D tensor with shape
                    # [num_requests, 1], where each row represents one generated
                    # token per request.
                    sampled_token_ids=sampled.unsqueeze(-1),
                    logprobs_tensors=logprobs_tensors,
                )
                return sampler_output
            
            # Patch the LLM init function
            Sampler.forward = hacked_logprob_forward
            
            def hacked_sample(
                self,
                logits,
                sampling_metadata,
            ) -> torch.Tensor:
                assert not (sampling_metadata.all_greedy
                            and sampling_metadata.all_random)
                if sampling_metadata.all_random:
                    greedy_sampled = None
                else:
                    greedy_sampled = self.greedy_sample(logits)
                    if sampling_metadata.all_greedy:
                        return greedy_sampled, logits

                assert sampling_metadata.temperature is not None

                # Apply temperature.
                logits = self.apply_temperature(logits, sampling_metadata.temperature)

                # Apply min_p.
                if sampling_metadata.min_p is not None:
                    logits = self.apply_min_p(logits, sampling_metadata.min_p)

                # Apply top_k and/or top_p.
                random_sampled, processed_logits = self.topk_topp_sampler(
                    logits,
                    sampling_metadata.generators,
                    sampling_metadata.top_k,
                    sampling_metadata.top_p,
                )

                if greedy_sampled is None:
                    return random_sampled, processed_logits

                sampled = torch.where(
                    sampling_metadata.temperature < 1e-5,
                    greedy_sampled,
                    random_sampled,
                    out=greedy_sampled,  # Reuse tensor
                )
                return sampled, processed_logits
            
            original_sample = Sampler.sample
            Sampler.beforeflashrl_sample = original_sample
            Sampler.sample = hacked_sample

            from .vllm_sampler import HackedTopKTopPSampler
            
            def hacked_init_(self):
                super(self.__class__, self).__init__()
                self.topk_topp_sampler = HackedTopKTopPSampler()
            original_init = Sampler.__init__
            Sampler.beforeflashrl_init = original_init
            Sampler.__init__ = hacked_init_
            
            logger.debug("Successfully patched Sampler at init")
        else:
            logger.debug("vllm Sampler already patched")
            
        status = True
    
    except Exception as e:
        logger.error(f"Error patching Sampler forward: {e}")
        status = False

    return status

def patch_vllm_llm_test_reload():
    try:
        if not hasattr(vllm.LLM, 'beforeflashrl_test__init__'):
            # Store the original LLM init function
            test_vllm_init = vllm.LLM.__init__
            vllm.LLM.beforeflashrl_test__init__ = test_vllm_init
        
            def test_reload_at_init_(
                self, 
                *args, 
                **kwargs
            ) -> None:
                # Call the parent's __init__ with the custom model
                init_return = test_vllm_init(self, *args, **kwargs)
                
                config = os.environ.get("FLASHRL_TEST_RELOAD", None)
                
                if config is not None:
                    config = config.split(',')
                    from transformers import AutoModelForCausalLM
                    for config_i in config:
                        model_to_be_reloaded = AutoModelForCausalLM.from_pretrained(config_i, device_map="cpu", torch_dtype="auto")
                        device = torch.cuda.current_device()
                        print(f"FLASH_RL re-loading model {config_i} to device {device}")
                        model = vllm_model_finder(self)
                        model.load_weights(
                            ((name, param.to(device)) for name, param in model_to_be_reloaded.named_parameters())
                        )
                else:
                    print(f"FLASH_RL re-loading model not detected.")
                
                return init_return
            
            # Patch the LLM init function
            vllm.LLM.__init__ = test_reload_at_init_
            
            print("Successfully patched vllm reload at init")
        else:
            print("vllm LLM reload already patched")
        return True
        
    except Exception as e:
        logger.error(f"Error patching vllm reload LLM: {e}")
        return False

