# Copyright 2015-2016 Yelp Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import abc
import collections
import csv
import logging
import random
from typing import Any
from typing import cast
from typing import Collection
from typing import DefaultDict
from typing import Dict
from typing import Iterable
from typing import List
from typing import MutableMapping
from typing import NamedTuple
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import TypeVar
from typing import Union

import requests
from kubernetes.client import V1Node
from kubernetes.client import V1Pod
from mypy_extensions import TypedDict

from paasta_tools import envoy_tools
from paasta_tools import kubernetes_tools
from paasta_tools import long_running_service_tools
from paasta_tools import mesos_tools
from paasta_tools.long_running_service_tools import LongRunningServiceConfig
from paasta_tools.mesos.exceptions import NoSlavesAvailableError
from paasta_tools.monitoring_tools import ReplicationChecker
from paasta_tools.utils import compose_job_id
from paasta_tools.utils import DEFAULT_SOA_DIR
from paasta_tools.utils import DeployBlacklist
from paasta_tools.utils import get_user_agent
from paasta_tools.utils import SystemPaastaConfig


class HaproxyBackend(TypedDict, total=False):
    check_code: str
    check_duration: str
    check_status: str
    lastchg: str
    pxname: str
    svname: str
    status: str


log = logging.getLogger(__name__)


def retrieve_haproxy_csv(
    synapse_host: str, synapse_port: int, synapse_haproxy_url_format: str, scope: str
) -> Iterable[Dict[str, str]]:
    """Retrieves the haproxy csv from the haproxy web interface

    :param synapse_host: A host that this check should contact for replication information.
    :param synapse_port: A integer that this check should contact for replication information.
    :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
    :param scope: scope
    :returns reader: a csv.DictReader object
    """
    synapse_uri = synapse_haproxy_url_format.format(
        host=synapse_host, port=synapse_port, scope=scope
    )

    # timeout after 1 second and retry 3 times
    haproxy_request = requests.Session()
    haproxy_request.headers.update({"User-Agent": get_user_agent()})
    haproxy_request.mount("http://", requests.adapters.HTTPAdapter(max_retries=3))
    haproxy_request.mount("https://", requests.adapters.HTTPAdapter(max_retries=3))
    haproxy_response = haproxy_request.get(synapse_uri, timeout=1)
    haproxy_data = haproxy_response.text
    reader = csv.DictReader(haproxy_data.splitlines())
    return reader


def get_backends(
    service: str, synapse_host: str, synapse_port: int, synapse_haproxy_url_format: str
) -> List[HaproxyBackend]:
    """Fetches the CSV from haproxy and returns a list of backends,
    regardless of their state.

    :param service: If None, return backends for all services, otherwise only return backends for this particular
                    service.
    :param synapse_host: A host that this check should contact for replication information.
    :param synapse_port: A integer that this check should contact for replication information.
    :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
    :returns backends: A list of dicts representing the backends of all
                       services or the requested service
    """
    if service:
        services = [service]
    else:
        services = None
    return get_multiple_backends(
        services,
        synapse_host=synapse_host,
        synapse_port=synapse_port,
        synapse_haproxy_url_format=synapse_haproxy_url_format,
    )


def get_multiple_backends(
    services: Optional[Collection[str]],
    synapse_host: str,
    synapse_port: int,
    synapse_haproxy_url_format: str,
) -> List[HaproxyBackend]:
    """Fetches the CSV from haproxy and returns a list of backends,
    regardless of their state.

    :param services: If None, return backends for all services, otherwise only return backends for these particular
                     services.
    :param synapse_host: A host that this check should contact for replication information.
    :param synapse_port: A integer that this check should contact for replication information.
    :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
    :returns backends: A list of dicts representing the backends of all
                       services or the requested service
    """

    if services is not None and len(services) == 1:
        (scope,) = services
    else:
        # Maybe if there's like two or three services we could make two queries, or find the longest common substring.
        # For now let's just hope this is rare and fetch all data.
        scope = ""

    reader = retrieve_haproxy_csv(
        synapse_host,
        synapse_port,
        synapse_haproxy_url_format=synapse_haproxy_url_format,
        scope=scope,
    )
    backends = []

    for line in reader:
        # clean up two irregularities of the CSV output, relative to
        # DictReader's behavior there's a leading "# " for no good reason:
        line["pxname"] = line.pop("# pxname")
        # and there's a trailing comma on every line:
        line.pop("")

        # Look for the service in question and ignore the fictional
        # FRONTEND/BACKEND hosts, use starts_with so that hosts that are UP
        # with 1/X healthchecks to go before going down get counted as UP:
        ha_slave, ha_service = line["svname"], line["pxname"]
        if (services is None or ha_service in services) and ha_slave not in (
            "FRONTEND",
            "BACKEND",
        ):
            backends.append(cast(HaproxyBackend, line))

    return backends


