import logging
import requests
from typing import Dict

from constants import constants as bdb_constants
from model_garden.model_config import ModelConfig
from model_garden.model_provider import ModelProvider
from utils.utils import Utils

logging.basicConfig(level=bdb_constants.LOGGING_LEVEL,
                    format="%(asctime)s - %(levelname)s - %(message)s")

logger = logging.getLogger(__name__)


class Model:
    def __init__(self, api_key: str, config: ModelConfig):
        self._api_key = api_key
        self.config = config

    def status(self):
        """
        Retrieves the current status of the model from BerryDB.

        This method queries BerryDB to get the operational status of the deployed
        or requested model.

        Returns:
        - `str`: A string representing the current status of the model (e.g., "AVAILABLE", "REQUESTED", "UNAVAILABLE" ...).

        Example:
        ```python
        from berrydb import BerryDB, ModelProvider

        # Assuming you have a Model instance, e.g., fetched from the ModelRepo
        repo = BerryDB.model_repo("BERRYDB_API_KEY")
        my_model = repo.get(provider=ModelProvider.CUSTOM_MODEL, model_name="my-deployed-model")

        # Check the model's status
        current_status = my_model.status()
        print(f"Model '{my_model.config.name}' status: {current_status}")

        # You might loop and check status until it's 'DEPLOYED'
        # while my_model.status() != "AVAILABLE":
        #     print("Model is not yet available, waiting...")
        #     time.sleep(10) # Wait for 10 seconds before checking again
        # print("Model is now available!")
        ```
        ---
        """
        print("Retrieving model status for model with name: ", self.config.name)
        url = bdb_constants.ML_BACKEND_BASE_URL + \
            bdb_constants.get_model_status_url.format(
                self.config.provider.value)
        params = {"apiKey": self._api_key, "modelId": self.config._id}
        if bdb_constants.debug_mode:
            print("url: ", url)
            print("params: ", params)
        response = requests.get(url, params=params)
        if bdb_constants.debug_mode:
            print("response.text: ", response.text)
            print("response.status_code: ", response.status_code)

        if response.status_code != 200:
            Utils.handleApiCallFailure(response.json(), response.status_code)
        data = response.json()
        status = data[0].get("status", None) if (
            isinstance(data, list) and len(data) > 0) else None
        if status is None:
            raise Exception(
                "Oops! couldn't fetch status, please try again later.")
        print("Model status:", status)
        return status

    def predict(self, input_data: any, annotation_config: str = None, model_args: dict = None) -> object:
        """
        Retrieves a prediction from the model hosted on BerryDB.

        This method sends input data to the deployed model and returns the
        prediction generated by the model. The structure and type of the
        input and output data depend on the specific model being used.

        Parameters:
        - **input_data** (`any`): The input data required by the model to generate a prediction. The format and type of this data are specific to the model's requirements.
        - **annotation_config** (`str`, optional): Annotation config name if the model's input and output needs to be transformed.
        - **model_args** (`dict`, optional): Additional arguments to pass to the model (eg: labels, etc).

        Returns:
        - `object`: The prediction result returned by the model. The structure and type of this object depend on the model's output format.

        Example:
        ```python
        from berrydb import BerryDB, ModelProvider

        # Assuming you have a Model instance, e.g., fetched from the ModelRepo
        repo = BerryDB.model_repo("BERRYDB_API_KEY")
        my_model = repo.get(provider=ModelProvider.HUGGING_FACE_MODEL, model_name="my-sentiment-analyzer")

        # Example input data (this will vary based on your model)
        text_to_predict = "This movie is absolutely fantastic!"

        # Get a prediction from the model
        prediction_result = my_model.predict(text_to_predict)
        print(f"Prediction for '{text_to_predict}': {prediction_result}")

        # Another example with structured input (if your model expects it)
        # structured_input = {"text": "The service was terrible.", "language": "en"}
        # prediction_result = my_model.predict(structured_input)
        # print(f"Prediction: {prediction_result}")
        ```
        ---
        """
        print("Fetching prediction for model with name: ", self.config.name)

        if model_args is not None and not isinstance(model_args, dict):
            raise ValueError("model_args must be None or a dictionary")
        if isinstance(model_args, dict) and "inputs" in model_args:
            raise ValueError("model_args cannot contain 'inputs' key")

        url = bdb_constants.ML_BACKEND_BASE_URL + \
            bdb_constants.model_predict_url.format(
                self.config.provider.value, self.config._id)
        payload = {"input": input_data}

        if model_args is not None:
            # copy so we don't mutate the original dict
            model_args = dict(model_args)

            # pull out label_studio_transform if it exists
            lst = model_args.pop("label_studio_transform", None)

            # merge remaining args into input
            if isinstance(payload["input"], dict):
                payload["input"] = {**payload["input"], **model_args}

            # add label_studio_transform as sibling if it was present
            if lst is not None:
                payload["label_studio_transform"] = lst

        params = {"apiKey": self._api_key}
        if annotation_config:
            params["annotationConfig"] = annotation_config
        if bdb_constants.debug_mode:
            print("url:", url)
            print("payload:", payload)
            print("params:", params)
        response = requests.post(url, json=payload, params=params)
        if bdb_constants.debug_mode:
            print("response.text:", response.text)
        if response.status_code != 200:
            Utils.handleApiCallFailure(response.json(), response.status_code)

        return response.json()


