import logging
from datetime import datetime, timedelta
from threading import Lock, Thread, Condition
from typing import Dict
from uuid import uuid4

from ._KeepAliveError import KeepAliveError

logger = logging.getLogger("ZeroMqServiceLocator.ServiceLocator")


class RegistrationInfo(object):
    def __init__(self, reg_id, service_name, service_address, registration_expiration):
        self.reg_id = reg_id
        self.service_name = service_name
        self.service_address = service_address
        self.registration_expiration = registration_expiration
        self.last_keep_alive = datetime.now()

    def keep_alive(self):
        self.last_keep_alive = datetime.now()

    def is_expired(self):
        return self.last_keep_alive + timedelta(seconds=self.registration_expiration) < datetime.now()

    def __str__(self):
        return str({
            "reg_id": self.reg_id,
            "service_name": self.service_name,
            "service_address": self.service_address,
            "registration_expiration": self.registration_expiration,
            "last_keep_alive": self.last_keep_alive.isoformat()
        })

    def __repr__(self) -> str:
        return self.__str__()


class ServiceLocator(object):
    def __init__(self, clean_interval=30, registration_expiration=20):
        self.clean_interval = clean_interval
        self.registration_expiration = registration_expiration

        self.lock = Lock()
        self.services = {}  # type: Dict[str, str]
        self.registrations = {}  # type: Dict[str, RegistrationInfo]

        self.clean_thread_running = True
        self.clean_thread_condition = Condition()
        self.clean_thread = Thread(target=self._run_clean_expired_registrations,
                                   name="ServiceLocatorCleanExpiredRegistrations")
        self.clean_thread.daemon = True
        self.clean_thread.start()

    def register_service(self, service_name, address):
        address = address.replace("*", "localhost")
        reg_id = str(uuid4())
        logger.debug("Created registration id: %s", reg_id)
        registration_info = RegistrationInfo(reg_id, service_name, address, self.registration_expiration)

        with self.lock:
            self.services[service_name] = address
            self.registrations[reg_id] = registration_info

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

        return reg_id

    def locate_service(self, service_name):
        with self.lock:
            return self.services.get(service_name, None)

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

            if current_address != registration_info.service_address:
                raise KeepAliveError("The address this registration refers to is not the current address. "
                                     "Current Address: {}".format(current_address))
            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 KeepAliveError("The registration id {} is not registered anymore".format(reg_id))

            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 dispose(self):
        self.clean_thread_running = False
        self.clean_thread_condition.notify_all()
        self.clean_thread.join()

    def _run_clean_expired_registrations(self):
        while self.clean_thread_condition:
            try:
                self._clean_expired_registrations()
                with self.clean_thread_condition:
                    self.clean_thread_condition.wait(self.clean_interval)
            except:
                logger.exception("Error cleaning expired registrations")

    def _clean_expired_registrations(self):
        logger.debug("Running clean")
        logger.debug("Current registrations: %s", self.registrations)
        logger.debug("Current services: %s", self.services)
        with self.lock:
            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)
                del self.registrations[registration_info.reg_id]
                if registration_info.service_name not in services_to_keep:
                    del self.services[registration_info.service_name]
        logger.debug("Cleaned services: %s", self.services)
        logger.debug("Cleaned registrations: %s", self.registrations)