def load_smartstack_info_for_service(
    service: str,
    namespace: str,
    blacklist: DeployBlacklist,
    system_paasta_config: SystemPaastaConfig,
    soa_dir: str = DEFAULT_SOA_DIR,
) -> Dict[str, Dict[str, int]]:
    """Retrieves number of available backends for given service

    :param service: A service name
    :param namespace: A Smartstack namespace
    :param blacklist: A list of blacklisted location tuples in the form (location, value)
    :param system_paasta_config: A SystemPaastaConfig object representing the system configuration.
    :param soa_dir: SOA dir
    :returns: a dictionary of the form

    ::

        {
          'location_type': {
              'unique_location_name': {
                  'service.instance': <# ofavailable backends>
              },
              'other_unique_location_name': ...
          }
        }

    """
    service_namespace_config = long_running_service_tools.load_service_namespace_config(
        service=service, namespace=namespace, soa_dir=soa_dir
    )
    discover_location_type = service_namespace_config.get_discover()
    return get_smartstack_replication_for_attribute(
        attribute=discover_location_type,
        service=service,
        namespace=namespace,
        blacklist=blacklist,
        system_paasta_config=system_paasta_config,
    )


def get_smartstack_replication_for_attribute(
    attribute: str,
    service: str,
    namespace: str,
    blacklist: DeployBlacklist,
    system_paasta_config: SystemPaastaConfig,
) -> Dict[str, Dict[str, int]]:
    """Loads smartstack replication from a host with the specified attribute

    :param attribute: a Mesos attribute
    :param service: A service name, like 'example_service'
    :param namespace: A particular smartstack namespace to inspect, like 'main'
    :param blacklist: A list of blacklisted location tuples in the form of (location, value)
    :param system_paasta_config: A SystemPaastaConfig object representing the system configuration.
    :returns: a dictionary of the form {'<unique_attribute_value>': <smartstack replication hash>}
              (the dictionary will contain keys for unique all attribute values)
    """
    replication_info = {}
    filtered_slaves = mesos_tools.get_all_slaves_for_blacklist_whitelist(
        blacklist=blacklist, whitelist=None
    )
    if not filtered_slaves:
        raise NoSlavesAvailableError

    attribute_slave_dict = mesos_tools.get_mesos_slaves_grouped_by_attribute(
        slaves=filtered_slaves, attribute=attribute
    )

    full_name = compose_job_id(service, namespace)

    for value, hosts in attribute_slave_dict.items():
        # arbitrarily choose the first host with a given attribute to query for replication stats
        synapse_host = hosts[0]["hostname"]
        repl_info = get_replication_for_services(
            synapse_host=synapse_host,
            synapse_port=system_paasta_config.get_synapse_port(),
            synapse_haproxy_url_format=system_paasta_config.get_synapse_haproxy_url_format(),
            services=[full_name],
        )
        replication_info[value] = repl_info

    return replication_info


def get_replication_for_all_services(
    synapse_host: str, synapse_port: int, synapse_haproxy_url_format: str
) -> Dict[str, int]:
    """Returns the replication level for all services known to this synapse haproxy

    :param synapse_host: The host that this check should contact for replication information.
    :param synapse_port: The port that this check should contact for replication information.
    :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
    :returns available_instance_counts: A dictionary mapping the service names
                                        to an integer number of available replicas.
    """
    backends = get_multiple_backends(
        services=None,
        synapse_host=synapse_host,
        synapse_port=synapse_port,
        synapse_haproxy_url_format=synapse_haproxy_url_format,
    )
    return collections.Counter([b["pxname"] for b in backends if backend_is_up(b)])


