#!/usr/bin/env python
# -*- coding: utf-8; -*-

# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
"""
APIs to interact with Oracle's Model Deployment service.

There are three main classes: ModelDeployment, ModelDeploymentDetails, ModelDeployer.

One creates a ModelDeployment and deploys it under the umbrella of the ModelDeployer class. This way
multiple ModelDeployments can be unified with one ModelDeployer. The ModelDeployer class also serves
as the interface to all the deployments. ModelDeploymentDetails holds information about the particular
details of a particular deployment, such as how many instances, etc. In this way multiple, independent
ModelDeployments with the same details can be created using the ModelDeployer class.

Examples
--------
>>> from model_deploy.model_deployer import ModelDeployer, ModelDeploymentDetails
>>> deployer = ModelDeployer("model_dep_conf.yaml")
>>> deployment_properties = ModelDeploymentProperties(
...             'ocid1.datasciencemodel.ocn.reg.xxxxxxxxxxxxxxxxxxxxxxxxx')
...                 .with_prop('display_name', "My model display name")
...                 .with_prop("project_id", project_id)
...                 .with_prop("compartment_id", compartment_id)
...                 .with_instance_configuration(
...                     config={"INSTANCE_SHAPE":"VM.Standard.E3.Flex",
...                             "INSTANCE_COUNT":"1",
...                             "bandwidth_mbps":10,
...                             "memory_in_gbs":10,
...                             "ocpus":2}
...                  ).build()
>>> deployment_info = deployer.deploy(deployment_properties,
...             max_wait_time=600, poll_interval=15)
>>> print(deployment_info.model_deployment_id)
>>> print(deployment_info.workflow_req_id)
>>> print(deployment_info.url)
>>> deployer.list_deployments() # Optionally pass in a status
"""
from typing import Dict, Union
import warnings

import oci.pagination
import pandas as pd
from ads.common.auth import default_signer
from oci.data_science.data_science_client import DataScienceClient
from oci.data_science import DataScienceClientCompositeOperations

from .common import utils
from .common.utils import OCIClientManager, State
from .model_deployment import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIME, ModelDeployment
from .model_deployment_properties import ModelDeploymentProperties


