import logging
from abc import abstractmethod
from collections.abc import Callable
from contextlib import AbstractContextManager
from threading import Thread
from time import sleep  # noqa
from urllib import request
from urllib.error import URLError

from blueapi.config import ApplicationConfig, ConfigLoader
from bluesky.callbacks import CallbackBase
from bluesky.callbacks.zmq import Proxy, RemoteDispatcher
from bluesky_stomp.messaging import StompClient
from bluesky_stomp.models import Broker
from dodal.log import LOGGER as DODAL_LOGGER
from dodal.log import set_up_all_logging_handlers

from mx_bluesky.common.external_interaction.alerting import set_alerting_service
from mx_bluesky.common.external_interaction.alerting.log_based_service import (
    LoggingAlertService,
)
from mx_bluesky.common.external_interaction.callbacks.common.log_uid_tag_callback import (
    LogUidTaggingCallback,
)
from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
    ZocaloCallback,
)
from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import (
    SampleHandlingCallback,
)
from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
    GridscanISPyBCallback,
    generate_start_info_from_omega_map,
)
from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import (
    GridscanNexusFileCallback,
)
from mx_bluesky.common.utils.log import (
    ISPYB_ZOCALO_CALLBACK_LOGGER,
    NEXUS_LOGGER,
    _get_logging_dirs,
    tag_filter,
)
from mx_bluesky.hyperion.external_interaction.callbacks.alert_on_container_change import (
    AlertOnContainerChange,
)
from mx_bluesky.hyperion.external_interaction.callbacks.robot_actions.ispyb_callback import (
    RobotLoadISPyBCallback,
)
from mx_bluesky.hyperion.external_interaction.callbacks.rotation.ispyb_callback import (
    RotationISPyBCallback,
    generate_start_info_from_ordered_runs,
)
from mx_bluesky.hyperion.external_interaction.callbacks.rotation.nexus_callback import (
    RotationNexusFileCallback,
)
from mx_bluesky.hyperion.external_interaction.callbacks.snapshot_callback import (
    BeamDrawingCallback,
)
from mx_bluesky.hyperion.external_interaction.callbacks.stomp.dispatcher import (
    StompDispatcher,
)
from mx_bluesky.hyperion.parameters.cli import CallbackArgs, parse_callback_args
from mx_bluesky.hyperion.parameters.constants import CONST
from mx_bluesky.hyperion.parameters.gridscan import (
    GridCommonWithHyperionDetectorParams,
    HyperionSpecifiedThreeDGridScan,
)

PING_TIMEOUT_S = 1

LIVENESS_POLL_SECONDS = 1
ERROR_LOG_BUFFER_LINES = 5000
HYPERION_PING_INTERVAL_S = 19


def create_gridscan_callbacks() -> tuple[
    GridscanNexusFileCallback, GridscanISPyBCallback
]:
    return (
        GridscanNexusFileCallback(param_type=HyperionSpecifiedThreeDGridScan),
        GridscanISPyBCallback(
            param_type=GridCommonWithHyperionDetectorParams,
            emit=ZocaloCallback(
                CONST.PLAN.DO_FGS, CONST.ZOCALO_ENV, generate_start_info_from_omega_map
            ),
        ),
    )


def create_rotation_callbacks() -> tuple[
    RotationNexusFileCallback, RotationISPyBCallback
]:
    return (
        RotationNexusFileCallback(),
        RotationISPyBCallback(
            emit=ZocaloCallback(
                CONST.PLAN.ROTATION_MULTI,
                CONST.ZOCALO_ENV,
                generate_start_info_from_ordered_runs,
            )
        ),
    )


def setup_callbacks() -> list[CallbackBase]:
    rot_nexus_cb, rot_ispyb_cb = create_rotation_callbacks()
    snapshot_cb = BeamDrawingCallback(emit=rot_ispyb_cb)
    return [
        *create_gridscan_callbacks(),
        rot_nexus_cb,
        snapshot_cb,
        LogUidTaggingCallback(),
        RobotLoadISPyBCallback(),
        SampleHandlingCallback(),
        AlertOnContainerChange(),
    ]


def setup_logging(dev_mode: bool):
    for logger, filename in [
        (ISPYB_ZOCALO_CALLBACK_LOGGER, "hyperion_ispyb_callback.log"),
        (NEXUS_LOGGER, "hyperion_nexus_callback.log"),
    ]:
        logging_path, debug_logging_path = _get_logging_dirs(dev_mode)
        if logger.handlers == []:
            handlers = set_up_all_logging_handlers(
                logger,
                logging_path,
                filename,
                dev_mode,
                ERROR_LOG_BUFFER_LINES,
                CONST.GRAYLOG_PORT,
                debug_logging_path,
            )
            handlers["graylog_handler"].addFilter(tag_filter)
    log_info(f"Loggers initialised with dev_mode={dev_mode}")
    nexgen_logger = logging.getLogger("nexgen")
    nexgen_logger.parent = NEXUS_LOGGER
    DODAL_LOGGER.parent = ISPYB_ZOCALO_CALLBACK_LOGGER
    log_debug("nexgen logger added to nexus logger")


def log_info(msg, *args, **kwargs):
    ISPYB_ZOCALO_CALLBACK_LOGGER.info(msg, *args, **kwargs)
    NEXUS_LOGGER.info(msg, *args, **kwargs)