def get_replication_for_services(
    synapse_host: str,
    synapse_port: int,
    synapse_haproxy_url_format: str,
    services: Collection[str],
) -> Dict[str, int]:
    """Returns the replication level for the provided services

    This check is intended to be used with an haproxy load balancer, and
    relies on the implementation details of that choice.

    :param synapse_host: The host that this check should contact for replication information.
    :param synapse_port: The port that this check should contact for replication information.
    :param synapse_haproxy_url_format: The format of the synapse haproxy URL.
    :param services: A list of strings that are the service names
                          that should be checked for replication.

    :returns available_instance_counts: A dictionary mapping the service names
                                  to an integer number of available
                                  replicas
    :returns None: If it cannot connect to the specified synapse host and port
    """
    backends = get_multiple_backends(
        services=services,
        synapse_host=synapse_host,
        synapse_port=synapse_port,
        synapse_haproxy_url_format=synapse_haproxy_url_format,
    )

    counter = collections.Counter([b["pxname"] for b in backends if backend_is_up(b)])
    return {sn: counter[sn] for sn in services}


def backend_is_up(backend: HaproxyBackend) -> bool:
    """Returns whether a server is receiving traffic in HAProxy.

    :param backend: backend dict, like one of those returned by smartstack_tools.get_multiple_backends.

    :returns is_up: Whether the backend is in a state that receives traffic.
    """
    return str(backend["status"]).startswith("UP")


def ip_port_hostname_from_svname(svname: str) -> Tuple[str, int, str]:
    """This parses the haproxy svname that smartstack creates.
    In old versions of synapse, this is in the format ip:port_hostname.
    In versions newer than dd5843c987740a5d5ce1c83b12b258b7253784a8 it is
    hostname_ip:port

    :param svname: A svname, in either of the formats described above
    :returns ip_port_hostname: A tuple of ip, port, hostname.
    """
    # split into parts
    parts = set(svname.split("_"))

    # find those that can be split by : - this is the ip:port
    # there will only be 1 of these
    ip_ports = {part for part in parts if len(part.split(":")) == 2}

    # the one *not* in the list is the hostname
    hostname = parts.difference(ip_ports).pop()

    ip, port = ip_ports.pop().split(":")
    return ip, int(port), hostname


def are_services_up_on_ip_port(
    synapse_host: str,
    synapse_port: int,
    synapse_haproxy_url_format: str,
    services: Collection[str],
    host_ip: str,
    host_port: int,
) -> bool:
    backends = get_multiple_backends(
        services,
        synapse_host=synapse_host,
        synapse_port=synapse_port,
        synapse_haproxy_url_format=synapse_haproxy_url_format,
    )
    backends_by_ip_port: DefaultDict[
        Tuple[str, int], List[HaproxyBackend]
    ] = collections.defaultdict(list)

    for backend in backends:
        ip, port, _ = ip_port_hostname_from_svname(backend["svname"])
        backends_by_ip_port[ip, port].append(backend)

    backends_on_ip = backends_by_ip_port[host_ip, host_port]
    # any backend being up is okay because a previous backend
    # may have had the same IP and synapse only removes them completely
    # after some time
    services_with_atleast_one_backend_up = {service: False for service in services}
    for service in services:
        for be in backends_on_ip:
            if be["pxname"] == service and backend_is_up(be):
                services_with_atleast_one_backend_up[service] = True
    return all(services_with_atleast_one_backend_up.values())


def match_backends_and_pods(
    backends: Iterable[HaproxyBackend], pods: Iterable[V1Pod]
) -> List[Tuple[Optional[HaproxyBackend], Optional[V1Pod]]]:
    """Returns tuples of matching (backend, pod) pairs, as matched by IP. Each backend will be listed exactly
    once. If a backend does not match with a pod, (backend, None) will be included.
    If a pod's IP does not match with any backends, (None, pod) will be included.

    :param backends: An iterable of haproxy backend dictionaries, e.g. the list returned by
                     smartstack_tools.get_multiple_backends.
    :param pods: An iterable of V1Pod objects.
    """

    # { ip : [backend1, backend2], ... }
    backends_by_ip: DefaultDict[str, List[HaproxyBackend]] = collections.defaultdict(
        list
    )
    backend_pod_pairs = []

    for backend in backends:
        ip, port, _ = ip_port_hostname_from_svname(backend["svname"])
        backends_by_ip[ip].append(backend)

    for pod in pods:
        ip = pod.status.pod_ip
        for backend in backends_by_ip.pop(ip, [None]):
            backend_pod_pairs.append((backend, pod))

    # we've been popping in the above loop, so anything left didn't match a k8s pod.
    for backends in backends_by_ip.values():
        for backend in backends:
            backend_pod_pairs.append((backend, None))

    return backend_pod_pairs


_MesosSlaveDict = TypeVar(
    "_MesosSlaveDict", bound=Dict
)  # no type has been defined in mesos_tools for these yet.


class DiscoveredHost(NamedTuple):
    hostname: str
    pool: str