class ModelDeployer:
    """ModelDeployer is the class responsible for deploying the ModelDeployment

    Attributes
    ----------
    config : dict
        ADS auth dictionary for OCI authentication.
    ds_client : DataScienceClient
        data science client
    ds_composite_client : DataScienceCompositeClient
        composite data science client

    Methods
    -------
    deploy(model_deployment_details, **kwargs)
        Deploy the model specified by `model_deployment_details`.
    get_model_deployment(model_deployment_id:str)
        Get the ModelDeployment specified by `model_deployment_id`.
    get_model_deployment_state(model_deployment_id)
        Get the state of the current deployment specified by id.
    delete(model_deployment_id, **kwargs)
        Remove the model deployment specified by the id or Model Deployment Object
    list_deployments(status)
        lists the model deployments associated with current compartment and data
        science client
    show_deployments(status)
        shows the deployments filtered by `status` in a Dataframe
    """

    def __init__(self, config: dict = None, ds_client: DataScienceClient = None):
        """Initializes model deployer.

        Parameters
        ----------
        config : dict, optional
            ADS auth dictionary for OCI authentication.
            This can be generated by calling ads.common.auth.api_keys() or ads.common.auth.resource_principal().
            If this is None, ads.common.default_signer(client_kwargs) will be used.

         ds_client: oci.data_science.data_science_client.DataScienceClient
            The Oracle DataScience client.
        """
        if config:
            warnings.warn(
                "`config` will be deprecated in 3.0.0 and will be ignored now. Please use `ads.set_auth()` to config the auth information."
            )

        self.ds_client = None
        self.ds_composite_client = None

        if ds_client:
            self.ds_client = ds_client
            self.ds_composite_client = DataScienceClientCompositeOperations(
                self.ds_client
            )

        if not (self.ds_client != None and self.ds_composite_client != None):
            self.client_manager = OCIClientManager()
            self.ds_client = self.client_manager.ds_client
            self.ds_composite_client = self.client_manager.ds_composite_client

    def deploy(
        self,
        properties: Union[ModelDeploymentProperties, Dict] = None,
        wait_for_completion: bool = True,
        max_wait_time: int = DEFAULT_WAIT_TIME,
        poll_interval: int = DEFAULT_POLL_INTERVAL,
        **kwargs,
    ) -> ModelDeployment:
        """Deploys a model.

        Parameters
        ----------
        properties : ModelDeploymentProperties or dict
            Properties to deploy the model.
            Properties can be None when kwargs are used for specifying properties.
        wait_for_completion : bool
            Flag set for whether to wait for deployment to complete before proceeding.
            Optional, defaults to True.
        max_wait_time : int
            Maximum amount of time to wait in seconds. Optional, defaults to 1200.
            Negative value implies infinite wait time.
        poll_interval : int
            Poll interval in seconds. Optional, defaults to 30.
        kwargs :
            Keyword arguments for initializing ModelDeploymentProperties.
            See ModelDeploymentProperties() for details.

        Returns
        -------
        ModelDeployment
            A ModelDeployment instance.
        """
        warnings.warn(
            "Make sure to provide `compartment_id`, `project_id`, `instance_count`, `instance_shape`, `memory_in_gbs` and `ocpus` either through `kwargs` or `properties`"
        )
        model_deployment = ModelDeployment(
            properties,
            **kwargs,
        )
        return model_deployment.deploy(
            wait_for_completion, max_wait_time, poll_interval
        )

    def deploy_from_model_uri(
        self,
        model_uri: str,
        properties: Union[ModelDeploymentProperties, Dict] = None,
        wait_for_completion: bool = True,
        max_wait_time: int = DEFAULT_WAIT_TIME,
        poll_interval: int = DEFAULT_POLL_INTERVAL,
        **kwargs,
    ) -> ModelDeployment:
        """Deploys a model.

        Parameters
        ----------
        model_uri : str
            uri to model files, can be local or in cloud storage
        properties : ModelDeploymentProperties or dict
            Properties to deploy the model.
            Properties can be None when kwargs are used for specifying properties.
        wait_for_completion : bool
            Flag set for whether to wait for deployment to complete before proceeding.
            Defaults to True
        max_wait_time : int
            Maximum amount of time to wait in seconds (Defaults to 1200).
            Negative implies infinite wait time.
        poll_interval : int
            Poll interval in seconds (Defaults to 30).
        kwargs :
            Keyword arguments for initializing ModelDeploymentProperties

        Returns
        -------
        ModelDeployment
            A ModelDeployment instance

        """
        warnings.warn(
            "Make sure to provide `compartment_id`, `project_id`, `instance_count`, `instance_shape`, `memory_in_gbs` and `ocpus` either through `kwargs` or `properties`"
        )
        if properties:
            model_id = self.client_manager.prepare_artifact(
                model_uri=model_uri, properties=properties
            )
            properties.model_deployment_configuration_details.model_configuration_details.model_id = (
                model_id
            )
        else:
            model_id = self.client_manager.prepare_artifact(
                model_uri=model_uri, properties=kwargs
            )
            kwargs["model_id"] = model_id

        return self.deploy(
            properties,
            wait_for_completion=wait_for_completion,
            max_wait_time=max_wait_time,
            poll_interval=poll_interval,
            **kwargs,
        )

    def update(
        self,
        model_deployment_id: str,
        properties: ModelDeploymentProperties = None,
        wait_for_completion: bool = True,
        max_wait_time: int = DEFAULT_WAIT_TIME,
        poll_interval: int = DEFAULT_POLL_INTERVAL,
        **kwargs,
    ) -> ModelDeployment:
        """Updates an existing model deployment.

        Parameters
        ----------
        model_deployment_id : str
            Model deployment OCID.
        properties : ModelDeploymentProperties
            An instance of ModelDeploymentProperties or dict to initialize the ModelDeploymentProperties.
            Defaults to None.
        wait_for_completion : bool
            Flag set for whether to wait for deployment to complete before proceeding.
            Defaults to True.
        max_wait_time : int
            Maximum amount of time to wait in seconds (Defaults to 1200).
        poll_interval : int
            Poll interval in seconds (Defaults to 30).
        kwargs :
            Keyword arguments for initializing ModelDeploymentProperties.


        Returns
        -------
        ModelDeployment
            A ModelDeployment instance

        """

        model_deployment = self.get_model_deployment(model_deployment_id)
        # Deployment properties will be refreshed within model_deployment.update() when update is done.
        return model_deployment.update(
            properties,
            wait_for_completion,
            max_wait_time=max_wait_time,
            poll_interval=poll_interval,
            **kwargs,
        )

    def get_model_deployment(self, model_deployment_id: str) -> ModelDeployment:
        """Gets a ModelDeployment by OCID.

        Parameters
        ----------
        model_deployment_id : str
            Model deployment OCID

        Returns
        -------
        ModelDeployment
            A ModelDeployment instance

        """

        try:
            oci_model_deployment_object = self.ds_client.get_model_deployment(
                model_deployment_id
            ).data
            model_deployment_object = ModelDeployment(
                oci_model_deployment_object,
            )
            model_deployment_object.set_spec(
                model_deployment_object.CONST_ID, oci_model_deployment_object.id
            )
            return model_deployment_object.sync()
        except Exception as e:
            utils.get_logger().error(
                "Getting model deployment failed with error: %s", e
            )
            raise e

    def get_model_deployment_state(self, model_deployment_id: str) -> State:
        """Gets the state of a deployment specified by OCID

        Parameters
        ----------
        model_deployment_id : str
            Model deployment OCID

        Returns
        -------
        str
            The state of the deployment
        """

        model_deployment = self.get_model_deployment(model_deployment_id)
        return model_deployment.state

    def delete(
        self,
        model_deployment_id,
        wait_for_completion: bool = True,
        max_wait_time: int = DEFAULT_WAIT_TIME,
        poll_interval: int = DEFAULT_POLL_INTERVAL,
    ) -> ModelDeployment:
        """Deletes the model deployment specified by OCID.

        Parameters
        ----------
        model_deployment_id : str
            Model deployment OCID.
        wait_for_completion : bool
            Wait for deletion to complete. Defaults to True.
        max_wait_time : int
            Maximum amount of time to wait in seconds (Defaults to 600).
            Negative implies infinite wait time.
        poll_interval : int
            Poll interval in seconds (Defaults to 60).

        Returns
        -------
        A ModelDeployment instance that was deleted

        """
        try:
            model_deployment_object = self.get_model_deployment(model_deployment_id)
            return model_deployment_object.delete(
                wait_for_completion, max_wait_time, poll_interval
            )
        except Exception as e:
            utils.get_logger().error(
                "Deleting model deployment failed with error: %s", format(e)
            )
            raise e

    def list_deployments(self, status=None, compartment_id=None, **kwargs) -> list:
        """Lists the model deployments associated with current compartment and data science client

        Parameters
        ----------
        status : str
            Status of deployment. Defaults to None.
        compartment_id : str
            Target compartment to list deployments from.
            Defaults to the compartment set in the environment variable "NB_SESSION_COMPARTMENT_OCID".
            If "NB_SESSION_COMPARTMENT_OCID" is not set, the root compartment ID will be used.
            An ValueError will be raised if root compartment ID cannot be determined.
        kwargs :
            The values are passed to oci.data_science.DataScienceClient.list_model_deployments.

        Returns
        -------
        list
            A list of ModelDeployment objects.

        Raises
        ------
        ValueError
            If compartment_id is not specified and cannot be determined from the environment.

        """
        if not compartment_id:
            compartment_id = self.client_manager.default_compartment_id()
        if not compartment_id:
            raise ValueError(
                "Unable to determine compartment ID from environment. Specify compartment_id."
            )

        if isinstance(status, State):
            status = status.name

        if status is not None:
            kwargs["lifecycle_state"] = status

        # https://oracle-cloud-infrastructure-python-sdk.readthedocs.io/en/latest/api/pagination.html#module-oci.pagination
        deployments = oci.pagination.list_call_get_all_results(
            self.ds_client.list_model_deployments, compartment_id, **kwargs
        ).data

        result = []
        for deployment in deployments:
            model_deployment = ModelDeployment(deployment)
            model_deployment.set_spec(model_deployment.CONST_ID, deployment.id)
            result.append(model_deployment.sync())

        return result

    def show_deployments(
        self,
        status=None,
        compartment_id=None,
    ) -> pd.DataFrame:
        """Returns the model deployments associated with current compartment and data science client
            as a Dataframe that can be easily visualized

        Parameters
        ----------
        status : str
            Status of deployment. Defaults to None.
        compartment_id : str
            Target compartment to list deployments from.
            Defaults to the compartment set in the environment variable "NB_SESSION_COMPARTMENT_OCID".
            If "NB_SESSION_COMPARTMENT_OCID" is not set, the root compartment ID will be used.
            An ValueError will be raised if root compartment ID cannot be determined.

        Returns
        -------
        DataFrame
            pandas Dataframe containing information about the ModelDeployments

        Raises
        ------
        ValueError
            If compartment_id is not specified and cannot be determined from the environment.
        """
        if not compartment_id:
            compartment_id = self.client_manager.default_compartment_id()
        if not compartment_id:
            raise ValueError(
                "Unable to determine compartment ID from environment. Specify compartment_id."
            )

        if type(status) == str or status == None:
            status = State._from_str(status)
        model_deployments = self.ds_client.list_model_deployments(compartment_id).data
        display = pd.DataFrame()
        ids, urls, status_list = [], [], []
        for oci_model_deployment_object in model_deployments:
            state_of_model = State._from_str(
                oci_model_deployment_object.lifecycle_state
            )
            if status == State.UNKNOWN or status.name == state_of_model.name:
                model_dep_id = oci_model_deployment_object.id
                model_dep_url = oci_model_deployment_object._model_deployment_url
                model_dep_state = oci_model_deployment_object.lifecycle_state
                ids.append(model_dep_id)
                urls.append(model_dep_url)
                status_list.append(model_dep_state)
        display["model_id"] = ids
        display["deployment_url"] = urls
        display["current_state"] = status_list
        return display
