from pathlib import Path

import mlflow


REMOTE_TRACKING_SCHEMES = ("file://", "http://", "https://", "databricks", "sqlite:")


def normalize_tracking_uri(tracking_uri):
    if not tracking_uri:
        raise ValueError("TRACKING_URI must be provided in the MLflow configuration.")

    if tracking_uri.startswith(REMOTE_TRACKING_SCHEMES):
        return tracking_uri

    tracking_path = Path(tracking_uri).expanduser().resolve()
    tracking_path.mkdir(parents=True, exist_ok=True)
    return tracking_path.as_uri()


def setup_mlflow(mlflow_config, config_file_path, optuna=False):
    config = dict(mlflow_config)
    tracking_uri = normalize_tracking_uri(config["TRACKING_URI"])
    mlflow.set_tracking_uri(tracking_uri)

    experiment_or_study_name = (
        config["STUDY_NAME"] if optuna else config["EXPERIMENT_NAME"]
    )
    existing_experiment = mlflow.get_experiment_by_name(experiment_or_study_name)
    if existing_experiment:
        config["EXPERIMENT_ID"] = existing_experiment.experiment_id
    else:
        config["EXPERIMENT_ID"] = mlflow.create_experiment(experiment_or_study_name)

    config["TRACKING_URI"] = tracking_uri
    config["config_file_path"] = config_file_path
    return config


def log_params(params, parent_key=""):
    for k, v in params.items():
        new_key = f"{parent_key} / {k}" if parent_key else k
        if isinstance(v, dict):
            for key, value in v.items():
                mlflow.log_param(f"{key}", value)
        elif isinstance(v, list):
            for i, item in enumerate(v):
                if isinstance(item, dict):
                    log_params(item, f"{new_key} / {i+1}")
                else:
                    mlflow.log_param(f"{new_key} / {i+1}", item)
        else:
            mlflow.log_param(new_key, v)
