import logging

from typing import Any

import mlflow
import numpy as np
import pandas as pd

from src.logger import setup_logger

logger = setup_logger(__name__)


def create_mlflow_experiment(experiment_name: str, artifact_location: str, tags: dict[str, Any]) -> str:
    """Creates a new experiment with the given name, artifact location, and tags."""

    try:
        experiment_id = mlflow.create_experiment(name=experiment_name, artifact_location=artifact_location, tags=tags)
        logging.info(f"Created new experiment with ID: {experiment_id}")
    except mlflow.exceptions.MlflowException as e:
        if "already exists" in str(e):
            logging.warning(f"Experiment {experiment_name} already exists")
            experiment_id = mlflow.get_experiment_by_name(experiment_name).experiment_id
        else:
            logging.error(f"Failed to create or retrieve experiment: {e}")
            raise

    return experiment_id


def get_mlflow_experiment(experiment_name: str = None, experiment_id: str = None) -> mlflow.entities.Experiment:
    """Get the MLflow experiment by name or id.

    Args:
        experiment_name (str, optional): Name of the experiment. Defaults to None.
        experiment_id (str, optional): ID of the experiment. Defaults to None.

    Returns:
        mlflow.entities.Experiment: Experiment object
    """

    if experiment_name:
        experiment = mlflow.get_experiment_by_name(experiment_name)
    elif experiment_id:
        experiment = mlflow.get_experiment(experiment_id)
    else:
        raise ValueError("Either experiment name or experiment id should be provided")

    return experiment


def model_signature(model, X: pd.DataFrame, y: pd.Series) -> mlflow.models.signature.ModelSignature:
    for i, col in enumerate(X.columns):
        X[i, col] = np.nan

    signature = mlflow.models.infer_signature(X, y)
    return signature


def mlflow_logged_model(logged_model_uri: str) -> mlflow.pyfunc.PyFuncModel:
    return mlflow.pyfunc.load_model(logged_model_uri)
