"""Contains Mixins for integrating OCI data models
"""
import inspect
import json
import logging
import os
import re
from datetime import date, datetime
from typing import Callable, Optional, Union

import oci
from dateutil import tz
from dateutil.parser import parse
from oci._vendor import six

from ads.common.auth import default_signer
from ads.common.oci_client import OCIClientFactory

logger = logging.getLogger(__name__)


class OCIClientMixin:
    """Mixin class for representing OCI resource/service with OCI client.

    Most OCI requests requires a client for making the connection.
    Usually the same client will be used for the requests related the same resource/service type.
    This Mixin adds a client property to simplify accessing the client.
    The actual client will be initialize lazily so that it is not required for a sub-class
    To use the client, sub-class should override the init_client() method pass in the "client" keyword argument.
    For example:

    @staticmethod
    def init_client(**kwargs) -> oci.logging.LoggingManagementClient:
        return cls._init_client(client=oci.logging.LoggingManagementClient, **kwargs)

    Authentication
    --------------
    The init_client() method is a class method. The call to this method may come from an instance or a class.
    When the method is called,
        the client will be initialized with the credentials passed in by the "auth" keyword argument.
    In order to allow the user to specify "auth", all methods using init_client() should expose "auth" as an argument.
    If "auth" is not specified,
        The default auth configured at ADS level with ads.set_auth() will be used.
    """

    @classmethod
    def init_client(cls, **kwargs):
        """Initializes the OCI client specified in the "client" keyword argument
        Sub-class should override this method and call cls._init_client(client=OCI_CLIENT)
        """
        return cls._init_client(**kwargs)

    @classmethod
    def _init_client(cls, auth: dict = None, **kwargs):
        """Initializes the OCI client specified in the "client" keyword argument

        Parameters
        ----------
        auth : dict
            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.
        **kwargs :
            client is required for specifying the class of OCI client
            Additional keyword arguments will be passed to ads.common.default_signer() when auth is None.

        Returns
        -------
            An instance of OCI client.

        """
        client_args = dict(retry_strategy=oci.retry.DEFAULT_RETRY_STRATEGY)
        client_class = kwargs.pop("client", None)
        if not client_class:
            raise ValueError("client is not specified in keyword argument.")

        client_args.update(kwargs)

        # If auth is empty at this time,
        # it mean the class authentication is not configured at the time when the instance was created, or
        # the call is from a class without auth configured.
        # The ADS default will be used.
        if not auth:
            auth = default_signer(client_args)

        return OCIClientFactory(
            **auth,
        ).create_client(client_class)

    def __init__(self) -> None:
        """Initializes a service/resource with OCI client as a property."""
        super().__init__()
        self._client = None

    @property
    def client(self):
        """OCI client"""
        if self._client is None:
            self._client = self.init_client()
        return self._client


