import logging
from threading import Lock
from typing import (
    Dict,
    Optional,
    Sequence,
)

from opentelemetry import trace
from opentelemetry.trace import Span
from service_locator_api.shared.exception import RegistrationNotFound
from ._KeepAliveError import KeepAliveError
from ._Registration import Registration
from ._RegistrationRepository import RegistrationRepository

logger = logging.getLogger("ZeroMqServiceLocator.InMemoryRegistrationRepository")
tracer = trace.get_tracer(__name__)


class InMemoryRegistrationRepository(RegistrationRepository):
    def __init__(self):
        self.services: Dict[str, RoundRobinList] = {}
        self.registrations: Dict[str, Registration] = {}

        self.lock = Lock()

    def add_registration(self, registration: Registration):
        with self.lock:
            if registration.service_name not in self.services:
                self.services[registration.service_name] = RoundRobinList()

            self.services[registration.service_name].append(registration.service_address)

            self.registrations[registration.reg_id] = registration

    def lookup(self, service_name: str) -> Optional[str]:
        with self.lock:
            if service_name not in self.services:
                return None

            address_to_return = self.services[service_name].get()

            return address_to_return

    def keep_alive(self, reg_id: str):
        try:
            with self.lock:
                registration_info = self.registrations[reg_id]
                registration_info.keep_alive()

            logger.debug("Keep alive processed: %s", registration_info)
            logger.debug("Current registrations: %s", self.registrations)
            logger.debug("Current services: %s", self.services)
        except KeyError as k:
            if k.args[0] == reg_id:
                raise RegistrationNotFound()

            logger.debug("Current registrations: %s", self.registrations)
            logger.debug("Current services: %s", self.services)
            raise KeepAliveError("The registration is active but the service is not registered.".format(reg_id))

    def clean_registrations(self):
        logger.debug("Running clean")
        logger.debug("Current registrations: %s", self.registrations)
        logger.debug("Current services: %s", self.services)
        with self.lock:
            with tracer.start_as_current_span("RegistrationsCleaned"):
                to_remove = []
                services_to_keep = {}
                for reg_id in self.registrations.keys():
                    registration_info = self.registrations[reg_id]
                    if registration_info.is_expired():
                        to_remove.append(registration_info)
                    else:
                        services_to_keep[registration_info.service_name] = registration_info.service_name

                for registration_info in to_remove:
                    logger.debug("Removing expired registration: %s", registration_info)

                    with tracer.start_as_current_span("RegistrationCleaned") as span:  # type: Span
                        span.set_attribute("registrationId", registration_info.reg_id)
                        del self.registrations[registration_info.reg_id]

                    self.services[registration_info.service_name].remove(registration_info.service_address)
                    if len(self.services[registration_info.service_name]) == 0:
                        del self.services[registration_info.service_name]

        logger.debug("Cleaned services: %s", self.services)
        logger.debug("Cleaned registrations: %s", self.registrations)

    def get_all_registrations(self) -> Sequence[Registration]:
        return tuple(list(self.registrations.values()))

    def get_registered_services(self) -> Sequence[str]:
        return tuple(self.services.keys())


class RoundRobinList(object):
    def __init__(self):
        self.index = -1
        self.data = []

    def append(self, data):
        self.data.append(data)

    def remove(self, data):
        self.data.remove(data)

    def get(self):
        self.index += 1
        if self.index >= len(self.data):
            self.index = 0

        return self.data[self.index]

    def __len__(self):
        return len(self.data)
