"""
Experimental.
Rewrite the models and experiments JSON export files generated by bulk export-models.
Allows you to post-process each model or experiment.

For example:
  - Model - rewrite model.json to only have one production version.
  - Experiment - rewrite experiment.json to have only the latest run.

Usage:

python -u -m mlflow_migration.tools.rewrite_export \
  --input-dir export-models \
  --custom_rewriters-module samples/custom_export_rewriters.py

The 'custom rewriters module' expects two methods:
  rewrite_model(dct) - processes the content of a model.json file.
  rewrite_experiment(dct) - processes the content of a experiment.json file.

"""

import os
import shutil
import click
from mlflow_migration.common import io_utils


def _load_func(func_name, module_dct, module_path):
    func = module_dct.get(func_name, None)
    if not func:
        raise TypeError(f"Cannot find function '{func_name}()' in '{module_path}'")
    return func


def _load_module(module_path):
    import runpy

    module_dct = runpy.run_path(module_path)
    rewrite_model = _load_func("rewrite_model", module_dct, module_path)
    rewrite_experiment = _load_func("rewrite_experiment", module_dct, module_path)
    return rewrite_model, rewrite_experiment


def _rewrite_object(object_dir, object_id, object_type, rewrite_object_func):
    """Process either a model or experiment"""
    object_dir = os.path.join(object_dir, object_id)
    path = os.path.join(object_dir, f"{object_type}.json")
    # path = os.path.join(object_dir, object_id, f"{object_type}.json")
    path_backup = f"{path}.backup"
    if os.path.exists(path_backup):
        print(
            f"WARNING: Not processing {object_type} '{object_id}' since '{path_backup}' already exists"
        )
    else:
        print(f"Processing {object_type} '{object_id}'")
        shutil.copyfile(path, path_backup)
        object_dct = io_utils.read_file(path)
        rewrite_object_func(object_dct, object_dir)
        io_utils.write_file(path, object_dct)


def rewrite_models(input_dir, rewrite_model):
    """Rewrite the models"""
    input_dir = os.path.join(input_dir, "models")
    models_dct = io_utils.read_file(os.path.join(input_dir, "models.json"))
    model_names = models_dct["mlflow"]["models"]

    print(f"Processing {len(model_names)} models")
    for model_name in model_names:
        _rewrite_object(input_dir, model_name, "model", rewrite_model)


def rewrite_experiments(input_dir, rewrite_experiment):
    """Rewrite the experiments"""
    input_dir = os.path.join(input_dir, "experiments")
    experiments_dct = io_utils.read_file(os.path.join(input_dir, "experiments.json"))
    experiments = experiments_dct["mlflow"]["experiments"]
    experiment_ids = [exp["id"] for exp in experiments]

    print(f"Processing {len(experiment_ids)} experiments")
    for experiment_id in experiment_ids:
        _rewrite_object(input_dir, experiment_id, "experiment", rewrite_experiment)


def rewrite_all(input_dir, custom_rewriters_module):
    """Rewrite all"""
    rewrite_model, rewrite_experiment = _load_module(custom_rewriters_module)
    rewrite_models(input_dir, rewrite_model)
    rewrite_experiments(input_dir, rewrite_experiment)


@click.command()
@click.option(
    "--input-dir",
    help="Export directory of export-models or export-all.",
    type=str,
    required=True,
)
@click.option(
    "--custom-rewriters-module",
    help="""
Python file containing user-provided custom model and experiment rewrite logic.
Module expects 2 methods: rewrite_models(dct) and rewrite_experiments(dct).
""",
    type=str,
    required=True,
)
def main(input_dir, custom_rewriters_module):
    print("Options:")
    for k, v in locals().items():
        print(f"  {k}: {v}")
    rewrite_all(input_dir, custom_rewriters_module)


if __name__ == "__main__":
    main()