class OCISerializableMixin(OCIClientMixin):
    """Mixin class containing OCI serialization/de-serialization methods.
    These methods are copied and modified from the OCI BaseClient.

    """

    type_mappings = None

    def serialize(self):
        """Serialize the model to a dictionary that is ready to be send to OCI API.

        Returns
        -------
        dict
            A dictionary that is ready to be send to OCI API.

        """
        return self.client.base_client.sanitize_for_serialization(self)

    @staticmethod
    def _parse_kwargs(attribute_map: dict, **kwargs):
        """Parse kwargs to make all key in camel format."""
        parsed_kwargs = {}
        for key, val in kwargs.items():
            if key in attribute_map:
                parsed_kwargs[attribute_map[key]] = val
            else:
                parsed_kwargs[key] = val

        return parsed_kwargs

    @classmethod
    def deserialize(cls, data, to_cls):
        """De-serialize data from dictionary to an OCI model"""
        if cls.type_mappings is None:
            cls.type_mappings = cls.init_client().base_client.type_mappings

        if data is None:
            return None

        # This is a work around for enums not being present
        # in the type mappings.
        # See OraclePythonSdkCodegen removeEnumsFromModelGeneration().
        if to_cls in cls.type_mappings:
            to_cls = cls.type_mappings[to_cls]

            if isinstance(data, dict):
                data = cls._parse_kwargs(to_cls().attribute_map, **data)
        else:
            return cls.__deserialize_primitive(data, to_cls)

        if hasattr(to_cls, "get_subtype"):
            # Return the object as is if it is already deserialized.
            if isinstance(data, to_cls) or issubclass(data.__class__, to_cls):
                return data

            # Use the discriminator value to get the correct subtype.
            to_cls = to_cls.get_subtype(data)  # get_subtype returns a str
            to_cls = cls.type_mappings[to_cls]

            # kwargs needs to be parsed again as there are more attributes in the sub types.
            if isinstance(data, dict):
                data = cls._parse_kwargs(to_cls().attribute_map, **data)

        if to_cls in [int, float, six.u, bool]:
            return cls.__deserialize_primitive(data, to_cls)
        elif to_cls == object:
            return data
        elif to_cls == date:
            return cls.__deserialize_date(data)
        elif to_cls == datetime:
            return cls.__deserialize_datetime(data)
        else:
            return cls.__deserialize_model(data, to_cls)

    @classmethod
    def __deserialize_model(cls, data, to_cls):
        """De-serializes list or dict to model."""
        if isinstance(data, to_cls):
            return data

        instance = to_cls()

        for attr, attr_type in instance.swagger_types.items():
            property = instance.attribute_map[attr]
            if property in data:
                value = data[property]
                setattr(instance, attr, cls.deserialize(value, attr_type))

        return instance

    @staticmethod
    def __deserialize_primitive(data, cls):
        """De-serializes string to primitive type."""
        try:
            value = cls(data)
        except UnicodeEncodeError:
            value = six.u(data)
        except TypeError:
            value = data
        return value

    @staticmethod
    def __deserialize_date(string):
        """De-serializes string to date."""
        try:
            return parse(string).date()
        except ImportError:
            return string
        except ValueError:
            raise Exception("Failed to parse `{0}` into a date object".format(string))

    @staticmethod
    def __deserialize_datetime(string):
        """De-serializes string to datetime.

        The string should be in iso8601 datetime format.
        """
        if isinstance(string, datetime):
            return string
        try:
            # If this parser creates a date without raising an exception
            # then the time zone is utc and needs to be set.
            naivedatetime = datetime.strptime(string, "%Y-%m-%dT%H:%M:%S.%fZ")
            awaredatetime = naivedatetime.replace(tzinfo=tz.tzutc())
            return awaredatetime

        except ValueError:
            try:
                return parse(string)
            except ImportError:
                return string
            except ValueError:
                raise Exception(
                    "Failed to parse `{0}` into a datetime object".format(string)
                )
        except ImportError:
            return string


