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

# Copyright (c) 2022, 2023 Oracle and/or its affiliates.
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/

import json
from typing import Dict, Union
from ads.common.auth import AuthContext, create_signer
from ads.common.oci_client import OCIClientFactory
from ads.model.deployment import ModelDeployment, ModelDeploymentInfrastructure
from ads.model.deployment.model_deployment_runtime import (
    ModelDeploymentCondaRuntime,
    ModelDeploymentContainerRuntime,
)
from ads.opctl.backend.base import Backend, RuntimeFactory


class ModelDeploymentBackend(Backend):
    def __init__(self, config: Dict) -> None:
        """
        Initialize a ModelDeployment object given config dictionary.

        Parameters
        ----------
        config: dict
            dictionary of configurations
        """
        self.config = config
        self.oci_auth = create_signer(
            config["execution"].get("auth"),
            config["execution"].get("oci_config", None),
            config["execution"].get("oci_profile", None),
        )
        self.auth_type = config["execution"].get("auth")
        self.profile = config["execution"].get("oci_profile", None)
        self.client = OCIClientFactory(**self.oci_auth).data_science

    def init(
        self,
        uri: Union[str, None] = None,
        overwrite: bool = False,
        runtime_type: Union[str, None] = None,
        **kwargs: Dict,
    ) -> Union[str, None]:
        """Generates a starter YAML specification for a Data Science Job.

        Parameters
        ----------
        overwrite: (bool, optional). Defaults to False.
            Overwrites the result specification YAML if exists.
        uri: (str, optional), Defaults to None.
            The filename to save the resulting specification template YAML.
        runtime_type: (str, optional). Defaults to None.
                The resource runtime type.
        **kwargs: Dict
            The optional arguments.

        Returns
        -------
        Union[str, None]
            The YAML specification for the given resource if `uri` was not provided.
            `None` otherwise.
        """
        with AuthContext(auth=self.auth_type, profile=self.profile):
            # define a job
            object = (
                ModelDeployment()
                .with_display_name("Model Deployment")
                .with_description("The model deployment description")
                .with_freeform_tags(key1="value1")
                .with_infrastructure(
                    ModelDeploymentInfrastructure(
                        **(self.config.get("infrastructure", {}) or {})
                    ).init()
                )
                .with_runtime(
                    ModelDeploymentRuntimeFactory.get_runtime(
                        key=runtime_type or ModelDeploymentCondaRuntime().type
                    ).init()
                )
            )

            note = (
                "# This YAML specification was auto generated by the `ads opctl init` command.\n"
                "# The more details about the jobs YAML specification can be found in the ADS documentation:\n"
                "# https://accelerated-data-science.readthedocs.io/en/latest/user_guide/model_registration/model_deploy_byoc.html \n\n"
            )

            return object.to_yaml(
                uri=uri,
                overwrite=overwrite,
                note=note,
                filter_by_attribute_map=True,
                **kwargs,
            )

    def apply(self) -> None:
        """
        Deploy model deployment from YAML.
        """
        wait_for_completion = self.config["execution"].get("wait_for_completion")
        max_wait_time = self.config["execution"].get("max_wait_time")
        poll_interval = self.config["execution"].get("poll_interval")
        with AuthContext(auth=self.auth_type, profile=self.profile):
            model_deployment = ModelDeployment.from_dict(self.config)
            model_deployment.deploy(
                wait_for_completion=wait_for_completion,
                max_wait_time=max_wait_time,
                poll_interval=poll_interval,
            )
            print("Model Deployment id: ", model_deployment.model_deployment_id)

    def delete(self) -> None:
        """
        Delete model deployment from OCID.
        """
        model_deployment_id = self.config["execution"].get("run_id")
        wait_for_completion = self.config["execution"].get("wait_for_completion")
        max_wait_time = self.config["execution"].get("max_wait_time")
        poll_interval = self.config["execution"].get("poll_interval")
        with AuthContext(auth=self.auth_type, profile=self.profile):
            model_deployment = ModelDeployment.from_id(model_deployment_id)
            model_deployment.delete(
                wait_for_completion=wait_for_completion,
                max_wait_time=max_wait_time,
                poll_interval=poll_interval,
            )
            print(
                f"Model Deployment {model_deployment.model_deployment_id} has been deleted."
            )

    def activate(self) -> None:
        """
        Activate model deployment from OCID.
        """
        model_deployment_id = self.config["execution"].get("run_id")
        wait_for_completion = self.config["execution"].get("wait_for_completion")
        max_wait_time = self.config["execution"].get("max_wait_time")
        poll_interval = self.config["execution"].get("poll_interval")
        with AuthContext(auth=self.auth_type, profile=self.profile):
            model_deployment = ModelDeployment.from_id(model_deployment_id)
            model_deployment.activate(
                wait_for_completion=wait_for_completion,
                max_wait_time=max_wait_time,
                poll_interval=poll_interval,
            )
            print(
                f"Model Deployment {model_deployment.model_deployment_id} has been activated."
            )

    def deactivate(self) -> None:
        """
        Deactivate model deployment from OCID.
        """
        model_deployment_id = self.config["execution"].get("run_id")
        wait_for_completion = self.config["execution"].get("wait_for_completion")
        max_wait_time = self.config["execution"].get("max_wait_time")
        poll_interval = self.config["execution"].get("poll_interval")
        with AuthContext(auth=self.auth_type, profile=self.profile):
            model_deployment = ModelDeployment.from_id(model_deployment_id)
            model_deployment.deactivate(
                wait_for_completion=wait_for_completion,
                max_wait_time=max_wait_time,
                poll_interval=poll_interval,
            )
            print(
                f"Model Deployment {model_deployment.model_deployment_id} has been deactivated."
            )

    def watch(self) -> None:
        """
        Watch Model Deployment from OCID.
        """
        model_deployment_id = self.config["execution"].get("run_id")
        log_type = self.config["execution"].get("log_type")
        interval = self.config["execution"].get("interval")
        log_filter = self.config["execution"].get("log_filter")
        with AuthContext(auth=self.auth_type, profile=self.profile):
            model_deployment = ModelDeployment.from_id(model_deployment_id)
            model_deployment.watch(
                log_type=log_type, interval=interval, log_filter=log_filter
            )

    def predict(self) -> None:
        ocid = self.config["execution"].get("ocid")
        data = self.config["execution"].get("payload")
        model_name = self.config["execution"].get("model_name")
        model_version = self.config["execution"].get("model_version")
        with AuthContext(auth=self.auth_type, profile=self.profile):
            model_deployment = ModelDeployment.from_id(ocid)
            try:
                data = json.loads(data)
            except:
                pass
            print(
                model_deployment.predict(
                    data=data, model_name=model_name, model_version=model_version
                )
            )


class ModelDeploymentRuntimeFactory(RuntimeFactory):
    """Mode Deployment runtime factory."""

    _MAP = {
        ModelDeploymentCondaRuntime().type: ModelDeploymentCondaRuntime,
        ModelDeploymentContainerRuntime().type: ModelDeploymentContainerRuntime,
    }