class BerryDBModel(Model):
    pass


class HuggingFaceModel(Model):

    def request(self):
        """
        Requests the setup or deployment of a Hugging Face model in BerryDB.

        This method is used for models obtained from the Hugging Face provider.
        If the model's current status is 'UNAVAILABLE', calling this method
        will initiate the process to make it available for predictions.
        If the model has already been requested or is in a state other than
        'UNAVAILABLE', this method may indicate that the action cannot be performed.

        Returns:
        - `dict`: A dictionary containing the response from BerryDB, typically
          indicating the success or status of the request.

        Raises:
        - `Exception`: If BerryDB returns an error, for example, if the model
          is not in a state where it can be requested (e.g., already 'AWAITING'
          or 'REQUESTED').

        Example:
        ```python
        from berrydb import BerryDB, ModelProvider

        berrydb_api_key = "BERRYDB_API_KEY"
        repo = BerryDB.model_repo(berrydb_api_key)

        try:
            # Get a Hugging Face model that might be 'UNAVAILABLE'
            hf_model = repo.get(
                provider=ModelProvider.HUGGING_FACE_MODEL,
                model_name="my-hf-model-to-request"
            )

            if hf_model and hf_model.config.status() == "UNAVAILABLE":
                print(f"Model '{hf_model.config.name}' is unavailable. Requesting deployment...")
                request_response = hf_model.request()
                print(f"Model request initiated: {request_response}")
                # You might want to check its status periodically after this
                # print(f"New status: {hf_model.status()}")
            elif hf_model:
                print(f"Model '{hf_model.config.name}' status is '{hf_model.config._hf_status}'. No request needed or possible.")
            else:
                print("Model not found.")

        except Exception as e:
            print(f"An error occurred: {e}")
        ```
        ---
        """
        # payload = {"id": self.config._id, "memory": memory}
        if (self.config._hf_status or "").lower() == "awaiting" or (self.config._hf_status or "").lower() == "requested":
            Utils.handleApiCallFailure(
                {"detail": "The model cannot be requested as it has already been requested before"}, 500)
        if not ((self.config._hf_status or "").lower() == "unavailable"):
            Utils.handleApiCallFailure(
                {"detail": "The model cannot be requested at this moment. Please try again later"}, 500)
        payload = {"id": self.config._id}
        url = bdb_constants.ML_BACKEND_BASE_URL + \
            bdb_constants.model_deploy_url.format(
                self.config.provider.value, self.config._id)
        params = {"apiKey": self._api_key}
        response = requests.post(url, json=payload, params=params)
        if response.status_code != 200:
            Utils.handleApiCallFailure(response.json(), response.status_code)
        return response.json()


