# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates

#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at

#      http://www.apache.org/licenses/LICENSE-2.0

#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import logging
import os
from pathlib import Path
from typing import Union


def setup_logger(
    log_dir: Path,
    log_file_name: str,
    level: Union[int, str] = logging.INFO,
    log_to_console: bool = True,
    propagate: bool = True,
) -> logging.Logger:
    if propagate and logging.root.handlers:
        return get_propagate_logger(log_dir, log_file_name, level)

    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    log_path = os.path.join(log_dir, log_file_name)
    handlers = [logging.FileHandler(log_path, encoding="utf-8")]
    if log_to_console:
        handlers.append(logging.StreamHandler())

    logging.basicConfig(
        level=level,
        format="[%(asctime)s] [%(filename)s] [%(levelname)s]: %(message)s",
        handlers=handlers,
    )

    logger = logging.getLogger(log_path)

    return logger


def get_propagate_logger(
    log_dir: Path,
    log_file: str,
    level: Union[int, str] = logging.INFO,
) -> logging.Logger:
    if not log_dir.exists():
        log_dir.mkdir(parents=True, exist_ok=True)

    log_path = os.path.join(log_dir, log_file)
    propagate_logger = logging.getLogger(log_path)
    propagate_logger.setLevel(level)

    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setLevel(logging.INFO)
    formatter = logging.Formatter(
        "[%(asctime)s] [%(filename)s] [%(levelname)s]: %(message)s"
    )
    file_handler.setFormatter(formatter)

    propagate_logger.addHandler(file_handler)

    propagate_logger.propagate = True

    return propagate_logger


def get_non_propagate_logger(
    log_dir: Path,
    log_file: str,
    level: Union[int, str] = logging.INFO,
    log_to_console: bool = True,
) -> logging.Logger:
    if not log_dir.exists():
        log_dir.mkdir(parents=True, exist_ok=True)

    log_path = os.path.join(log_dir, log_file)
    non_propagate_logger = logging.getLogger(log_path)
    if non_propagate_logger.handlers:
        return non_propagate_logger

    non_propagate_logger.setLevel(level)

    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setLevel(logging.INFO)
    formatter = logging.Formatter(
        "[%(asctime)s] [%(filename)s] [%(levelname)s]: %(message)s"
    )
    file_handler.setFormatter(formatter)

    non_propagate_logger.addHandler(file_handler)
    if log_to_console:
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        non_propagate_logger.addHandler(console_handler)

    non_propagate_logger.propagate = False

    return non_propagate_logger