class OCIModelMixin(OCISerializableMixin):
    """Mixin class to operate OCI model.
    OCI resources are represented by models in the OCI Python SDK.

    Unifying OCI models for the same resource
    -----------------------------------------
    OCI SDK uses different models to represent the same resource for different operations.
    For example, CreateLogDetails is used when creating a log resource,
        while LogSummary is returned when listing the log resources.
    However, both CreateLogDetails and LogSummary have the same commonly used attribute like
        compartment_id, display_name, log_type, etc.
    In general, there is a class with a super set of all properties.
    For example, the Log class contains all properties of CreateLogDetails and LogSummary,
        as well as other classes representing an OCI log resource.
    A subclass can be implemented with this Mixin to unify the OCI models,
        so that all properties are available to the user.
    For example, if we define the Mixin model as ``class OCILog(OCIModelMixin, oci.logging.models.Log)``,
        users will be able to access properties like ``OCILog().display_name``
    Since this sub-class contains all the properties, it can be converted to any related OCI model in an operation.
    For example, we can create ``CreateLogDetails`` from ``OCILog`` by extracting a subset of the properties.
    When listing the resources, properties from ``LogSummary`` can be used to update
        the corresponding properties of ``OCILog``.
    Such convertion can be done be the generic methods provided by this Mixin.
    Although OCI SDK accepts dictionary (JSON) data instead of objects like CreateLogDetails when creating or
        updating the resource, the structure and keys of the dictionary is not easy for a user to remember.
    It is also unnecessary for the users to construct the entire dictionary if they only want to update a single value.

    This Mixin class should be the first parent as the class from OCI SDK does not call ``super().__init__()``
        in its ``__init__()`` constructor.
    Mixin properties may not be intialized correctly if ``super().__init__()`` is not called.


    Provide generic methods for CRUDL operations
    --------------------------------------------
    Since OCI SDK uses different models in CRUDL operations,
        this Mixin provides the following method to convert between them.
    An OCI model instance can be any OCI model of the resource containing some properties, e.g. LogSummary
    ``from_oci_model()`` static method can be used to initialize a new instance from an OCI model instance.
    ``update_from_oci_model()`` can be used to update the existing properties from an OCI model instance.
    ``to_oci_model()`` can be used to extract properties from the Mixin model to OCI model.

    """

    # Regex pattern matching the module name of an OCI model.
    OCI_MODEL_PATTERN = r"oci.[^.]+\.models[\..*]?"
    OCI_LIST_METHOD = None

    CONS_COMPARTMENT_ID = "compartment_id"

    @staticmethod
    def check_compartment_id(compartment_id: Optional[str]) -> str:
        """Checks if a compartment ID has value and
            return the value from NB_SESSION_COMPARTMENT_OCID environment variable if it is not specified.

        Parameters
        ----------
        compartment_id : str
            Compartment OCID or None

        Returns
        -------
        type
            str: Compartment OCID

        Raises
        ------
        ValueError
            compartment_id is not specified and NB_SESSION_COMPARTMENT_OCID environment variable is not set

        """
        compartment_id = (
            compartment_id
            if compartment_id
            else os.environ.get("NB_SESSION_COMPARTMENT_OCID")
        )
        if not compartment_id:
            raise ValueError("Please specify compartment OCID.")
        return compartment_id

    @staticmethod
    def camel_to_snake(s: str) -> str:
        """Converts camel style string to snake style.

        Parameters
        ----------
        s : str
            The string to be converted


        Returns
        -------
        str
            A string in snake format.

        """
        return re.sub(r"(?<!^)(?=[A-Z])", "_", s).lower()

    @classmethod
    def list_resource(cls, compartment_id: str = None, **kwargs) -> list:
        """Generic method to list OCI resources

        Parameters
        ----------
        compartment_id : str
            Compartment ID of the OCI resources. Defaults to None.
            If compartment_id is not specified,
            the value of NB_SESSION_COMPARTMENT_OCID in environment variable will be used.
        **kwargs :
            Additional keyword arguments to filter the resource.
            The kwargs are passed into OCI API.

        Returns
        -------
        list
            A list of OCI resources

        Raises
        ------
        NotImplementedError
            List method is not supported or implemented.

        """
        items = oci.pagination.list_call_get_all_results(
            cls._find_oci_method("list"),
            cls.check_compartment_id(compartment_id),
            **kwargs,
        ).data
        return [cls.from_oci_model(item) for item in items]

    @classmethod
    def _find_oci_parent(cls):
        """Finds the parent OCI model class

        Parameters
        ----------

        Returns
        -------
        class
            An OCI model class

        Raises
        ------
        AttributeError
            If the class if not inherit from an OCI model.

        """
        oci_model = None
        for parent in inspect.getmro(cls):
            if re.match(OCIModelMixin.OCI_MODEL_PATTERN, parent.__module__):
                oci_model = parent
                break
        if not oci_model:
            raise AttributeError(f"{cls.__name__} is not inherited from an OCI model.")
        return oci_model

    @classmethod
    def _find_oci_method(cls, method: str) -> Callable:
        """Finds the OCI method for operations like get/list

        Parameters
        ----------
        method : str
            The operation to be performed, e.g. get, list, delete


        Returns
        -------
        callable
            The method from OCI client to perform the operation.

        """
        client = cls.init_client()

        user_define_method = f"oci_{method}_method".upper()
        if hasattr(cls, user_define_method) and getattr(cls, user_define_method):
            return getattr(cls, user_define_method).__get__(client, client.__class__)

        oci_class = cls._find_oci_parent()
        if method == "list":
            method_name = f"{method}_{cls.camel_to_snake(oci_class.__name__)}s"
        else:
            method_name = f"{method}_{cls.camel_to_snake(oci_class.__name__)}"

        if hasattr(client, method_name):
            return getattr(client, method_name)

        raise NotImplementedError(
            f"{str(oci_class)} does not have {method_name} method. "
            f"Please define {user_define_method} in {str(cls)}"
        )

    @classmethod
    def from_oci_model(cls, oci_instance):
        """Initialize an instance from an instance of OCI model.

        Parameters
        ----------
        oci_instance :
            An instance of an OCI model.

        """
        oci_model = cls._find_oci_parent()
        kwargs = {}
        for attr in oci_model().swagger_types.keys():
            if hasattr(oci_instance, attr):
                kwargs[attr] = getattr(oci_instance, attr)
        return cls(**kwargs)

    @classmethod
    def from_dict(cls, data):
        """Initialize an instance from a dictionary.

        Parameters
        ----------
        data : dict
            A dictionary containing the properties to initialize the class.

        """
        instance = cls()
        oci_instance = cls.deserialize(data)
        instance.update_from_oci_model(oci_instance)
        return instance

    @classmethod
    def deserialize(cls, data, to_cls=None):
        """Deserialize data

        Parameters
        ----------
        data :

        to_cls :
             (Default value = None)

        Returns
        -------

        """
        if to_cls is None:
            oci_model = cls._find_oci_parent()
            to_cls = oci_model.__name__
        return super().deserialize(data, to_cls)

    @classmethod
    def _get(cls, ocid):
        get_method = cls._find_oci_method("get")
        return get_method(ocid).data

    @classmethod
    def from_ocid(cls, ocid: str):
        """Initializes an object from OCID

        Parameters
        ----------
        ocid : str
            The OCID of the object

        """
        oci_instance = cls._get(ocid)
        return cls.from_oci_model(oci_instance)

    def __init__(self, auth: dict = None, **kwargs) -> None:
        # Initialize an empty object
        super().__init__()
        if auth:
            self._client_auth.update(auth)
        # Load properties from environment
        self.load_properties_from_env()
        if kwargs:
            self.update_from_oci_model(self.deserialize(kwargs))
        # Retrieve the compartment ID if user specified the OCID of the resource but not compartment ID
        if (
            self.id
            and hasattr(self, self.CONS_COMPARTMENT_ID)
            and not getattr(self, self.CONS_COMPARTMENT_ID)
        ):
            from .oci_resource import OCIResource

            setattr(
                self, self.CONS_COMPARTMENT_ID, OCIResource.get_compartment_id(self.id)
            )

    @property
    def name(self):
        """Display name of the object"""
        if hasattr(self, "display_name"):
            return getattr(self, "display_name")
        raise NotImplementedError()

    def load_properties_from_env(self):
        """Loads properties from the environment"""
        env_var_mapping = {
            "project_id": "PROJECT_OCID",
            OCIModelMixin.CONS_COMPARTMENT_ID: "NB_SESSION_COMPARTMENT_OCID",
        }
        for attr, env_name in env_var_mapping.items():
            if hasattr(self, attr) and not getattr(self, attr):
                setattr(self, attr, os.environ.get(env_name))

    def to_oci_model(self, oci_model):
        """Converts the object into an instance of OCI data model.

        Parameters
        ----------
        oci_model : class or str
            The OCI model to be converted to. This can be a string of the model name.
        type_mapping : dict
            A dictionary mapping the models.
            Returns: An instance of the oci_model

        Returns
        -------

        """
        if isinstance(oci_model, str):
            oci_model_name = oci_model
        else:
            oci_model_name = oci_model.__name__
        data = json.dumps(self.to_dict()).encode("utf8")
        return self.client.base_client.deserialize_response_data(data, oci_model_name)

    @staticmethod
    def flatten(data: dict) -> dict:
        """Flattens a nested dictionary.

        Parameters
        ----------
        data : A nested dictionary


        Returns
        -------
        dict
            The flattened dictionary.

        """
        flatten_dict = {}
        for key, value in data.items():
            if isinstance(value, dict):
                flatten_dict.update(OCIModelMixin.flatten(value))
            else:
                flatten_dict[key] = value
        return flatten_dict

    def to_dict(self, flatten: bool = False) -> dict:
        """Converts the properties to a dictionary

        Parameters
        ----------
        flatten :
             (Default value = False)

        Returns
        -------

        """
        data = self.serialize()
        if flatten:
            data = OCIModelMixin.flatten(data)
        return data

    def update_from_oci_model(self, oci_model_instance):
        """Updates the properties from OCI model with the same properties.

        Parameters
        ----------
        oci_model_instance :
            An instance of OCI model, which should have the same properties of this class.

        """
        for attr in self.swagger_types.keys():
            if hasattr(oci_model_instance, attr) and getattr(oci_model_instance, attr):
                setattr(self, attr, getattr(oci_model_instance, attr))
        return self

    def sync(self):
        """Refreshes the properties of the object from OCI"""
        return self.update_from_oci_model(self.from_ocid(self.id))

    def delete(self):
        """Deletes the resource"""
        delete_method = self._find_oci_method("delete")
        delete_method(self.id)
        return self


