import argparse
import logging
import os
import sys

from ludwig.globals import MODEL_FILE_NAME, MODEL_HYPERPARAMETERS_FILE_NAME, model_weights_exist
from ludwig.utils.print_utils import get_logging_level_registry
from ludwig.utils.upload_utils import HuggingFaceHub

logger = logging.getLogger(__name__)


def get_upload_registry():
    return {
        "hf_hub": HuggingFaceHub,
    }


def upload_cli(
    service: str,
    repo_id: str,
    model_path: str,
    repo_type: str = "model",
    private: bool = False,
    commit_message: str = "Upload trained [Ludwig](https://ludwig.ai/latest/) model weights",
    commit_description: str | None = None,
    dataset_file: str | None = None,
    dataset_name: str | None = None,
    **kwargs,
) -> None:
    """Create an empty repo on the HuggingFace Hub and upload trained model artifacts to that repo.

    Args:
        service (`str`):
            Name of the hosted model service to push the trained artifacts to.
            Currently only `hf_hub` is supported.
        repo_id (`str`):
            A namespace (user or an organization) and a repo name separated
            by a `/`.
        model_path (`str`):
            The path of the saved model. This is the parent-folder of the folder
            where the 'model_weights' folder and the 'model_hyperparameters.json' file
            are stored.
        private (`bool`, *optional*, defaults to `False`):
            Whether the model repo should be private.
        repo_type (`str`, *optional*):
            Set to `"dataset"` or `"space"` if uploading to a dataset or
            space, `None` or `"model"` if uploading to a model. Default is
            `None`.
        commit_message (`str`, *optional*):
            The summary / title / first line of the generated commit. Defaults to:
            `f"Upload {path_in_repo} with huggingface_hub"`
        commit_description (`str` *optional*):
            The description of the generated commit
        dataset_file (`str`, *optional*):
            The path to the dataset file.
        dataset_name (`str`, *optional*):
            The name of the dataset.
    """
    model_service = get_upload_registry().get(service, "hf_hub")
    hub: HuggingFaceHub = model_service()
    if model_weights_exist(os.path.join(model_path, MODEL_FILE_NAME)) and os.path.exists(
        os.path.join(model_path, MODEL_FILE_NAME, MODEL_HYPERPARAMETERS_FILE_NAME)
    ):
        experiment_path = model_path
    elif model_weights_exist(model_path) and os.path.exists(os.path.join(model_path, MODEL_HYPERPARAMETERS_FILE_NAME)):
        experiment_path = os.path.normpath(os.path.join(model_path, ".."))
    else:
        raise ValueError(
            f"Can't find model weights and '{MODEL_HYPERPARAMETERS_FILE_NAME}' either at "
            f"'{model_path}' or at '{model_path}/model'"
        )
    hub.upload(
        repo_id=repo_id,
        model_path=experiment_path,
        repo_type=repo_type,
        private=private,
        commit_message=commit_message,
        commit_description=commit_description,
        dataset_file=dataset_file,
        dataset_name=dataset_name,
    )


def cli(sys_argv):
    parser = argparse.ArgumentParser(
        description="This script pushes a trained model to a hosted model repository service",
        prog="ludwig upload",
        usage="%(prog)s [options]",
    )

    # ---------------
    # Required parameters
    # ---------------
    parser.add_argument(
        "service",
        help="Name of the model repository service.",
        default="hf_hub",
        choices=["hf_hub"],
    )

    parser.add_argument(
        "-r",
        "--repo_id",
        help="Name of the repo. This will be created if it doesn't exist. Format: username/repo_name",
        required=True,
    )

    parser.add_argument("-m", "--model_path", help="Path of the trained model on disk", required=True)

    # ---------------
    # Optional parameters
    # ---------------
    parser.add_argument("-p", "--private", help="Make the repo private", default=False, choices=[True, False])

    parser.add_argument(
        "-t", "--repo_type", help="Type of repo", default="model", choices=["model", "space", "dataset"]
    )

    parser.add_argument(
        "-c",
        "--commit_message",
        help="The summary / title / first line of the generated commit.",
        default="Upload trained [Ludwig](https://ludwig.ai/latest/) model weights",
    )

    parser.add_argument("-d", "--commit_description", help="The description of the generated commit", default=None)

    parser.add_argument(
        "-l",
        "--logging_level",
        default="info",
        help="The level of logging to use",
        choices=["critical", "error", "warning", "info", "debug", "notset"],
    )

    parser.add_argument("-df", "--dataset_file", help="The location of the dataset file", default=None)
    parser.add_argument(
        "-dn", "--dataset_name", help="(Optional) The name of the dataset in the Provider", default=None
    )

    args = parser.parse_args(sys_argv)

    args.logging_level = get_logging_level_registry()[args.logging_level]
    logging.getLogger("ludwig").setLevel(args.logging_level)
    global logger
    logger = logging.getLogger("ludwig.upload")

    upload_cli(**vars(args))


if __name__ == "__main__":
    cli(sys.argv[1:])