class VertexAIModel(Model):

    def deploy(self, machine_type: str = None, min_replicas: int = None, max_replicas: int = None):
        """
        Deploys a VertexAI model hosted by BerryDB.

        This method initiates the deployment process for a VertexAI model
        that has been requested. It allows you to specify
        the machine type and scaling configuration for the deployment.

        This method is available for `Model` instances obtained for the
        `ModelProvider.VERTEX_AI_MODEL` providers.

        Parameters:
        - **machine_type** (`str`, optional): The type of machine to use for the model replicas (e.g., "n1-standard-4"). Defaults to a standard machine type if not specified.
        - **min_replicas** (`int`, optional): The minimum number of model replicas to maintain. Defaults to 1. Must be between 1 and 3.
        - **max_replicas** (`int`, optional): The maximum number of model replicas to scale up to. Defaults to 1. Must be between 1 and 3.

        Returns:
        - `dict`: A dictionary containing the response from BerryDB, typically
          confirming the deployment request and providing details about the initiated process.

        Raises:
        - `ValueError`: If `min_replicas` or `max_replicas` are outside the allowed range (1-3) or if `min_replicas` is greater than `max_replicas`.

        Example:
        ```python
        from berrydb import BerryDB, ModelProvider, ModelConfig

        # Initialize BerryDB (if not already done)
        # BerryDB.init("your-berrydb-host")

        berrydb_api_key = "BERRYDB_API_KEY"
        repo = BerryDB.model_repo(berrydb_api_key)

        # Assuming you have previously requested a VertexAI model and now want to deploy it.
        # You can retrieve the model instance using repo.get()
        try:
            vertexai_model_to_deploy = repo.get(
                provider=ModelProvider.VERTEX_AI_MODEL,
                model_name="my-image-classifier" # Name of your saved VertexAI model
            )

            if vertexai_model_to_deploy:
                print(f"Found VertexAI model: {vertexai_model_to_deploy.config.name}")

                # Deploy the model with specific scaling settings
                print("Requesting deployment...")
                deployment_response = vertexai_model_to_deploy.deploy(
                    machine_type="n1-standard-4",
                    min_replicas=1,
                    max_replicas=2
                )
                print(f"Deployment initiated: {deployment_response}")

                # You can check the status periodically using vertexai_model_to_deploy.status()
                # until it shows as 'AVAILABLE'.

            else:
                print("VertexAI model not found.")

        except ValueError as ve:
            print(f"Deployment configuration error: {ve}")
        except Exception as e:
            print(f"An error occurred during deployment: {e}")
        ```
        ---
        """
        min_replicas, max_replicas, machine_type = _validate_min_max_instances(
            min_replicas, max_replicas, machine_type)

        payload = {
            "name": self.config.name,
            "machineType": machine_type,
            "minReplicaCount": min_replicas,
            "maxReplicaCount": max_replicas,
        }
        url = bdb_constants.ML_BACKEND_BASE_URL + \
            bdb_constants.model_deploy_url.format(
                self.config.provider.value, self.config._id)
        params = {"apiKey": self._api_key}
        response = requests.post(url, json=payload, params=params)
        if response.status_code != 200:
            Utils.handleApiCallFailure(response.json(), response.status_code)
        json_res = response.json()
        return json_res["message"] if "message" in json_res else json_res