class OCIWorkRequestMixin:
    """Mixin class containing methods related to OCI work request"""

    def wait_for_work_request(
        self,
        work_request_id: str,
        wait_for_state: Union[str, tuple],
        max_wait_seconds: int = None,
        wait_interval_seconds: int = None,
    ):
        """Wait for a work request to be completed.

        Parameters
        ----------
        work_request_id : str
            OCI work request ID
        wait_for_state : str or tuple
            The state to wait for. Must be a tuple for multiple states.
        max_wait_seconds : int
            Max wait seconds for the work request. Defaults to None (Default value from OCI SDK will be used).
        wait_interval_seconds : int
            Interval in seconds between each status check. Defaults to None (Default value from OCI SDK will be used).

        Returns
        -------
        Response
            OCI API Response

        """
        result = None
        if hasattr(self.client, "get_work_request") and callable(
            getattr(self.client, "get_work_request")
        ):
            try:
                wait_period_kwargs = {}
                if max_wait_seconds is not None:
                    wait_period_kwargs["max_wait_seconds"] = max_wait_seconds
                if wait_interval_seconds is not None:
                    wait_period_kwargs["max_interval_seconds"] = wait_interval_seconds

                logger.debug(
                    "Action completed. Waiting until the work request has entered state: %s",
                    wait_for_state,
                )
                result = oci.wait_until(
                    self.client,
                    self.client.get_work_request(work_request_id),
                    "status",
                    wait_for_state,
                    **wait_period_kwargs,
                )
            except oci.exceptions.MaximumWaitTimeExceeded:
                # If we fail, we should show an error, but we should still provide the information to the customer
                logger.error(
                    "Failed to wait until the work request %s entered the specified state. Maximum wait time exceeded",
                    work_request_id,
                )
            except Exception:
                logger.error(
                    "Encountered error while waiting for work request to enter the specified state."
                )
                raise
        else:
            logger.error(
                "%s does not support wait for the work request to enter the specified state",
                str(self.client.__class__),
            )
        return result
