"""
InferenceUtil and construct_prompt are adapted from the ClassEval project by FudanSELab.

https://github.com/FudanSELab/ClassEval/blob/master/generation/inference_util.py
"""

from enum import Enum


class ModelName(Enum):
    Instruct_CodeGen = 0
    WizardCoder = 1
    Instruct_StarCoder = 2
    InCoder = 3
    PolyCoder = 4
    SantaCoder = 5
    Vicuna = 6
    ChatGLM = 7
    GPT_3_5 = 8
    GPT_4 = 9
    others = 10
    Magicoder = 11
    CodeGeeX2 = 12
    DeepSeekCoder_inst = 13
    Gemini_Pro = 14
    CodeLlama_13b_inst = 15


class GenerationStrategy(Enum):
    Holistic = 0
    Incremental = 1
    Compositional = 2


class InferenceUtil:
    @staticmethod
    def generate_prompt(instruction: str, model: ModelName) -> str:
        if model in (ModelName.DeepSeekCoder_inst, ModelName.Gemini_Pro):
            return instruction

        elif model == ModelName.Magicoder:
            return f"""You are an exceptionally intelligent coding assistant that consistently delivers accurate and reliable responses to user instructions.

            @@ Instruction:
            {instruction}

            @@ Response:
            """
        else:
            return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

            ### Instruction:
            {instruction}

            ### Response:
            """


def construct_prompt(
    record: dict[str, str],
    model: ModelName = ModelName.GPT_3_5,
    strategy: GenerationStrategy = GenerationStrategy.Holistic,
) -> str:
    """
    Construct a prompt for the given record, model, and strategy.

    From https://github.com/FudanSELab/ClassEval/blob/master/generation/inference_pipeline.py
    FudanSELab/ClassEval allows passing the model as an argument, using integers.
    For now, we don't do that here and just default to using the prompts that work for GPT3.5 and the holistic strategy.
    """
    prompt = ""
    if strategy == GenerationStrategy.Holistic:
        if model == ModelName.PolyCoder or model == ModelName.SantaCoder:
            skeleton = record["skeleton"]
            prompt = skeleton
        else:
            class_name = record["class_name"]
            skeleton = record["skeleton"]
            instruction = (
                f"Please complete the class {class_name} in the following code."
            )
            instruction = instruction + "\n" + skeleton
            prompt = InferenceUtil.generate_prompt(instruction, model)

    elif strategy == GenerationStrategy.Incremental:
        if model == ModelName.PolyCoder or model == ModelName.SantaCoder:
            prompt = record["skeleton"]
        else:
            prompt = record["instruction"] + record["skeleton"]
            prompt = InferenceUtil.generate_prompt(prompt, model)

    elif strategy == GenerationStrategy.Compositional:
        if model == ModelName.PolyCoder or model == ModelName.SantaCoder:
            prompt = record["skeleton"]
        else:
            prompt = record["instruction"] + record["skeleton"]
            prompt = InferenceUtil.generate_prompt(prompt, model)

    return prompt