class CustomModel(Model):

    def deploy(self, machine_type: str = None, min_replicas: int = None, max_replicas: int = None):
        """
        Deploys a Custom model hosted by BerryDB.

        This method initiates the deployment process for a Custom model that has
        been saved to BerryDB (i.e., not self-hosted). It allows you to specify
        the machine type and scaling configuration for the deployment.

        This method is available for `Model` instances obtained for the
        `ModelProvider.CUSTOM_MODEL` (when not self-hosted).

        Parameters:
        - **machine_type** (`str`, optional): The type of machine to use for the model replicas (e.g., "n1-standard-4"). Defaults to a standard machine type if not specified.
        - **min_replicas** (`int`, optional): The minimum number of model replicas to maintain. Defaults to 1. Must be between 1 and 3.
        - **max_replicas** (`int`, optional): The maximum number of model replicas to scale up to. Defaults to 1. Must be between 1 and 3.

        Returns:
        - `dict`: A dictionary containing the response from BerryDB, typically
          confirming the deployment request and providing details about the initiated process.

        Raises:
        - `ValueError`: If `min_replicas` or `max_replicas` are outside the allowed range (1-3) or if `min_replicas` is greater than `max_replicas`.

        Example:
        ```python
        from berrydb import BerryDB, ModelProvider, ModelConfig

        # Initialize BerryDB (if not already done)
        # BerryDB.init("your-berrydb-host")

        berrydb_api_key = "BERRYDB_API_KEY"
        repo = BerryDB.model_repo(berrydb_api_key)

        # Assuming you have previously saved a non-self-hosted Custom model
        # and now want to deploy it.
        # You can retrieve the model instance using repo.get()
        try:
            custom_model_to_deploy = repo.get(
                provider=ModelProvider.CUSTOM_MODEL,
                model_name="my-image-classifier" # Name of your saved custom model
            )

            if custom_model_to_deploy:
                print(f"Found Custom model: {custom_model_to_deploy.config.name}")

                # Deploy the model with specific scaling settings
                print("Requesting deployment...")
                deployment_response = custom_model_to_deploy.deploy(
                    machine_type="n1-standard-4",
                    min_replicas=1,
                    max_replicas=2
                )
                print(f"Deployment initiated: {deployment_response}")

                # You can check the status periodically using custom_model_to_deploy.status()
                # until it shows as 'AVAILABLE'.

            else:
                print("Custom model not found.")

        except ValueError as ve:
            print(f"Deployment configuration error: {ve}")
        except Exception as e:
            print(f"An error occurred during deployment: {e}")
        ```
        ---
        """

        min_replicas, max_replicas, machine_type = _validate_min_max_instances(
            min_replicas, max_replicas, machine_type)

        payload = {
            "id": self.config._id,
            "machineType": machine_type,
            "minReplicaCount": min_replicas,
            "maxReplicaCount": max_replicas,
        }
        url = bdb_constants.ML_BACKEND_BASE_URL + \
            bdb_constants.model_deploy_url.format(
                self.config.provider.value, self.config._id)
        params = {"apiKey": self._api_key}
        response = requests.post(url, json=payload, params=params)
        if response.status_code != 200:
            Utils.handleApiCallFailure(response.json(), response.status_code)
        json_res = response.json()
        return json_res["message"] if "message" in json_res else json_res

    def shutdown(self):
        """
        Shuts down a deployed Custom model hosted by BerryDB.

        This method terminates the running instances (replicas) of a Custom model
        that was previously deployed on BerryDB. After calling this method, the
        model will no longer be available for predictions until it is deployed again.

        This method is available for `Model` instances obtained for the
        `ModelProvider.CUSTOM_MODEL` provider (specifically, non-self-hosted models
        that have been deployed).

        Returns:
        - `dict`: A dictionary containing the response from BerryDB, typically
          confirming the shutdown request.

        Example:
        ```python
        from berrydb import BerryDB, ModelProvider

        # Initialize BerryDB (if not already done)
        # BerryDB.init("your-berrydb-host")

        berrydb_api_key = "BERRYDB_API_KEY"
        repo = BerryDB.model_repo(berrydb_api_key)

        # Assuming you have a deployed non-self-hosted Custom model
        # and now want to shut it down.
        # You can retrieve the model instance using repo.get()
        try:
            custom_model_to_shutdown = repo.get(
                provider=ModelProvider.CUSTOM_MODEL,
                model_name="my-image-classifier" # Name of your deployed custom model
            )

            if custom_model_to_shutdown:
                print(f"Found Custom model: {custom_model_to_shutdown.config.name}")

                # Shut down the model
                print("Requesting shutdown...")
                shutdown_response = custom_model_to_shutdown.shutdown()
                print(f"Shutdown initiated: {shutdown_response}")

                # You can check the status periodically using custom_model_to_shutdown.status()
                # until it shows as 'UNAVAILABLE'.

            else:
                print("Custom model not found.")

        except Exception as e:
            print(f"An error occurred during shutdown: {e}")
        ```
        """
        print("Shutting down model with name: ", self.config.name)
        if self.config.self_hosted:
            raise ValueError("Shutdown not supported for self-hosted models")
        url = bdb_constants.ML_BACKEND_BASE_URL + \
            bdb_constants.model_shutdown_url.format(
                self.config.provider.value, self.config._id)
        params = {"apiKey": self._api_key}
        response = requests.post(url, params=params)
        if response.status_code != 200:
            Utils.handleApiCallFailure(response.json(), response.status_code)

        return response.json()


def _validate_min_max_instances(min_replicas: str | None, max_replicas: str | None, machine_type: str | None):
    if min_replicas is None:
        min_replicas = bdb_constants.DEFAULT_MIN_REPLICAS
    if max_replicas is None:
        max_replicas = bdb_constants.DEFAULT_MAX_REPLICAS
    if bdb_constants.MIN_REPLICAS <= min_replicas >= bdb_constants.MAX_REPLICAS:
        raise ValueError("min_replicas must be in the range of 1-3")
    if bdb_constants.MIN_REPLICAS <= max_replicas >= bdb_constants.MAX_REPLICAS:
        raise ValueError("max_replicas must be in the range of 1-3")
    if min_replicas > max_replicas:
        raise ValueError(
            "min_replicas must be less than or equal to max_replicas")
    if machine_type is None:
        machine_type = bdb_constants.DEFAULT_MACHINE_TYPE
    return min_replicas, max_replicas, machine_type
