import json
from abc import ABC, abstractmethod
from dataclasses import dataclass

import numpy as np
import requests
import torch
from omegaconf import DictConfig, OmegaConf

from openelm.codegen.codegen_utilities import model_setup, sample, set_seed, truncate


class MutationModel(ABC):
    """Base model class for all mutation models."""

    @abstractmethod
    def generate_program(self, seed_str: str) -> dict:
        pass


@dataclass(frozen=True)
class FunctionTemplate:
    """
    A function template for a mutation model.

    Attributes:
        func_name: (str) The name of the function that we want to execute.
        import_line: (str) The import lines we add to the code.
        func_preamble: (str) The function definition, as well as potentially a
        few initial lines to generate code.
        return_line: (str) The return line we add to the end of the code.
    """

    func_name: str
    import_line: str
    func_preamble: str
    return_line: str


class PromptMutationModel(MutationModel):
    """Mutation model that uses prompts to change a seed."""

    def __init__(
        self,
        cfg,
        function_template: FunctionTemplate,
        sandbox_server="http://localhost:5000",
    ) -> None:
        if isinstance(cfg, str):
            self.cfg = OmegaConf.load(cfg)
        elif isinstance(cfg, (dict, DictConfig)):
            self.cfg = DictConfig(cfg)
        else:
            raise ValueError

        set_seed(self.cfg.seed)
        # Use RNG to rotate random seeds during inference.
        self.rng = np.random.default_rng(seed=self.cfg.seed)
        self.sandbox_server = sandbox_server
        self.model, self.tokenizer = model_setup(self.cfg)
        self.func_template: FunctionTemplate = function_template

    def generate_prompt_str(
        self,
        seed: str,
        tokenizer=None,
        batch_size=None,
        append_return=True,
        without_trunc=True,
    ) -> list[str]:
        """
        Generate a mutation string from prompting a model with a seed and function template.

        Args:
            seed: (str) the seed text.
            tokenizer: (Optional) assign only if you want to use a different
            tokenizer (default: None).
            batch_size: (Optional) override the batch size in config.
            append_return: (Optional) append a return line to the code in the end.
            without_trunc: (Optional) True if we don't apply the `truncate` function.

        Returns:
            A list of code generated by the model.
        """
        tokenizer = self.tokenizer if tokenizer is None else tokenizer
        encoding = tokenizer(
            [seed + "\n\n" + self.func_template.func_preamble],
            truncation=True,
            padding=True,
            max_length=self.cfg.gen_max_len,
            return_tensors="pt",
        )

        cfg = OmegaConf.merge(
            self.cfg,
            {"batch_size": self.cfg.batch_size if batch_size is None else batch_size},
        )
        with torch.no_grad():
            completion = sample(cfg, self.model, self.tokenizer, encoding)
        # Reset random seed
        set_seed(int(self.rng.integers(0, 1e8)))

        if without_trunc:
            truncation = completion
        else:
            truncation = [
                truncate(code, print_num=float("inf"), only_local_scope=True)
                for code in completion
            ]

        truncation = [
            self.func_template.import_line
            + "\n"
            + self.func_template.func_preamble
            + "\n"
            + code
            for code in truncation
        ]

        if append_return:
            truncation = [
                code + "\n" + self.func_template.return_line for code in truncation
            ]

        return truncation

    def generate_program(self, code: str) -> list[dict]:
        """
        Given a piece of code, do prompt mutation, call the sandbox server to execute the code and return the result.

        Args:
            code: the full code string.
        Returns:
            a numpy array (if successful) or the exception object.
        """
        results = []
        for code in self.generate_prompt_str(code):
            resp = self._get_response(code, self.cfg.timeout)
            if resp.status_code == 200:
                return_dict = json.loads(resp.text)
                self._post_process(return_dict)
                error_code = "0"
            elif resp.status_code == 500:  # Bad request
                try:
                    msg = json.loads(resp.text)
                    return_dict = {"program_str": code, "result_obj": msg["message"]}
                    error_code = msg["unsafe_execute_error_code"]
                except Exception as e:
                    return_dict = {"program_str": code, "result_obj": str(e)}
                    error_code = 6
            else:
                return_dict = {"program_str": code, "result_obj": resp.text}
                error_code = 6

            results.append({**return_dict, "error_code": error_code})

        return results

    @abstractmethod
    def _get_response(self, code: str, timeout: int) -> requests.models.Response:
        pass

    @abstractmethod
    def _post_process(self, response_dict: dict) -> dict:
        pass


class PromptMutationForSodarace(PromptMutationModel):
    def __init__(self, cfg, sandbox_server="http://localhost:5000") -> None:
        function_template = FunctionTemplate(
            func_name="make_walker",
            import_line="from .walker import walker_creator",
            func_preamble="def make_walker():\n\twc = walker_creator()\n",
            return_line="\treturn wc.get_walker()\n",
        )
        super().__init__(cfg, function_template, sandbox_server)

    def _get_response(self, code: str, timeout: int) -> requests.models.Response:
        return requests.post(
            f"{self.sandbox_server}/gen_racer",
            json={"code": code, "timeout": timeout},
            timeout=timeout,
        )

    def _post_process(self, response_dict: dict) -> dict:
        pass


class PromptMutationForImgTask(PromptMutationModel):
    def __init__(self, cfg, sandbox_server="http://localhost:5000") -> None:
        func_name = "draw"
        func_preamble = (
            f'def {func_name}():\n\t"""Draw a yellow circle.\n'
            '\t"""\n\tpic = np.zeros((32, 32, 3))\n'
        )
        function_template = FunctionTemplate(
            func_name=func_name,
            import_line="import math\nimport numpy as np",
            func_preamble=func_preamble,
            return_line="",
        )
        super().__init__(cfg, function_template, sandbox_server)

    def reset_shape(self, shape: tuple):
        self.func_preamble = f'def {self.func_name}():\n\t"""Draw a yellow circle.\n\t"""\n\tpic = np.zeros({shape})\n'

    def _get_response(self, code: str, timeout: int) -> requests.models.Response:
        return requests.post(
            f"{self.sandbox_server}/eval_imageoptim_func",
            json={"code": code, "func_name": self.func_name, "timeout": timeout},
            timeout=timeout,
        )

    def _post_process(self, response_dict: dict) -> dict:
        response_dict["result_obj"] = np.array(response_dict["result_obj"])
        return response_dict


# TODO: complete diff model (when it's available)
class DiffModel(MutationModel):
    def __init__(self, cfg) -> None:
        raise NotImplementedError()

    def generate_prompt_str(self, seed, tokenizer):
        raise NotImplementedError()

    def generate_program(self, seed_str: str) -> dict:
        raise NotImplementedError()
