Logger

def add_console_handler(
    self, format: str = DEFAULT_FORMAT, level: Any = DEFAULT_LOG_LEVEL
):
    """Adds a console handler

    Args:
        format: the format string of the logger
        level: initial level of the logger for this handler (optional) if not given, the default level is set
    """
    handler: logging.Handler
    if is_ipython():
        handler = _IpythonConsoleHandler()
    else:
        handler = logging.StreamHandler()

    handler.setLevel(self._internal_level_translator(level))

    # Prevent security audit logs from appearing in interactive console (e.g., IPython).
    handler.addFilter(_ExcludeSecurityFilter())

    self._internal_add_handler("CONSOLE", handler, format)

    pass

def add_file_handler(
    self,
    filename: str = DEFAULT_LOG_FILE,
    format: str = DEFAULT_FORMAT,
    level: Any = DEFAULT_LOG_LEVEL,
):
    """Adds a file handler

    Args:
        filename: File to log to
        format: Log format
        level: Initial level of the logger
    """

    handler = logging.FileHandler(filename=filename, mode="a")
    handler.setLevel(self._internal_level_translator(level))

    self._internal_add_handler("FILE", handler, format)

def add_grpc_handler(
    self, on_log: Callable, node_id: Optional[str] = None, level: Any = logging.INFO
):
    """Adds a gRPC handler, to publish error message on a topic

    Args:
        on_log: Provided by higher level GRPC implementation
        node_id: id of the caller (necessary for msg formatting to the researcher)
        level: level of this handler (non-mandatory) level must be lower than ERROR to ensure that the
            research get all ERROR/CRITICAL messages
    """

    handler = _GrpcHandler(
        on_log=on_log,
        node_id=node_id,
    )

    # may be not necessary ?
    handler.setLevel(self._internal_level_translator(level))
    formatter = _GrpcFormatter(node_id)

    handler.setFormatter(formatter)

    # Never transmit security audit logs to the researcher via gRPC.
    handler.addFilter(_ExcludeSecurityFilter())

    self._internal_add_handler("GRPC", handler)

    # as a side effect this will set the minimal level to ERROR
    # FIXME: alitolga: This could cause problems in the logging level,
    # I believe the level should be set when node and researcher get started.
    # self.setLevel(level, "GRPC")

    pass

def add_security_file_handler(
    self,
    filename: str = DEFAULT_SECURITY_LOG_FILE,
    level: Any = logging.INFO,
) -> None:
    """
    Adds a dedicated SECURITY_FILE handler that writes JSONL only.
    Only records emitted with extra {"is_security": True} will be written.
    Automatically rotates daily at midnight. Old logs are never deleted.

    Args:
        filename: Security log file path. Defaults to 'security_audit.log'
        level: Logging level for security events. Defaults to INFO
    """
    handler = TimedRotatingFileHandler(
        filename=filename,
        when="midnight",
        interval=1,
        backupCount=0,  # Keep all old logs, never delete
    )
    handler.setLevel(self._internal_level_translator(level))
    handler.setFormatter(_SecurityFormatter(self._security_defaults))
    handler.addFilter(_SecurityOnlyFilter())
    # Disable buffering to ensure immediate writes
    handler.stream.reconfigure(line_buffering=True)

    handler_name = "SECURITY_FILE"
    # Register under its own key so it doesn't collide with FILE handler
    if handler_name not in self._handlers:
        self._logger.debug(" adding handler for: " + handler_name)
        self._handlers[handler_name] = handler
        self._logger.addHandler(handler)
        self._original_format[handler_name] = None
        self._handler_prefix[handler_name] = ""

def add_syslog_handler(
    self,
    host,
    port,
    protocol,
    facility,
    level,
):
    """Adds a remote syslog handler.

    Args:
        host: remote syslog host
        port: remote syslog port
        protocol: 'udp' or 'tcp'
        facility: syslog facility
        level: fedbiomed log level
    """
    if protocol not in {"udp", "tcp"}:
        raise FedbiomedError(f"Unsupported syslog protocol: {protocol}")

    socktype = socket.SOCK_DGRAM if protocol == "udp" else socket.SOCK_STREAM

    kwargs = {
        "address": (host, port),
        "facility": facility,
        "socktype": socktype,
    }

    handler = SysLogHandler(**kwargs)

    # Optional compatibility knobs from stdlib docs
    handler.append_nul = False
    handler.ident = SYSLOG_IDENT
    handler_key = "SYSLOG"

    handler.setLevel(self._internal_level_translator(level))

    # Forward only security events
    handler.addFilter(_SecurityOnlyFilter())

    # Keep the same JSON structure as the security audit file
    handler.setFormatter(_SecurityFormatter(self._security_defaults))

    self._internal_add_handler(handler_key, handler)

