import logging
import re
import time
from dataclasses import dataclass
from packaging.version import Version
from typing import Any, Dict, List, Optional, Tuple, Union

import ray
import yaml
from ray.experimental.state.api import get_node, get_placement_group
from ray.util.placement_group import placement_group, placement_group_table
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

from deltacat import logs

logger = logs.configure_deltacat_logger(logging.getLogger(__name__))

# Limitation of current node group or placement group manager
# Must run on driver or head node bc state.api needs to query dashboard api server at 127.0.0.1.
# Issue: https://github.com/ray-project/ray/issues/29959


def _get_available_resources_per_node():
    # This API changed after this commit
    # https://github.com/ray-project/ray/pull/43252
    # TODO: Use this method from a durable State API once it's available
    if Version(ray.__version__) >= Version("2.10.0"):
        return ray._private.state.available_resources_per_node()
    else:
        return ray._private.state.state._available_resources_per_node()


@dataclass
class PlacementGroupConfig:
    def __init__(self, opts, resource, node_ips):
        self.opts = opts
        self.resource = resource
        self.node_ips = node_ips


class NodeGroupManager:
    def __init__(self, path: str, gname: str):
        """Node Group Manager
        Args:
                path: cluster yaml file
                gname: node group prefix, e.g., 'partition'
        """
        # cluster init status:
        self.NODE_GROUP_PREFIX = gname
        self.cluster_config = self._read_yaml(path)
        self.init_groups = self._cluster_node_groups(self.cluster_config)
        self.init_group_res = self._parse_node_resources()

    def _cluster_node_groups(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """Get Worker Groups
        Args:
                config: cluster yaml data
        Returns:
                worker groups: a dict of worker node group

        """
        avail_node_types = list(config["available_node_types"].items())
        # exclude head node type
        head_node_types = [
            nt
            for nt in avail_node_types
            if "resources" in nt[1]
            and "CPU" in nt[1]["resources"]
            and nt[1]["resources"]["CPU"] == 0
        ][0]
        worker_node_types = [x for x in avail_node_types if x != head_node_types]
        # assuming homogenous cluster
        # in future, update with fleet resource
        if len(worker_node_types) > 0:
            self.INSTANCE_TYPE = worker_node_types[0][1]["node_config"]["InstanceType"]
        return worker_node_types

    def _read_yaml(self, path: str) -> Dict[str, Any]:
        with open(path, "rt") as f:
            return yaml.safe_load(f)

    def _update_groups(self) -> List[Tuple[str, float]]:
        """
                Node groups can come and go during runtime, whenever a node group is needed, we need to check the current available groups
        Returns:
                current_groups: dict of custom resource groups
        """
        # Add 1.1 second latency to avoid inconsistency issue between raylet and head
        time.sleep(1.1)
        all_available_res = ray.available_resources()
        current_groups = [
            (k, all_available_res[k])
            for k in all_available_res.keys()
            if self.NODE_GROUP_PREFIX in k
        ]
        return current_groups

    def _parse_node_resources(self) -> Dict[str, Dict[str, float]]:
        """
        Parse resources per node to get detailed resource tighted to each node group
        Returns:
                group_res: a dict of resources, e.g., {'CPU':0,'memory':0,'object_store_memory':0}
        """
        all_available_resources = _get_available_resources_per_node()
        group_keys = [x[0] for x in self.init_groups]
        group_res = {}
        for k in group_keys:
            group_res[k] = {
                "CPU": 0,
                "memory": 0,
                "object_store_memory": 0,
                "node_id": [],
            }
        for v in all_available_resources.values():
            keys = v.keys()
            r = re.compile(self.NODE_GROUP_PREFIX)
            partition = list(filter(r.match, list(keys)))
            r = re.compile("node:")
            node_id = list(filter(r.match, list(keys)))
            if len(partition) > 0:
                partition = partition[0]
            if len(node_id) > 0:
                node_id = node_id[0]
            if self.NODE_GROUP_PREFIX in partition:
                group_res[partition]["CPU"] += v["CPU"]
                group_res[partition]["memory"] += v["memory"]
                group_res[partition]["object_store_memory"] += v["object_store_memory"]
                group_res[partition]["node_id"].append(node_id)
        return group_res

    def _update_group_res(self, gname: str) -> Dict[str, Union[str, float]]:
        """
        Get the realtime resource of a node group
        Args:
                gname: name of node group
        Returns:
                group_res: dict of updated resource(cpu, memory, object store memory) for a given group
        """
        all_available_resources = _get_available_resources_per_node()
        group_res = {"CPU": 0, "memory": 0, "object_store_memory": 0, "node_id": []}
        for v in all_available_resources.values():
            keys = v.keys()
            r = re.compile("node:")
            node_id = list(filter(r.match, list(keys)))
            if len(node_id) > 0:
                node_id = node_id[0]
            if gname in v.keys():
                group_res["CPU"] += v["CPU"]
                group_res["memory"] += v["memory"]
                group_res["object_store_memory"] += v["object_store_memory"]
                group_res["node_id"].append(node_id)
        return group_res

    def get_one_group(self) -> Optional[Dict[str, Union[str, float]]]:
        """
        Pop up one node group
        Returns:
                group_res: dict of node group resource, {"group":"partition_1","CPU":2,...}
        """
        current_groups = self._update_groups()
        if len(current_groups) > 0:
            gname = current_groups[-1][0]
            group_res = self._update_group_res(gname)
            group_res["group"] = gname
            try:
                group_res["group_res"] = ray.available_resources()[gname]
            except Exception as e:
                logger.info(f"Error: {e}. There is no available resources for {gname}")
                return None
            return group_res
        else:
            return None

    def get_group_by_name(self, gname: str) -> Optional[Dict[str, Union[str, float]]]:
        """
        Get the specific node group given its pre-filled name
        Args:
                gname: name of the node group
        Returns:
                group_res: dict of node group resource

        """
        group_res = self._update_group_res(gname)
        group_res["group"] = gname
        try:
            group_res["group_res"] = ray.available_resources()[gname]
        except Exception as e:
            logger.info(f"Error: {e}. There is no available resources for {gname}")
            return None
        return group_res


class PlacementGroupManager:
    """Placement Group Manager
    Create a list of placement group with the desired number of cpus
    e.g., create a pg with 32 cpus, then this class will look for a node that has 32 cpus, and collect all
    resources, including cpu, memory, and object store;
    How to use:
            ```
                    from deltacat.utils.placement import PlacementGroupManager as pgm
                    pgm = pgm(10, 32)
                    pg_configs = pgm.pgs
                    opts = pg_configs[0][0]
                    fun.options(**opts).remote()
            ```
    Args:
            num_pgs: number of placement groups to be created
            instance_cpus: number of cpus per instance
    """

    def __init__(
        self,
        num_pgs: int,
        total_cpus_per_pg: int,
        cpu_per_bundle: int,
        strategy="SPREAD",
        capture_child_tasks=True,
        memory_per_bundle=None,
    ):
        head_res_key = self.get_current_node_resource_key()
        # run the task on head and consume a fractional cpu, so that pg can be created on non-head node
        # if cpu_per_bundle is less than the cpus per node, the pg can still be created on head
        # curent assumption is that the cpu_per_bundle = cpus per node
        # TODO: figure out how to create pg on non-head explicitly
        self._pg_configs = ray.get(
            [
                _config.options(resources={head_res_key: 0.01}).remote(
                    total_cpus_per_pg,
                    cpu_per_bundle,
                    strategy,
                    capture_child_tasks,
                    memory_per_bundle=memory_per_bundle,
                )
                for i in range(num_pgs)
            ]
        )
        # TODO: handle the cases where cpu_per_bundle is larger than max cpus per node, support it on ec2/flex/manta

    @property
    def pgs(self):
        return self._pg_configs

    def get_current_node_resource_key(self) -> str:
        # on ec2: address="172.31.34.51:6379"
        # on AWS Glue for Ray: address = "2600:1f10:4674:6815:aadb:2dc8:de61:bc8e:6379"
        (
            current_node_name,
            _,
        ) = ray.experimental.internal_kv.global_gcs_client.address.rsplit(
            ":", 1
        )  # using rsplit split on the last occurence of delimiter ":"
        for node in ray.nodes():
            if node["NodeName"] == current_node_name:
                # Found the node.
                for key in node["Resources"].keys():
                    if key.startswith("node:"):
                        return key


@ray.remote(num_cpus=0.01)
def _config(
    total_cpus_per_pg: int,
    cpu_per_node: int,
    strategy="SPREAD",
    capture_child_tasks=True,
    time_out: Optional[float] = None,
    memory_per_bundle: Optional[float] = None,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    pg_config = None
    opts = {}
    cluster_resources = {}
    num_bundles = (int)(total_cpus_per_pg / cpu_per_node)
    bundles = [{"CPU": cpu_per_node} for i in range(num_bundles)]

    if memory_per_bundle:
        for bundle in bundles:
            bundle["memory"] = memory_per_bundle

    pg = placement_group(bundles, strategy=strategy)
    ray.get(pg.ready(), timeout=time_out)
    if not pg:
        return None
    opts = {
        "scheduling_strategy": PlacementGroupSchedulingStrategy(
            placement_group=pg, placement_group_capture_child_tasks=capture_child_tasks
        )
    }
    pg_id = placement_group_table(pg)["placement_group_id"]
    pg_details = get_placement_group(pg_id)
    bundles = pg_details["bundles"]
    node_ids = []
    for bd in bundles:
        node_ids.append(bd["node_id"])
    # query available resources given list of node id
    all_nodes_available_res = _get_available_resources_per_node()
    pg_res = {"CPU": 0, "memory": 0, "object_store_memory": 0}
    node_ips = []
    for node_id in node_ids:
        if node_id in all_nodes_available_res:
            v = all_nodes_available_res[node_id]
            node_detail = get_node(node_id)
            pg_res["CPU"] += node_detail["resources_total"]["CPU"]
            pg_res["memory"] += v["memory"]
            pg_res["object_store_memory"] += v["object_store_memory"]
            node_ips.append(node_detail["node_ip"])
    cluster_resources["CPU"] = int(pg_res["CPU"])
    cluster_resources["memory"] = float(pg_res["memory"])
    cluster_resources["object_store_memory"] = float(pg_res["object_store_memory"])
    pg_config = PlacementGroupConfig(
        opts=opts, resource=cluster_resources, node_ips=node_ips
    )
    logger.info(f"pg has resources:{cluster_resources}")
    logger.debug(f"pg has node ips:{node_ips}")

    return pg_config
