import logging
import sys
import uuid

from omegaconf import OmegaConf
from hydra.utils import get_original_cwd

logger = logging.getLogger(__name__)
OmegaConf.register_new_resolver("uuid", lambda: str(uuid.uuid4()))


def _do_model_op(
    mode,
    model,
    data,
    trainer=None,
    seed=42,
    mlflow=None,
    full_conf={},
    test=False,
    multiprocessing_strategy=None,
    make_notebook=None,
    **_,
):

    from pathlib import Path

    import pytorch_lightning as pl
    from hydra.utils import instantiate

    from .mlflow_utils import mlflow_fit, mlflow_predict, mlflow_test
    from .utils import make_notebook as mk_notebook
    from .utils import save_model_predictions

    if multiprocessing_strategy is not None:
        import torch

        torch.multiprocessing.set_sharing_strategy(multiprocessing_strategy)

    if mode in ["train", "test", "predict"]:
        if make_notebook is not None:
            mk_notebook(full_conf, make_notebook)
            return

        pl.seed_everything(seed)
        logger.info("Instantiating model, datamodule and trainer")
        model = instantiate(model)
        data = instantiate(data)
        trainer = instantiate(trainer)

        if mode == "train":
            if mlflow is not None and mlflow.get("tracking_uri") is not None:
                mlflow_fit(mlflow, trainer, model, data, full_conf=full_conf, test=test)
            else:
                logger.info("Calling trainer.fit")
                trainer.fit(model, data)

        elif mode == "test":
            if mlflow is not None and mlflow.get("tracking_uri") is not None:
                mlflow_test(mlflow, trainer, model, data, full_conf=full_conf)
            else:
                if "ckpt_path" not in full_conf:
                    raise NotImplementedError(
                        "Cannot `serotiny test` without "
                        "an MLFlow config, or a local ckpt_path."
                    )
                ckpt_path = full_conf["ckpt_path"]
                trainer.test(model, data, ckpt_path=ckpt_path)

        elif mode == "predict":
            if mlflow is not None and mlflow.get("tracking_uri") is not None:
                mlflow_predict(mlflow, trainer, model, data, full_conf=full_conf)
            else:
                if "ckpt_path" not in full_conf:
                    raise NotImplementedError(
                        "Cannot `serotiny predict` without "
                        "an MLFlow config, or a local ckpt_path."
                    )

                preds_dir_default = Path(get_original_cwd()) / "predictions"
                preds_dir = Path(
                    full_conf.get("predictions_output_dir", preds_dir_default)
                )
                preds_dir.mkdir(exist_ok=True, parents=True)

                ckpt_path = full_conf["ckpt_path"]
                preds = trainer.predict(model, data, ckpt_path=ckpt_path)
                save_model_predictions(model, preds, preds_dir)

    else:
        raise ValueError(f"`mode` must be 'train', 'test' or 'predict'. Got '{mode}'")


def _do_model_op_wrapper(cfg):
    if isinstance(cfg, dict):
        cfg = OmegaConf.create(cfg)

    # there might be other dots in the
    # executable path, ours is the last
    mode = sys.argv[0].split(".")[-1]
    if mode not in ["train", "predict", "test"]:
        try:
            # if we can't get mode from the command line,
            # try getting it from the config. this might be
            # needed for runs on distributed clusters like ray
            mode = cfg.pop("mode")
        except (AttributeError, KeyError) as e:
            logger.error(
                f"serotiny was misconfigured. Unable to infer "
                f"mode from either config or sys.argv:\n"
                f"{sys.argv}"
            )
            raise e
    else:
        cfg.pop("mode", None)

    _do_model_op(
        mode,
        **cfg,
        full_conf=cfg,
    )
