import logging
import os
from pathlib import Path

import typer

from solarnet.tasks.dataset import make_dataset
from solarnet.tasks.download_dataset import download_dataset
from solarnet.utils.log import init_log, set_log_level
from solarnet.utils.yaml import load_yaml

init_log()
set_log_level(logging.WARNING)
logger = logging.getLogger()
app = typer.Typer(context_settings={"help_option_names": ["-h", "--help"]})
os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "gloo"


@app.command("download")
def download_command(
    dataset: str = typer.Argument("sdo-benchmark"),
    destination: Path = typer.Argument(Path("./")),
    verbose: bool = typer.Option(False, "--verbose", "-v"),
):
    if verbose:
        set_log_level(logging.INFO)

    download_dataset(dataset, destination)


@app.command("sdo-dataset")
def sdo_dataset_command(
    config_file: Path,
    verbose: bool = typer.Option(False, "--verbose", "-v"),
):
    if verbose:
        set_log_level(logging.INFO)

    params = load_yaml(config_file)

    make_dataset(params)


# Command to add options before the command (-v train ...)
@app.callback()
def main(verbose: bool = typer.Option(False, "--verbose", "-v")):
    if verbose:
        set_log_level(logging.INFO)


if __name__ == "__main__":
    app()
