import os
import json
import requests

from nmtwizard.framework import Framework
from nmtwizard.logger import get_logger
from nmtwizard import utils, serving

logger = get_logger(__name__)

_MODEL_NAME = "model.pt"
_RELEASED_MODEL_NAME = "model_released.pt"


class OpenNMTPYFramework(Framework):
    def train(
        self,
        config,
        src_file,
        tgt_file,
        src_vocab_info,
        tgt_vocab_info,
        align_file=None,
        example_weights_file=None,
        model_path=None,
        gpuid=0,
        models_to_average=None,
    ):
        if models_to_average:
            logger.warning("Checkpoint averaging is not implemented in this framework")

        # Preprocess training files.
        options = config["options"].get("config", {})
        options_preprocess = options.get("preprocess", {}).copy()
        options_preprocess["src_vocab"] = self._convert_vocab(
            config["vocabulary"]["source"]["path"]
        )
        options_preprocess["tgt_vocab"] = self._convert_vocab(
            config["vocabulary"]["target"]["path"]
        )
        bin_file = os.path.join(self._data_dir, "bin")
        cmd = [
            "onmt_preprocess",
            "-train_src",
            src_file,
            "-train_tgt",
            tgt_file,
            "-save_data",
            bin_file,
        ]
        cmd += _build_cmd(options_preprocess)
        utils.run_cmd(cmd)

        # Train.
        options_train = options.get("train", {}).copy()
        if "train_steps" not in options_train:
            options_train["single_pass"] = True
            options_train["train_steps"] = 0
        options_train["save_checkpoint_steps"] = 0
        options_train["data"] = bin_file
        options_train["save_model"] = self._output_dir + "/model"
        if isinstance(gpuid, list):
            options_train["world_size"] = len(gpuid)
            options_train["gpu_ranks"] = " ".join(str(i - 1) for i in gpuid)
        elif gpuid > 0:
            options_train["gpu_ranks"] = gpuid - 1
        if model_path is not None:
            options_train["train_from"] = os.path.join(model_path, _MODEL_NAME)
        utils.run_cmd(["onmt_train"] + _build_cmd(options_train))

        # Select model.
        models = os.listdir(self._output_dir)
        if not models:
            raise RuntimeError("no model generated by the training")
        if len(models) > 1:
            raise RuntimeError("more than one model generated by the training")
        model_file = os.path.join(self._output_dir, models[0])
        return {_MODEL_NAME: model_file}

    def trans(self, config, model_path, input, output, gpuid=0):
        options_trans = _trans_options(config, gpuid)
        options_trans["model"] = os.path.join(model_path, _MODEL_NAME)
        options_trans["src"] = input
        options_trans["output"] = output
        utils.run_cmd(["onmt_translate"] + _build_cmd(options_trans))

    def serve(self, config, model_path, gpuid=0):
        server_config_path = os.path.join(self._output_dir, "conf.json")
        with open(server_config_path, "w") as server_config_file:
            json.dump(
                {
                    "models_root": model_path,
                    "models": [
                        {
                            "id": 0,
                            "model": _RELEASED_MODEL_NAME,
                            "opt": _trans_options(config, gpuid),
                        }
                    ],
                },
                server_config_file,
            )
        port = serving.pick_free_port()
        process = utils.run_cmd(
            [
                "onmt_server",
                "--ip",
                "127.0.0.1",
                "--port",
                str(port),
                "--url_root",
                "/translator-backend",
                "--config",
                server_config_path,
            ],
            background=True,
        )
        return process, {"port": port}

    def release(self, config, model_path, gpuid=0):
        model = os.path.join(model_path, _MODEL_NAME)
        released_model = os.path.join(self._output_dir, _RELEASED_MODEL_NAME)
        utils.run_cmd(["onmt_release_model", "-m", model, "-o", released_model])
        return {_RELEASED_MODEL_NAME: released_model}

    def forward_request(self, model_info, inputs, outputs=None, options=None):
        if options is None:
            options = {}
        data = [{"src": " ".join(tokens), "id": 0} for tokens in inputs]
        try:
            result = requests.post(
                "http://127.0.0.1:%d/translator-backend/translate" % model_info["port"],
                json=data,
                timeout=options.get("timeout"),
            )
            return [
                [serving.TranslationOutput(r["tgt"].split(), score=r["pred_score"])]
                for r in result.json()[0]
            ]
        except requests.exceptions.Timeout as e:
            logger.error("%s", e)
            return None

    def _map_vocab_entry(self, index, token, vocab):
        if index == 0:
            vocab.write("<unk>\n")
            vocab.write("<blank>\n")
            vocab.write("<s>\n")
            vocab.write("</s>\n")
        vocab.write("%s\n" % token)


def _trans_options(config, gpuid):
    options = config["options"].get("config", {})
    opt = options.get("trans", {}).copy()
    if gpuid > 0:
        opt["gpu"] = gpuid - 1
    return opt


def _build_cmd(options):
    opts = []
    for k, v in options.items():
        if isinstance(v, bool):
            if v:
                opts.append("-%s" % k)
        else:
            opts.append("-%s" % k)
            opts.append(str(v))
    return opts


if __name__ == "__main__":
    OpenNMTPYFramework().run()