def log_debug(msg, *args, **kwargs):
    ISPYB_ZOCALO_CALLBACK_LOGGER.debug(msg, *args, **kwargs)
    NEXUS_LOGGER.debug(msg, *args, **kwargs)


class HyperionCallbackRunner:
    """Runs Nexus, ISPyB and Zocalo callbacks in their own process."""

    def __init__(self, callback_args: CallbackArgs) -> None:
        setup_logging(callback_args.dev_mode)
        log_info("Hyperion callback process started.")
        set_alerting_service(LoggingAlertService(CONST.GRAYLOG_STREAM_ID))

        self.callbacks = setup_callbacks()

        self.watchdog_thread = Thread(
            target=run_watchdog,
            daemon=True,
            name="Watchdog",
            args=[callback_args.watchdog_port],
        )

        self._dispatcher_cm: DispatcherContextMgr
        if callback_args.stomp_config:
            self._dispatcher_cm = StompDispatcherContextMgr(
                callback_args, self.callbacks
            )
        else:
            self._dispatcher_cm = RemoteDispatcherContextMgr(self.callbacks)

    def start(self):
        log_info(f"Launching threads, with callbacks: {self.callbacks}")
        self.watchdog_thread.start()
        with self._dispatcher_cm:
            ping_watchdog_while_alive(self._dispatcher_cm, self.watchdog_thread)


def run_watchdog(watchdog_port: int):
    log_info("Hyperion watchdog keepalive running")
    while True:
        try:
            with request.urlopen(
                f"http://localhost:{watchdog_port}/callbackPing",
                timeout=PING_TIMEOUT_S,
            ) as response:
                if response.status != 200:
                    log_debug(
                        f"Unable to ping Hyperion liveness endpoint, status {response.status}"
                    )
        except URLError as e:
            log_debug("Unable to ping Hyperion liveness endpoint", exc_info=e)
        sleep(HYPERION_PING_INTERVAL_S)


def main(dev_mode=False) -> None:
    callback_args = parse_callback_args()
    callback_args.dev_mode = dev_mode or callback_args.dev_mode
    print(f"In dev mode: {dev_mode}")
    runner = HyperionCallbackRunner(callback_args)
    runner.start()


class DispatcherContextMgr(AbstractContextManager):
    @abstractmethod
    def is_alive(self) -> bool: ...


class RemoteDispatcherContextMgr(DispatcherContextMgr):
    def __init__(self, callbacks: list[CallbackBase]):
        super().__init__()

        self.proxy = Proxy(*CONST.CALLBACK_0MQ_PROXY_PORTS)
        self.proxy_thread = Thread(
            target=self.proxy.start, daemon=True, name="0MQ Proxy"
        )

        self.dispatcher = RemoteDispatcher(
            f"localhost:{CONST.CALLBACK_0MQ_PROXY_PORTS[1]}"
        )

        def start_dispatcher(callbacks: list[Callable]):
            for cb in callbacks:
                self.dispatcher.subscribe(cb)
            self.dispatcher.start()

        self.dispatcher_thread = Thread(
            target=start_dispatcher,
            args=[callbacks],
            daemon=True,
            name="0MQ Dispatcher",
        )
        log_info("Created 0MQ proxy and local RemoteDispatcher.")

    def __enter__(self):
        log_info("Proxy and dispatcher thread launched.")
        self.proxy_thread.start()
        self.dispatcher_thread.start()
        return self

    def __exit__(self, exc_type, exc_value, traceback, /):
        self.dispatcher.stop()
        # proxy has no way to stop

    def is_alive(self):
        return self.proxy_thread.is_alive() and self.dispatcher_thread.is_alive()


class StompDispatcherContextMgr(DispatcherContextMgr):
    def __init__(self, args: CallbackArgs, callbacks: list[CallbackBase]):
        super().__init__()
        loader = ConfigLoader(ApplicationConfig)
        assert args.stomp_config is not None
        loader.use_values_from_yaml(args.stomp_config)
        config = loader.load()
        log_info(
            f"Stomp client configured on {config.stomp.url.host}:{config.stomp.url.port}"
        )
        assert config.stomp.url.host is not None
        assert config.stomp.url.port is not None
        self._stomp_client = StompClient.for_broker(
            broker=Broker(
                host=config.stomp.url.host,
                port=config.stomp.url.port,
                auth=config.stomp.auth,
            )
        )
        self._dispatcher = StompDispatcher(self._stomp_client)
        for cb in callbacks:
            self._dispatcher.subscribe(cb)

    def is_alive(self) -> bool:
        return self._stomp_client.is_connected()

    def __enter__(self):
        self._dispatcher.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback, /):
        self._dispatcher.__exit__(exc_type, exc_value, traceback)


def ping_watchdog_while_alive(
    dispatcher_cm: DispatcherContextMgr, watchdog_thread: Thread
):
    alive = watchdog_thread.is_alive() and dispatcher_cm.is_alive()
    try:
        log_debug("Trying to wait forever on callback and dispatcher threads")
        while alive:
            sleep(LIVENESS_POLL_SECONDS)
            alive = watchdog_thread.is_alive() and dispatcher_cm.is_alive()
    except KeyboardInterrupt:
        log_info("Main thread received interrupt - exiting.")
    else:
        log_info("Proxy or dispatcher thread ended - exiting.")


if __name__ == "__main__":
    main()
