from typing import List, Union

class VLLMError(Exception):
    def __init__(self, status_code, message):
        self.status_code = status_code
        self.message = message
        super().__init__(self.message)


try:
    from vllm import LLM, SamplingParams
except ImportError as e:
    raise VLLMError(
        status_code=1,
        message="Failed to import 'vllm' package. Make sure it is installed correctly.",
    ) from e


class VLLM:
    def __init__(
        self,
        model: str,
        ngpus: int = 1,
        model_save_dir: str = None,
        **vllm_kwargs,
    ) -> None:
        """
        A wrapper class for the VLLM model.

        Arguments:
            model (str): The name of the VLLM model to use.
            **vllm_kwargs: Additional keyword arguments to pass to the VLLM model.  
        """
        self.model: str = model
        self.ngpus: int = ngpus
        self.model_save_dir: str = model_save_dir

        self.model: LLM = LLM(
            model=self.model,
            tensor_parallel_size=self.ngpus,
            download_dir=self.model_save_dir,
            tokenizer_mode='auto',
            **vllm_kwargs,
            trust_remote_code=True
        )

    def validate_vllm(self):
        return True

    def completions(
        self,
        prompts: List[str],
        use_tqdm: bool = True,
        **kwargs: Union[int, float, str],
    ) -> List[str]:
        """
        Generate completions for the given prompts using the VLLM model.

        Arguments:
            prompts (List[str]): A list of prompts to generate completions for.
            use_tqdm (bool): Whether to display a progress bar for the completions generation process.
            **kwargs: Additional keyword arguments to pass to the VLLM model.
        
        Returns:
            List[str]: A list of completions generated by the VLLM model.
        """
        prompts = [prompt.strip() for prompt in prompts]
        params = SamplingParams(**kwargs)

        outputs = self.model.generate(prompts, params, use_tqdm=use_tqdm)
        outputs = [output.outputs[0].text for output in outputs]
        return outputs