import logging
import os
from typing import Text, Dict
import typing
import convo.utils.common as convo_utils
from convo.constants import (
    ENV_GPU_CONFIG,
    ENV_CPU_INTER_OP_CONFIG,
    ENV_CPU_INTRA_OP_CONFIG,
)

if typing.TYPE_CHECKING:
    from tensorflow import config as tf_config

logger = logging.getLogger(__name__)


def _setup_gpu_environment() -> None:
    """Set configuration for TensorFlow GPU environment based on the environment variable set."""

    gpu_memory_config = os.getenv(ENV_GPU_CONFIG)

    if not gpu_memory_config:
        return

    # Import from tensorflow only if necessary (environment variable was set)
    from tensorflow import config as tf_config

    parsed_gpu_config = _parse_gpu_config(gpu_memory_config)
    physical_gpus = tf_config.list_physical_devices("GPU")

    # Logic taken from https://www.tensorflow.org/guide/gpu
    if physical_gpus:
        for gpu_id, gpu_id_memory in parsed_gpu_config.items():
            _allocate_gpu_memory(physical_gpus[gpu_id], gpu_id_memory)

    else:
        convo_utils.raise_warning(
            f"You have an environment variable '{ENV_GPU_CONFIG}' set but no GPUs were detected to configure."
        )


def _allocate_gpu_memory(
    gpu_instance: "tf_config.PhysicalDevice", logical_memory: int
) -> None:
    """Create a new logical device for the requested amount of memory.

    Args:
        gpu_instance: PhysicalDevice instance of a GPU device.
        logical_memory: Absolute amount of memory to be allocated to the new logical device.
    """

    from tensorflow import config as tf_config

    try:
        tf_config.experimental.set_virtual_device_configuration(
            gpu_instance,
            [
                tf_config.experimental.VirtualDeviceConfiguration(
                    memory_limit=logical_memory
                )
            ],
        )

    except RuntimeError:
        # Helper explanation of where the error comes from
        raise RuntimeError(
            "Error while setting up tensorflow environment. "
            "Virtual devices must be set before GPUs have been initialized."
        )


def _parse_gpu_config(gpu_memory_config: Text) -> Dict[int, int]:
    """Parse GPU configuration variable from a string to a dict.

    Args:
        gpu_memory_config: String containing the configuration for GPU usage.

    Returns:
        Parsed configuration as a dictionary with GPU IDs as keys and requested memory as the value.
    """

    # gpu_config is of format "gpu_id_1:gpu_id_1_memory, gpu_id_2: gpu_id_2_memory"
    # Parse it and store in a dictionary
    parsed_gpu_config = {}

    try:
        for instance in gpu_memory_config.split(","):
            instance_gpu_id, instance_gpu_mem = instance.split(":")
            instance_gpu_id = int(instance_gpu_id)
            instance_gpu_mem = int(instance_gpu_mem)

            parsed_gpu_config[instance_gpu_id] = instance_gpu_mem
    except ValueError:
        # Helper explanation of where the error comes from
        raise ValueError(
            f"Error parsing GPU configuration. Please cross-check the format of '{ENV_GPU_CONFIG}' "
            f"at https://convo.com/docs/convo/api/tensorflow_usage.html#restricting-absolute-gpu-memory-available ."
        )

    return parsed_gpu_config


def _setup_cpu_environment() -> None:
    """Set configuration for the CPU environment based on the environment variable set."""

    inter_op_parallel_threads = os.getenv(ENV_CPU_INTER_OP_CONFIG)
    intra_op_parallel_threads = os.getenv(ENV_CPU_INTRA_OP_CONFIG)

    if not inter_op_parallel_threads and not intra_op_parallel_threads:
        return

    from tensorflow import config as tf_config

    if inter_op_parallel_threads:
        try:
            inter_op_parallel_threads = int(inter_op_parallel_threads.strip())
        except ValueError:
            raise ValueError(
                f"Error parsing the environment variable '{ENV_CPU_INTER_OP_CONFIG}'. Please "
                f"cross-check the value."
            )

        tf_config.threading.set_inter_op_parallelism_threads(inter_op_parallel_threads)

    if intra_op_parallel_threads:
        try:
            intra_op_parallel_threads = int(intra_op_parallel_threads.strip())
        except ValueError:
            raise ValueError(
                f"Error parsing the environment variable '{ENV_CPU_INTRA_OP_CONFIG}'. Please "
                f"cross-check the value."
            )

        tf_config.threading.set_intra_op_parallelism_threads(intra_op_parallel_threads)


def setup_tf_environment() -> None:
    """Setup CPU and GPU related environment settings for TensorFlow."""

    _setup_cpu_environment()
    _setup_gpu_environment()