class ServiceDiscoveryProvider(abc.ABC):

    NAME = "..."

    @abc.abstractmethod
    def get_replication_for_all_services(self, hostname: str) -> Dict[str, int]:
        ...


class SmartstackServiceDiscovery(ServiceDiscoveryProvider):

    NAME = "Smartstack"

    def __init__(self, system_paasta_config: SystemPaastaConfig) -> None:
        self._synapse_port = system_paasta_config.get_synapse_port()
        self._synapse_haproxy_url_format = (
            system_paasta_config.get_synapse_haproxy_url_format()
        )

    def get_replication_for_all_services(self, hostname: str) -> Dict[str, int]:
        return get_replication_for_all_services(
            synapse_host=hostname,
            synapse_port=self._synapse_port,
            synapse_haproxy_url_format=self._synapse_haproxy_url_format,
        )


class EnvoyServiceDiscovery(ServiceDiscoveryProvider):

    NAME = "Envoy"

    def __init__(self, system_paasta_config: SystemPaastaConfig) -> None:
        self._envoy_admin_port = system_paasta_config.get_envoy_admin_port()
        self._envoy_admin_endpoint_format = (
            system_paasta_config.get_envoy_admin_endpoint_format()
        )

    def get_replication_for_all_services(self, hostname: str) -> Dict[str, int]:
        return envoy_tools.get_replication_for_all_services(
            envoy_host=hostname,
            envoy_admin_port=self._envoy_admin_port,
            envoy_admin_endpoint_format=self._envoy_admin_endpoint_format,
        )


def get_service_discovery_providers(
    system_paasta_config: SystemPaastaConfig,
) -> List[ServiceDiscoveryProvider]:
    providers: List[ServiceDiscoveryProvider] = []
    for name, _ in system_paasta_config.get_service_discovery_providers().items():
        if name == "smartstack":
            providers.append(SmartstackServiceDiscovery(system_paasta_config))
        elif name == "envoy":
            providers.append(EnvoyServiceDiscovery(system_paasta_config))
        else:
            log.warn("unknown provider")
    return providers


class BaseReplicationChecker(ReplicationChecker):
    """Base class for checking replication. Extendable for different frameworks.

    Optimized for multiple queries. Gets the list of backends from service
    discovery provider only once per location and reuse it in all subsequent
    calls of BaseReplicationChecker.get_replication_for_instance().

    get_allowed_locations_and_hosts must be implemented in sub class

    A list of service discovery providers to collect information about
    instances and their status must be provided as
    `service_discovery_providers`.
    """

    def __init__(
        self,
        system_paasta_config: SystemPaastaConfig,
        service_discovery_providers: Iterable[ServiceDiscoveryProvider],
    ) -> None:
        self._system_paasta_config = system_paasta_config
        self._cache: Dict[Tuple[str, str], Dict[str, int]] = {}
        self._service_discovery_providers = service_discovery_providers

    @abc.abstractmethod
    def get_allowed_locations_and_hosts(
        self, instance_config: LongRunningServiceConfig
    ) -> Dict[str, Sequence[DiscoveredHost]]:
        ...

    def get_replication_for_instance(
        self, instance_config: LongRunningServiceConfig
    ) -> Dict[str, Dict[str, Dict[str, int]]]:
        """Returns the number of registered instances in each discoverable
        location for each service dicrovery provider.

        :param instance_config: An instance of LongRunningServiceConfig.
        :returns: a dict {'service_discovery_provider': {'location_type': {'service.instance': int}}}
        """
        replication_infos = {}
        for provider in self._service_discovery_providers:
            replication_info = {}
            attribute_host_dict = self.get_allowed_locations_and_hosts(instance_config)
            instance_pool = instance_config.get_pool()
            for location, hosts in attribute_host_dict.items():
                # Try to get information from all available hosts in the pool before giving up
                hostnames = self.get_hostnames_in_pool(hosts, instance_pool)
                for hostname in hostnames:
                    try:
                        replication_info[location] = self._get_replication_info(
                            location, hostname, instance_config, provider
                        )
                        break
                    except Exception as e:
                        log.warn(
                            f"Error while getting replication info for {location} from {hostname}: {e}"
                        )
                        if hostname == hostnames[-1]:
                            # Last hostname failed, giving up
                            raise
            replication_infos[provider.NAME] = replication_info
        return replication_infos

    def get_first_host_in_pool(self, hosts: Sequence[DiscoveredHost], pool: str) -> str:
        for host in hosts:
            if host.pool == pool:
                return host.hostname
        return hosts[0].hostname

    def get_hostname_in_pool(self, hosts: Sequence[DiscoveredHost], pool: str) -> str:
        return random.choice(self.get_hostnames_in_pool(hosts, pool))

    def get_hostnames_in_pool(
        self, hosts: Sequence[DiscoveredHost], pool: str
    ) -> Sequence[str]:
        hostnames = []
        for host in hosts:
            if host.pool == pool:
                hostnames.append(host.hostname)
        if len(hostnames) == 0:
            hostnames.append(hosts[0].hostname)
        return hostnames

    def _get_replication_info(
        self,
        location: str,
        hostname: str,
        instance_config: LongRunningServiceConfig,
        provider: ServiceDiscoveryProvider,
    ) -> Dict[str, int]:
        """Returns service.instance and the number of instances registered in smartstack
        at the location as a dict.

        :param location: A string that identifies a habitat, a region and etc.
        :param hostname: A mesos slave hostname to read replication information from.
        :param instance_config: An instance of LongRunningServiceConfig.
        :returns: A dict {"service.instance": number_of_instances}.
        """
        full_name = compose_job_id(instance_config.service, instance_config.instance)
        key = (location, provider.NAME)
        replication_info = self._cache.get(key)
        if replication_info is None:
            replication_info = provider.get_replication_for_all_services(hostname)
            self._cache[key] = replication_info
        return {full_name: replication_info[full_name]}


