import json
import logging
from logging.handlers import TimedRotatingFileHandler
from os import environ
from pathlib import Path

from dodal.log import (
    ERROR_LOG_BUFFER_LINES,
    CircularMemoryHandler,
    DodalLogHandlers,
    integrate_bluesky_and_ophyd_logging,
    set_up_all_logging_handlers,
    set_up_debug_memory_handler,
    set_up_info_file_handler,
)
from dodal.log import LOGGER as DODAL_LOGGER

LOGGER = logging.getLogger("MX-Bluesky")
LOGGER.setLevel("DEBUG")
LOGGER.parent = DODAL_LOGGER

ISPYB_ZOCALO_CALLBACK_LOGGER = logging.getLogger("ISPyB and Zocalo callbacks")
ISPYB_ZOCALO_CALLBACK_LOGGER.setLevel(logging.DEBUG)

NEXUS_LOGGER = logging.getLogger("NeXus callbacks")
NEXUS_LOGGER.setLevel(logging.DEBUG)

ALL_LOGGERS = [LOGGER, ISPYB_ZOCALO_CALLBACK_LOGGER, NEXUS_LOGGER]

__logger_handlers: DodalLogHandlers | None = None


def format_doc_for_log(doc):
    class _BestEffortEncoder(json.JSONEncoder):
        def default(self, o):
            return repr(o)

    return json.dumps(doc, indent=2, cls=_BestEffortEncoder)


class ExperimentMetadataTagFilter(logging.Filter):
    """When an instance of this custom filter is added to a logging handler, dc_group_id
    and run_id will be tagged in that handlers' log messages."""

    dc_group_id: str | None = None
    run_uid: str | None = None

    def filter(self, record):
        if self.dc_group_id:
            record.dc_group_id = self.dc_group_id
        if self.run_uid:
            record.run_uid = self.run_uid
        return True


tag_filter = ExperimentMetadataTagFilter()


def set_dcgid_tag(dcgid):
    """Set the datacollection group id as a tag on all subsequent log messages.
    Setting to None will remove the tag."""
    tag_filter.dc_group_id = dcgid


def set_uid_tag(uid):
    """Set the unique id as a tag on all subsequent log messages.
    Setting to None will remove the tag."""
    tag_filter.run_uid = uid


def setup_hyperion_blueapi_logging(log_file_name: str):
    """Configure debug logging for hyperion-blueapi.
    Args:
        log_file_name: Base name of the log file.
    """
    dodal_logger = DODAL_LOGGER
    logging_path, debug_logging_path = _get_logging_dirs(False)
    set_up_debug_memory_handler(
        dodal_logger, debug_logging_path, log_file_name, ERROR_LOG_BUFFER_LINES
    )
    set_up_info_file_handler(
        dodal_logger,
        logging_path,
        log_file_name,
    )


def do_default_logging_setup(
    file_name: str,
    graylog_port: int,
    dev_mode: bool = False,
    integrate_all_logs: bool = True,
):
    """Configures dodal logger so that separate debug and info log files are created,
    info logs are sent to Graylog, info logs are streamed to sys.sterr, and logs from ophyd
    and bluesky and ophyd-async are optionally included."""
    logging_path, debug_logging_path = _get_logging_dirs(dev_mode)
    handlers = set_up_all_logging_handlers(
        DODAL_LOGGER,
        logging_path,
        file_name,
        dev_mode,
        ERROR_LOG_BUFFER_LINES,
        graylog_port,
        debug_logging_path,
    )

    if integrate_all_logs:
        integrate_bluesky_and_ophyd_logging(DODAL_LOGGER)

    handlers["graylog_handler"].addFilter(tag_filter)

    global __logger_handlers
    __logger_handlers = handlers


def _get_debug_handler() -> CircularMemoryHandler:
    assert __logger_handlers is not None, (
        "You can only use this after running the default logging setup"
    )
    return __logger_handlers["debug_memory_handler"]


def flush_debug_handler() -> str:
    """Writes the contents of the circular debug log buffer to disk and returns the written filename"""
    handler = _get_debug_handler()
    assert isinstance(handler.target, TimedRotatingFileHandler), (
        "Circular memory handler doesn't have an appropriate fileHandler target"
    )
    handler.flush()
    return handler.target.baseFilename


def _get_logging_dirs(dev_mode: bool) -> tuple[Path, Path]:
    """Get the paths to write the mx_bluesky log files to.

    Log location must be specified in the LOG_DIR environment variable,
    and the debug log location specified in the DEBUG_LOG_DIR environment variable.
    This directory will be created if it is not found.

    Logs are written to ./tmp/logs/bluesky if BEAMLINE environment variable is not found

    Args:
        dev_mode (bool): If True, the logs will be written to /tmp/logs/bluesky if the environment variables are not set
    Returns:
        tuple[Path, Path]: Paths to the standard log file and to the debug log file, for the file handlers to write to
    Raises:
        ValueError: If LOG_DIR or DEBUG_LOG_DIR environment variable is not set and dev_mode is False
    """
    if dev_mode:
        logging_path = Path(environ.get("LOG_DIR", "/tmp/logs/bluesky"))
        debug_logging_path = Path(environ.get("DEBUG_LOG_DIR", "/tmp/logs/bluesky"))
    else:
        logging_dir = environ.get("LOG_DIR")
        if not logging_dir:
            raise ValueError("LOG_DIR environment variable is not set")
        debug_logging_dir = environ.get("DEBUG_LOG_DIR")
        if not debug_logging_dir:
            raise ValueError("DEBUG_LOG_DIR environment variable is not set")
        logging_path = Path(logging_dir)
        debug_logging_path = Path(debug_logging_dir)

    Path.mkdir(logging_path, exist_ok=True, parents=True)
    Path.mkdir(debug_logging_path, exist_ok=True, parents=True)
    return logging_path, debug_logging_path