def configure_security(
    self,
    *,
    component_id: Optional[str] = None,
    component_name: Optional[ComponentType] = None,
    fedbiomed_version: Optional[str] = None,
) -> None:
    """Configure default fields that must appear in every security log entry."""
    if component_id is not None:
        self._security_defaults["component_id"] = component_id
    if component_name is not None:
        self._security_defaults["component_name"] = component_name
    if fedbiomed_version is not None:
        self._security_defaults["fedbiomed_version"] = fedbiomed_version

def critical(self, msg, *args, broadcast=False, researcher_id=None, **kwargs):
    """Same as info message"""
    # Merge extra dicts properly to support is_security flag
    extra = kwargs.pop("extra", {})
    extra.update({"researcher_id": researcher_id, "broadcast": broadcast})

    stacklevel = kwargs.pop("stacklevel", 2)
    self._logger.critical(
        msg,
        *args,
        **kwargs,
        extra=extra,
        stacklevel=stacklevel,
    )

def debug(self, msg, *args, broadcast=False, researcher_id=None, **kwargs):
    """Same as info message"""
    # Merge extra dicts properly to support is_security flag
    extra = kwargs.pop("extra", {})
    extra.update({"researcher_id": researcher_id, "broadcast": broadcast})

    stacklevel = kwargs.pop("stacklevel", 2)
    self._logger.debug(
        msg,
        *args,
        **kwargs,
        extra=extra,
        stacklevel=stacklevel,
    )

def del_syslog_handler(self, handler_key: str = "SYSLOG"):
    self._internal_add_handler(handler_key, None)

def error(self, msg, *args, broadcast=False, researcher_id=None, **kwargs):
    """Same as info message"""
    # Merge extra dicts properly to support is_security flag
    extra = kwargs.pop("extra", {})
    extra.update({"researcher_id": researcher_id, "broadcast": broadcast})

    stacklevel = kwargs.pop("stacklevel", 2)
    self._logger.error(
        msg,
        *args,
        **kwargs,
        extra=extra,
        stacklevel=stacklevel,
    )

def info(self, msg, *args, broadcast=False, researcher_id=None, **kwargs):
    """Extends arguments of info message.

    Valid only GrpcHandler is existing

    Args:
        msg: Message to log
        broadcast: Broadcast message to all available researchers
        researcher_id: ID of the researcher that the message will be sent.
            If broadcast True researcher id will be ignored
        **kwargs: Additional keyword arguments. Can include extra={"is_security": True}
            to also write to security log file
    """
    # Merge extra dicts properly to support is_security flag
    extra = kwargs.pop("extra", {})
    extra.update({"researcher_id": researcher_id, "broadcast": broadcast})

    # Ensure caller attribution points to the real call site, not this wrapper.
    # Users may override by explicitly passing stacklevel=...
    stacklevel = kwargs.pop("stacklevel", 2)
    self._logger.info(
        msg,
        *args,
        **kwargs,
        extra=extra,
        stacklevel=stacklevel,
    )

def log(self, level: Any, msg: str):
    """Overrides the logging.log() method to allow the use of string instead of a logging.* level"""

    level = logger._internal_level_translator(level)
    self._logger.log(level, msg)

@contextmanager
def security_context(self, **ctx: Any):
    """
    Bind security context for downstream logging.

    How reset works:
    - SECURITY_CONTEXT.set(new_dict) returns a token referencing the previous value.
    - In the `finally` block we call SECURITY_CONTEXT.reset(token) so the old context
      is restored even if exceptions occur (prevents context leaks).
    """
    current = dict(
        SECURITY_CONTEXT.get() or {}
    )  # copy to avoid mutation across calls
    current.update({k: v for k, v in ctx.items() if v is not None})
    token = SECURITY_CONTEXT.set(current)
    try:
        yield
    finally:
        SECURITY_CONTEXT.reset(token)