class KubeSmartstackEnvoyReplicationChecker(BaseReplicationChecker):
    def __init__(
        self, nodes: Sequence[V1Node], system_paasta_config: SystemPaastaConfig
    ) -> None:
        self.nodes = nodes
        super().__init__(
            system_paasta_config=system_paasta_config,
            service_discovery_providers=get_service_discovery_providers(
                system_paasta_config
            ),
        )

    def get_allowed_locations_and_hosts(
        self, instance_config: LongRunningServiceConfig
    ) -> Dict[str, Sequence[DiscoveredHost]]:
        discover_location_type = kubernetes_tools.load_service_namespace_config(
            service=instance_config.service,
            namespace=instance_config.get_nerve_namespace(),
            soa_dir=instance_config.soa_dir,
        ).get_discover()

        attribute_to_nodes = kubernetes_tools.get_nodes_grouped_by_attribute(
            nodes=self.nodes, attribute=discover_location_type
        )
        ret: Dict[str, Sequence[DiscoveredHost]] = {}
        for attr, nodes in attribute_to_nodes.items():
            ret[attr] = [
                DiscoveredHost(
                    hostname=node.metadata.labels["yelp.com/hostname"],
                    pool=node.metadata.labels["yelp.com/pool"],
                )
                for node in nodes
            ]
        return ret


def build_smartstack_location_dict(
    location: str,
    matched_backends_and_tasks: List[
        Tuple[
            Optional[HaproxyBackend],
            Optional[V1Pod],
        ]
    ],
    should_return_individual_backends: bool = False,
) -> MutableMapping[str, Any]:
    running_backends_count = 0
    backends = []
    for backend, task in matched_backends_and_tasks:
        if backend is None:
            continue
        if backend_is_up(backend):
            running_backends_count += 1
        if should_return_individual_backends:
            backends.append(build_smartstack_backend_dict(backend, task))

    return {
        "name": location,
        "running_backends_count": running_backends_count,
        "backends": backends,
    }


def build_smartstack_backend_dict(
    smartstack_backend: HaproxyBackend,
    task: Union[V1Pod],
) -> MutableMapping[str, Any]:
    svname = smartstack_backend["svname"]
    if isinstance(task, V1Pod):
        node_hostname = svname.split("_")[0]
        pod_ip = svname.split("_")[1].split(":")[0]
        hostname = f"{node_hostname}:{pod_ip}"
    else:
        hostname = svname.split("_")[0]
    port = svname.split("_")[-1].split(":")[-1]

    smartstack_backend_dict = {
        "hostname": hostname,
        "port": int(port),
        "status": smartstack_backend["status"],
        "check_status": smartstack_backend["check_status"],
        "check_code": smartstack_backend["check_code"],
        "last_change": int(smartstack_backend["lastchg"]),
        "has_associated_task": task is not None,
    }

    check_duration = smartstack_backend["check_duration"]
    if check_duration:
        smartstack_backend_dict["check_duration"] = int(check_duration)

    return smartstack_backend_dict
