import asyncio
import logging
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Awaitable, Callable, Iterable, Optional, Union

from kubernetes import client, config as k8s_config
from kubernetes.client import ApiException
from kubernetes.client.models import (
    V1Container,
    V1DaemonSet,
    V1Deployment,
    V1Job,
    V1Pod,
    V1PodList,
    V1StatefulSet,
    V2HorizontalPodAutoscaler,
)

from server.recommendation.vertical_rightsizing.models.objects import HPAData, K8sObjectData, KindLiteral, PodData
from server.recommendation.vertical_rightsizing.models.result import ResourceAllocations
from server.recommendation.vertical_rightsizing.models.object_like_dict import ObjectLikeDict
from server.recommendation.vertical_rightsizing.models.config import Config

logger = logging.getLogger("krr")

AnyKubernetesAPIObject = Union[V1Deployment, V1DaemonSet, V1StatefulSet, V1Pod, V1Job]
HPAKey = tuple[str, str, str]


class ClusterService:
    def __init__(self, cluster: str, executor: ThreadPoolExecutor, config: Config):
        self.cluster = cluster
        self._config = config
        # This executor will be running requests to Kubernetes API
        self.executor = executor
        k8s_config.load_config()
        self.api_client = client.ApiClient()
        self.apps = client.AppsV1Api(api_client=self.api_client)
        self.custom_objects = client.CustomObjectsApi(api_client=self.api_client)
        self.batch = client.BatchV1Api(api_client=self.api_client)
        self.core = client.CoreV1Api(api_client=self.api_client)
        self.autoscaling_v1 = client.AutoscalingV1Api(api_client=self.api_client)
        self.autoscaling_v2 = client.AutoscalingV2Api(api_client=self.api_client)

        self.__kind_available: defaultdict[KindLiteral, bool] = defaultdict(lambda: True)

        self.__jobs_for_cronjobs: dict[str, list[V1Job]] = {}
        self.__jobs_loading_locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock)

    async def list_scannable_objects(self) -> list[K8sObjectData]:
        """List all scannable objects.

        Returns:
            A list of scannable objects.
        """

        logger.info(f"Listing scannable objects in {self.cluster}")
        logger.debug(f"Namespaces: {self._config.namespaces}")
        logger.debug(f"Resource Kinds: {self._config.resource_kinds}")

        self.__hpa_list = await self._try_list_hpa()
        workload_object_lists = await asyncio.gather(
            self._list_deployments(),
            self._list_rollouts(),
            self._list_deploymentconfig(),
            self._list_all_statefulsets(),
            self._list_all_daemon_set(),
            self._list_all_jobs(),
            self._list_all_cronjobs(),
        )

        filtered1 = [
            object
            for workload_objects in workload_object_lists
            for object in workload_objects
            # NOTE: By default we will filter out kube-system namespace
            if not (
                (self._config.namespaces is None or len(self._config.namespaces) == 0)
                and object.namespace == "kube-system"
            )
        ]

        if self._config.resource_names is not None and len(self._config.resource_names) > 0:
            filtered1 = [object for object in filtered1 if object.name in self._config.resource_names]

        return filtered1

    async def _list_jobs_for_cronjobs(self, namespace: str) -> list[V1Job]:
        if namespace not in self.__jobs_for_cronjobs:
            loop = asyncio.get_running_loop()

            async with self.__jobs_loading_locks[namespace]:
                logging.debug(f"Loading jobs for cronjobs in {namespace}")
                ret = await loop.run_in_executor(
                    self.executor,
                    lambda: self.batch.list_namespaced_job(namespace=namespace),
                )
                self.__jobs_for_cronjobs[namespace] = ret.items

        return self.__jobs_for_cronjobs[namespace]

    async def list_pods(self, object: K8sObjectData) -> list[PodData]:
        loop = asyncio.get_running_loop()

        if object.kind == "CronJob":
            namespace_jobs = await self._list_jobs_for_cronjobs(object.namespace)
            ownered_jobs_uids = [
                job.metadata.uid
                for job in namespace_jobs
                if any(
                    owner.kind == "CronJob" and owner.uid == object._api_resource.metadata.uid
                    for owner in job.metadata.owner_references or []
                )
            ]
            selector = f"batch.kubernetes.io/controller-uid in ({','.join(ownered_jobs_uids)})"

        else:
            if object.selector is None:
                return []

            selector = self._build_selector_query(object.selector)
            if selector is None:
                return []

        ret: V1PodList = await loop.run_in_executor(
            self.executor,
            lambda: self.core.list_namespaced_pod(
                namespace=object._api_resource.metadata.namespace, label_selector=selector
            ),
        )

        return [PodData(name=pod.metadata.name, deleted=False) for pod in ret.items]

    @staticmethod
    def _get_match_expression_filter(expression) -> str:
        if expression.operator.lower() == "exists":
            return expression.key
        elif expression.operator.lower() == "doesnotexist":
            return f"!{expression.key}"

        values = ",".join(expression.values)
        return f"{expression.key} {expression.operator} ({values})"

    @staticmethod
    def _build_selector_query(selector: Any) -> Union[str, None]:
        label_filters = []

        if selector.match_labels is not None:
            label_filters += [f"{label[0]}={label[1]}" for label in selector.match_labels.items()]

        if hasattr(selector, "matchLabels") and selector.matchLabels is not None:
            label_filters += [f"{label[0]}={label[1]}" for label in selector.matchLabels.items()]

        if selector.match_expressions is not None:
            label_filters += [
                ClusterService._get_match_expression_filter(expression) for expression in selector.match_expressions
            ]

        if label_filters == []:
            # NOTE: This might mean that we have DeploymentConfig,
            # which uses ReplicationController and it has a dict like matchLabels
            if len(selector) != 0:
                label_filters += [f"{label[0]}={label[1]}" for label in selector.items()]
            else:
                return None

        return ",".join(label_filters)

    def __build_scannable_object(
        self, item: AnyKubernetesAPIObject, container: V1Container, kind: Optional[str] = None
    ) -> K8sObjectData:
        name = item.metadata.name
        namespace = item.metadata.namespace
        kind = kind or item.__class__.__name__[2:]

        obj = K8sObjectData(
            cluster=self.cluster,
            namespace=namespace,
            name=name,
            kind=kind,
            container=container.name,
            allocations=ResourceAllocations.from_container(container),
            hpa=self.__hpa_list.get((namespace, kind, name)),
        )
        obj._api_resource = item
        return obj

    def _should_list_resource_kind(self, resource_kind: str) -> bool:
        if self._config.resource_kinds is None or self._config.resource_kinds == "*":
            return True
        return resource_kind in self._config.resource_kinds

    async def _list_namespaced_or_global_objects(
        self, kind: KindLiteral, all_namespaces_request: Callable, namespaced_request: Callable
    ) -> list[Any]:
        logger.debug(f"Listing {kind}s in {self.cluster}")
        loop = asyncio.get_running_loop()

        if self._config.namespaces == "*":
            requests = [
                loop.run_in_executor(
                    self.executor,
                    lambda: all_namespaces_request(
                        watch=False,
                        label_selector=self._config.selector,
                    ),
                )
            ]
        else:
            requests = [
                loop.run_in_executor(
                    self.executor,
                    lambda ns=namespace: namespaced_request(
                        namespace=ns,
                        watch=False,
                        label_selector=self._config.selector,
                    ),
                )
                for namespace in self._config.namespaces
            ]

        result = [item for request_result in await asyncio.gather(*requests) for item in request_result.items]

        logger.debug(f"Found {len(result)} {kind} in {self.cluster}")
        return result

    async def _list_scannable_objects(
        self,
        kind: KindLiteral,
        all_namespaces_request: Callable,
        namespaced_request: Callable,
        extract_containers: Callable[[Any], Union[Iterable[V1Container], Awaitable[Iterable[V1Container]]]],
        filter_workflows: Optional[Callable[[Any], bool]] = None,
    ) -> list[K8sObjectData]:
        if not self._should_list_resource_kind(kind):
            logger.debug(f"Skipping {kind}s in {self.cluster}")
            return []

        if not self.__kind_available[kind]:
            return []

        result = []
        try:
            for item in await self._list_namespaced_or_global_objects(kind, all_namespaces_request, namespaced_request):
                if filter_workflows is not None and not filter_workflows(item):
                    continue

                containers = extract_containers(item)
                if asyncio.iscoroutine(containers):
                    containers = await containers

                result.extend(self.__build_scannable_object(item, container, kind) for container in containers)
        except ApiException as e:
            if kind in ("Rollout", "DeploymentConfig") and e.status in [400, 401, 403, 404]:
                if self.__kind_available[kind]:
                    logger.debug(f"{kind} API not available in {self.cluster}")
                self.__kind_available[kind] = False
            else:
                logger.exception(f"Error {e.status} listing {kind} in cluster {self.cluster}: {e.reason}")
                logger.error("Will skip this object type and continue.")

        return result

    def _list_deployments(self) -> list[K8sObjectData]:
        return self._list_scannable_objects(
            kind="Deployment",
            all_namespaces_request=self.apps.list_deployment_for_all_namespaces,
            namespaced_request=self.apps.list_namespaced_deployment,
            extract_containers=lambda item: item.spec.template.spec.containers,
        )

    def _list_rollouts(self) -> list[K8sObjectData]:
        async def _extract_containers(item: Any) -> list[V1Container]:
            if item.spec.template is not None:
                return item.spec.template.spec.containers

            loop = asyncio.get_running_loop()

            logging.debug(
                f"Rollout has workloadRef, fetching template for {item.metadata.name} in {item.metadata.namespace}"
            )

            # Template can be None and object might have workloadRef
            workloadRef = item.spec.workloadRef
            if workloadRef is not None:
                ret = await loop.run_in_executor(
                    self.executor,
                    lambda: self.apps.read_namespaced_deployment(
                        namespace=item.metadata.namespace, name=workloadRef.name
                    ),
                )
                return ret.spec.template.spec.containers

            return []

        # NOTE: Using custom objects API returns dicts, but all other APIs return objects
        # We need to handle this difference using a small wrapper
        return self._list_scannable_objects(
            kind="Rollout",
            all_namespaces_request=lambda **kwargs: ObjectLikeDict(
                self.custom_objects.list_cluster_custom_object(
                    group="argoproj.io",
                    version="v1alpha1",
                    plural="rollouts",
                    **kwargs,
                )
            ),
            namespaced_request=lambda **kwargs: ObjectLikeDict(
                self.custom_objects.list_namespaced_custom_object(
                    group="argoproj.io",
                    version="v1alpha1",
                    plural="rollouts",
                    **kwargs,
                )
            ),
            extract_containers=_extract_containers,
        )

    def _list_deploymentconfig(self) -> list[K8sObjectData]:
        # NOTE: Using custom objects API returns dicts, but all other APIs return objects
        # We need to handle this difference using a small wrapper
        return self._list_scannable_objects(
            kind="DeploymentConfig",
            all_namespaces_request=lambda **kwargs: ObjectLikeDict(
                self.custom_objects.list_cluster_custom_object(
                    group="apps.openshift.io",
                    version="v1",
                    plural="deploymentconfigs",
                    **kwargs,
                )
            ),
            namespaced_request=lambda **kwargs: ObjectLikeDict(
                self.custom_objects.list_namespaced_custom_object(
                    group="apps.openshift.io",
                    version="v1",
                    plural="deploymentconfigs",
                    **kwargs,
                )
            ),
            extract_containers=lambda item: item.spec.template.spec.containers,
        )

    def _list_all_statefulsets(self) -> list[K8sObjectData]:
        return self._list_scannable_objects(
            kind="StatefulSet",
            all_namespaces_request=self.apps.list_stateful_set_for_all_namespaces,
            namespaced_request=self.apps.list_namespaced_stateful_set,
            extract_containers=lambda item: item.spec.template.spec.containers,
        )

    def _list_all_daemon_set(self) -> list[K8sObjectData]:
        return self._list_scannable_objects(
            kind="DaemonSet",
            all_namespaces_request=self.apps.list_daemon_set_for_all_namespaces,
            namespaced_request=self.apps.list_namespaced_daemon_set,
            extract_containers=lambda item: item.spec.template.spec.containers,
        )

    def _list_all_jobs(self) -> list[K8sObjectData]:
        return self._list_scannable_objects(
            kind="Job",
            all_namespaces_request=self.batch.list_job_for_all_namespaces,
            namespaced_request=self.batch.list_namespaced_job,
            extract_containers=lambda item: item.spec.template.spec.containers,
            # NOTE: If the job has ownerReference and it is a CronJob, then we should skip it
            filter_workflows=lambda item: not any(
                owner.kind == "CronJob" for owner in item.metadata.owner_references or []
            ),
        )

    def _list_all_cronjobs(self) -> list[K8sObjectData]:
        return self._list_scannable_objects(
            kind="CronJob",
            all_namespaces_request=self.batch.list_cron_job_for_all_namespaces,
            namespaced_request=self.batch.list_namespaced_cron_job,
            extract_containers=lambda item: item.spec.job_template.spec.template.spec.containers,
        )

    async def __list_hpa_v1(self) -> dict[HPAKey, HPAData]:
        loop = asyncio.get_running_loop()
        res = await loop.run_in_executor(
            self.executor,
            lambda: self._list_namespaced_or_global_objects(
                kind="HPA-v1",
                all_namespaces_request=self.autoscaling_v1.list_horizontal_pod_autoscaler_for_all_namespaces,
                namespaced_request=self.autoscaling_v1.list_namespaced_horizontal_pod_autoscaler,
            ),
        )
        return {
            (
                hpa.metadata.namespace,
                hpa.spec.scale_target_ref.kind,
                hpa.spec.scale_target_ref.name,
            ): HPAData(
                min_replicas=hpa.spec.min_replicas,
                max_replicas=hpa.spec.max_replicas,
                current_replicas=hpa.status.current_replicas,
                desired_replicas=hpa.status.desired_replicas,
                target_cpu_utilization_percentage=hpa.spec.target_cpu_utilization_percentage,
                target_memory_utilization_percentage=None,
            )
            async for hpa in res
        }

    async def __list_hpa_v2(self) -> dict[HPAKey, HPAData]:
        res = await self._list_namespaced_or_global_objects(
            kind="HPA-v2",
            all_namespaces_request=self.autoscaling_v2.list_horizontal_pod_autoscaler_for_all_namespaces,
            namespaced_request=self.autoscaling_v2.list_namespaced_horizontal_pod_autoscaler,
        )

        def __get_metric(hpa: V2HorizontalPodAutoscaler, metric_name: str) -> Optional[float]:
            return next(
                (
                    metric.resource.target.average_utilization
                    for metric in hpa.spec.metrics
                    if metric.type == "Resource" and metric.resource.name == metric_name
                ),
                None,
            )

        return {
            (
                hpa.metadata.namespace,
                hpa.spec.scale_target_ref.kind,
                hpa.spec.scale_target_ref.name,
            ): HPAData(
                min_replicas=hpa.spec.min_replicas,
                max_replicas=hpa.spec.max_replicas,
                current_replicas=hpa.status.current_replicas,
                desired_replicas=hpa.status.desired_replicas,
                target_cpu_utilization_percentage=__get_metric(hpa, "cpu"),
                target_memory_utilization_percentage=__get_metric(hpa, "memory"),
            )
            for hpa in res
        }

    # TODO: What should we do in case of other metrics bound to the HPA?
    async def __list_hpa(self) -> dict[HPAKey, HPAData]:
        """List all HPA objects in the cluster.

        Returns:
            dict[tuple[str, str], HPAData]: A dictionary of HPA objects, indexed by scaleTargetRef (kind, name).
        """

        try:
            # Try to use V2 API first
            return await self.__list_hpa_v2()
        except ApiException as e:
            if e.status != 404:
                # If the error is other than not found, then re-raise it.
                raise

            # If V2 API does not exist, fall back to V1
            return await self.__list_hpa_v1()

    async def _try_list_hpa(self) -> dict[HPAKey, HPAData]:
        try:
            return await self.__list_hpa()
        except Exception as e:
            logger.exception(f"Error trying to list hpa in cluster {self.cluster}: {e}")
            logger.error(
                "Will assume that there are no HPA. "
                "Be careful as this may lead to inaccurate results if object actually has HPA."
            )
            return {}