def security_event(
    self,
    *,
    operation: Optional[str] = None,
    status: Optional[str] = None,
    researcher_id: Optional[str] = None,
    stacklevel: int = 1,
    **fields: Any,
) -> None:
    """
    Writes one JSON security/audit log line to security file only.

    Always present:
      - node_id, researcher_id, timestamp, operation, status, fedbiomed_version
      - caller_function, caller_module, caller_file, caller_line (automatically captured)

    Values resolved as:
      - explicit args > bound SECURITY_CONTEXT > defaults (for node_id/version)

    Args:
        operation: Name of the operation being logged (e.g., 'dataset_search', 'training_execute')
        status: Status of the operation (e.g., 'success', 'failure', 'pending')
        researcher_id: ID of the researcher performing the operation
        stacklevel: How many frames to skip when attributing the caller.
            Use higher values when calling `security_event` from within wrappers.
        **fields: Additional fields to log (e.g., dataset_id, experiment_id, node_id)
    """
    ctx = dict(SECURITY_CONTEXT.get() or {})

    op = operation or ctx.get("operation")
    st = status or ctx.get("status")
    rid = researcher_id or ctx.get("researcher_id")

    # Capture caller information from the stack.
    # `stacklevel` follows Python logging semantics: stacklevel=2 points to the
    # caller of this wrapper method.
    try:
        frame = inspect.currentframe()
        caller_frame = frame
        steps = max(int(stacklevel), 0)
        for _ in range(steps):
            caller_frame = caller_frame.f_back if caller_frame else None
        caller_info = {
            "caller_function": caller_frame.f_code.co_name
            if caller_frame
            else "unknown",
            "caller_module": os.path.basename(caller_frame.f_code.co_filename)
            if caller_frame
            else "unknown",
            "caller_file": caller_frame.f_code.co_filename
            if caller_frame
            else "unknown",
            "caller_line": caller_frame.f_lineno if caller_frame else 0,
        }
    except Exception as e:
        FedbiomedError(
            "Failed to capture caller information for security log entry. "
            "Caller info will be set to 'unknown'. "
            "Caller parser error: " + str(e)
        )
        caller_info = {
            "caller_function": "unknown",
            "caller_module": "unknown",
            "caller_file": "unknown",
            "caller_line": 0,
        }
        self._logger.warning(
            json.dumps(
                {
                    "timestamp": _utc_timestamp(),
                    "node_id": self._security_defaults.get("component_id"),
                    "caller_parsing": "failure",
                    "caller_parser_error": str(e),
                },
                default=str,
                separators=(",", ":"),
            ),
            extra={"is_security": True},
        )

    entry = {
        "timestamp": _utc_timestamp(),
        "component_id": self._security_defaults.get("component_id"),
        "component_name": self._security_defaults.get("component_name"),
        "researcher_id": rid,
        "operation": op,
        "status": st,
        "fedbiomed_version": self._security_defaults.get("fedbiomed_version"),
        **caller_info,
    }

    # Merge extra fields: context first, then explicit fields override
    merged = {}
    merged.update(ctx)
    merged.update(fields)

    # Avoid duplicating required keys in payload
    for k in ("operation", "status", "researcher_id"):
        merged.pop(k, None)

    entry.update(merged)

    # Write only to security file (is_security=True filters out other handlers)
    self._logger.info(
        json.dumps(entry, default=str, separators=(",", ":")),
        extra={"is_security": True},
        stacklevel=stacklevel,
    )

def setLevel(self, level: Any, htype: Any = None):
    """Overrides the setLevel method, to deal with level given as a string and to change le level of
    one or all known handlers

    This also change the default level for all future handlers.

    !!! info "Remark"

        Level should not be lower than CRITICAL (meaning CRITICAL errors are always displayed)

        Example:
        ```python
        setLevel( logging.DEBUG, 'FILE')
        ```

    Args:
        level : level to modify, can be a string or a logging.* level (mandatory)
        htype : if provided (non-mandatory), change the level of the given handler. if not  provided (or None),
            change the level of all known handlers
    """

    level = self._internal_level_translator(level)

    if htype is None:
        # store this level (for future handler adding)
        self._logger.setLevel(level)

        for h in self._handlers:
            self._handlers[h].setLevel(level)
            self._set_handler_formatter(h)
        return

    if htype in self._handlers:
        self._handlers[htype].setLevel(level)
        self._set_handler_formatter(htype)
        return

    # htype provided but no handler for this type exists
    self._logger.warning(htype + " handler not initialized yet")

def setPrefix(self, prefix: str = "") -> None:
    """Sets a log prefix for all handlers.

    Args:
        prefix: Prefix to add to all log messages
    """
    for h in self._handlers:
        self._handler_prefix[h] = prefix
        self._set_handler_formatter(h)

def set_security_logs(self, root_path: str) -> None:
    security_log_dir = os.path.join(root_path, "log")
    os.makedirs(security_log_dir, exist_ok=True)
    security_log_path = os.path.join(security_log_dir, "security_audit.log")
    self.add_security_file_handler(filename=security_log_path)

def warning(self, msg, *args, broadcast=False, researcher_id=None, **kwargs):
    """Same as info message"""
    # Merge extra dicts properly to support is_security flag
    extra = kwargs.pop("extra", {})
    extra.update({"researcher_id": researcher_id, "broadcast": broadcast})

    stacklevel = kwargs.pop("stacklevel", 2)
    self._logger.warning(
        msg,
        *args,
        **kwargs,
        extra=extra,
        stacklevel=stacklevel,
    )