import requests
from datetime import datetime
from typing import Dict, Optional

import pandas as pd
from mlflow.tracking import MlflowClient

from mindsdb.integrations.libs.base import BaseMLEngine


class MLflowHandler(BaseMLEngine):
    """
    The MLflow integration engine needs to have a working connection to MLFlow. For this:
        - All models to use should be previously served
        - An MLflow server should be running, to access its model registry

    Example:
        1. Run `mlflow server -p 5001 --backend-store-uri sqlite:////path/to/mlflow.db --default-artifact-root ./artifacts --host 0.0.0.0`
        2. Run `mlflow models serve --model-uri ./model_path`
        3. Run MindsDB

    Note: above, `artifacts` is a folder to store artifacts for new experiments that do not specify an artifact store.
    """  # noqa

    name = "mlflow"

    def create(
        self,
        target: str,
        df: Optional[pd.DataFrame] = None,
        args: Optional[Dict] = None,
    ) -> None:
        args = args["using"]  # ignore the rest of the problem definition
        connection = MlflowClient(args["mlflow_server_url"], args["mlflow_server_path"])
        model_name = args["model_name"]
        mlflow_models = [model.name for model in connection.search_registered_models()]

        if model_name not in mlflow_models:
            raise Exception(
                f"Error: model '{model_name}' not found in mlflow. Check serving and try again."
            )

        args["target"] = target
        self._check_model_url(args["predict_url"])
        self.model_storage.json_set("args", args)

    def predict(self, df, args=None):
        args = self.model_storage.json_get("args")  # override any incoming args for now
        self._check_model_url(args["predict_url"])
        resp = requests.post(
            args["predict_url"],
            data=df.to_json(orient="records"),
            headers={"content-type": "application/json; format=pandas-records"},
        )
        answer = resp.json()
        predictions = pd.DataFrame({args["target"]: answer})
        return predictions

    def describe(self, key: Optional[str] = None) -> pd.DataFrame:
        if key == "info":
            args = self.model_storage.json_get("args")
            connection = MlflowClient(
                args["mlflow_server_url"], args["self.mlflow_server_path"]
            )
            models = {
                model.name: model for model in connection.search_registered_models()
            }
            model = models[key]
            latest_version = model.latest_versions[-1]
            description = {
                "NAME": [model.name],
                "USER_DESCRIPTION": [model.description],
                "LAST_STATUS": [latest_version.status],
                "CREATED_AT": [
                    datetime.fromtimestamp(model.creation_timestamp // 1000).strftime(
                        "%m/%d/%Y, %H:%M:%S"
                    )
                ],
                "LAST_UPDATED": [
                    datetime.fromtimestamp(
                        model.last_updated_timestamp // 1000
                    ).strftime("%m/%d/%Y, %H:%M:%S")
                ],
                "TAGS": [model.tags],
                "LAST_RUN_ID": [latest_version.run_id],
                "LAST_SOURCE_PATH": [latest_version.source],
                "LAST_USER_ID": [latest_version.user_id],
                "LAST_VERSION": [latest_version.version],
            }
            return pd.DataFrame.from_dict(description)
        else:
            tables = ["info"]
            return pd.DataFrame(tables, columns=["tables"])

    @staticmethod
    def _check_model_url(url):
        """try post without data, check status code not in (not_found, method_not_allowed)"""
        try:
            resp = requests.post(url)
            if resp.status_code in (404, 405):
                raise Exception(
                    f"Model url is incorrect, status_code: {resp.status_code}"
                )
        except requests.RequestException as e:
            raise Exception(f"Model url is incorrect: {str(e)}")
