# Copyright 2023 Agnostiq Inc.


import re
from dataclasses import asdict
from datetime import timedelta
from typing import Optional, Union

import arrow
from pydantic import model_validator
from pydantic.dataclasses import dataclass

executor_plugin_name = "cloud"

unit_multipliers = {"GB": 1024, "GiB": 954, "MB": 1}


@dataclass
class CloudExecutor:
    """
    CloudExecutor represents a configuration for executing a Covalent workflow on the Covalent Cloud.
    This class allows users to configure the resources (such as the number of CPUs, memory, GPUs, and GPU type) and the software environment for a Covalent workflow that will be executed on the Covalent Cloud. The time limit for the workflow execution can also be set.

    Attributes:

        num_cpus (int, optional): Number of CPUs to be used for the workflow. Defaults to 1.
        memory (int, optional): Amount of memory (in MB) to be used for the workflow. Defaults to 1024.
        num_gpus (int, optional): Number of GPUs to be used for the workflow. Defaults to 0.
        gpu_type (str, optional): Type of GPU to be used for the workflow. Defaults to an empty string.
        env (str, optional): Name of the software environment to be used for the workflow. Defaults to "default".
        time_limit (Union[int, timedelta], optional): Time limit for the workflow execution, in seconds or as a timedelta. Defaults to 1800s (30 mins). Alternatively can take human readable string in format 'in <number> <unit(s)>'

    Examples:

        # create a CloudExecutor with default resource configuration
        # executor = CloudExecutor()

        # create a custom CloudExecutor with specified resources and environment
        # executor = CloudExecutor(
        #     num_cpus=4,
        #     memory=2048,
        #     num_gpus=1,
        #     gpu_type="NVIDIA-Tesla-V100",
        #     env="my_custom_env",
        #     time_limit="in 30 minutes"  # 30 minutes
        # )

        import covalent as ct
        from covalent_cloud import CloudExecutor

        cloud_executor1 = CloudExecutor(num_cpus=1, memory=1024)
        cloud_executor2 = CloudExecutor(num_cpus=2, memory=2048)
        cloud_executor3 = CloudExecutor(num_cpus=1, memory=512)

        # Define manageable tasks as electrons with different cloud executors
        @ct.electron(executor=cloud_executor1)
        def add(x, y):
            return x + y

        @ct.electron(executor=cloud_executor2)
        def multiply(x, y):
            return x * y

        @ct.electron(executor=cloud_executor3)
        def divide(x, y):
            return x / y

        # Define the workflow as a lattice
        @ct.lattice
        def workflow(x, y):
            r1 = add(x, y)
            r2 = [multiply(r1, y) for _ in range(4)]
            r3 = [divide(x, value) for value in r2]
            return r3

        # Import the Covalent Cloud module
        import covalent_cloud as ctc

        # Dispatch the workflow to the Covalent Cloud
        dispatch_id = ctc.dispatch(workflow)(1, 2)
        result = ctc.get_result(dispatch_id, wait=True)
        print(result)
    """

    num_cpus: int = 1
    memory: Union[int, str] = 1024
    num_gpus: int = 0
    gpu_type: str = ""
    env: str = "default"
    time_limit: Union[int, timedelta, str] = 60 * 30
    volume_id: Optional[int] = None

    @property
    def short_name(self) -> str:
        """
        Property which returns the short name
        of the executor used by Covalent for identification.

        Args:
            None

        Returns:
            The short name of the executor

        """

        return executor_plugin_name

    # Validators:
    @model_validator(mode="after")
    def gt_than_zero_validator(self) -> None:
        """
        Validator which ensures that the values of the executor
        are greater than 0.

        Args:
            None

        Returns:
            None

        """

        # memory
        if isinstance(self.memory, str):

            # grab the number and the unit
            match = re.match(r"(\d+)\s*([A-Za-z]+)", self.memory)
            if match:
                memory_value, unit = match.groups()
            else:
                raise ValueError("Invalid memory string format")

            # convert to MB (int)
            try:
                memory_value = int(memory_value)
                unit = unit.strip()
                self.memory = int(memory_value * unit_multipliers[unit])
            except ValueError as e:
                raise ValueError(f"Invalid memory value: {self.memory}.") from e

        # time_limit
        if isinstance(self.time_limit, timedelta):
            self.time_limit = self.time_limit.total_seconds()

        elif isinstance(self.time_limit, str):
            try:
                now = arrow.utcnow()

                later = now.dehumanize(self.time_limit)

                # arrow dehumanize method calculates time relative to another time
                difference = later - now

                self.time_limit = difference.total_seconds()
            except Exception as e:
                raise ValueError(
                    f"Invalid time limit: {self.time_limit}. Please provide a valid time limit with proper grammar."
                ) from e

        self.gt_than_zero(self.num_cpus)
        self.gt_than_zero(self.memory)
        self.gt_than_zero(self.time_limit)

    def gt_than_zero(self, v: int) -> int:
        """
        Validator which ensures that the value is greater than 0.

        Args:
            v: The value to validate

        Returns:
            The validated value

        """

        if v <= 0:
            raise ValueError(f"{v} must be greater than 0")
        return v

    def to_dict(self) -> dict:
        """
        Return a JSON-serializable dictionary representation of this object.

        Args:
            None

        Returns:
            The JSON-serializable dictionary representation of this object.
        """

        return {
            "type": str(self.__class__),
            "short_name": self.short_name,
            "attributes": asdict(self),
        }
