from typing import IO, Type

import asyncio
import io
import logging
import sys
import time
from copy import deepcopy
from datetime import datetime, timezone, timedelta

import pytest

from mode.utils.logging import (
    LOG_RECORD_BUILTINS,
    CompositeLogger,
    DefaultFormatter,
    ExtensionFormatter,
    FileLogProxy,
    LogMessage,
    Logwrapped,
    _FlightRecorderProxy,
    _formatter_registry,
    _setup_logging,
    current_flight_recorder,
    flight_recorder,
    formatter,
    get_logger,
    level_name,
    level_number,
    on_timeout,
    print_task_name,
    redirect_stdouts,
    setup_logging,
)
from mode.utils.mocks import ANY, AsyncMock, Mock, call, patch


def log_called_with(logger, *args, stacklevel, **kwargs):
    logger.log.assert_called_once_with(*args, stacklevel=stacklevel, **kwargs)


def formatter_called_with(formatter, *args, stacklevel, **kwargs):
    formatter.assert_called_once_with(*args, stacklevel=stacklevel, **kwargs)


def test_log_record_builtins():
    rv = logging.makeLogRecord({})
    assert LOG_RECORD_BUILTINS - {"message", "asctime", "extra"} == set(
        rv.__dict__.keys()
    )


class test_CompositeLogger:
    @pytest.fixture()
    def logger(self):
        return Mock(name="logger")

    @pytest.fixture()
    def formatter(self):
        return Mock(name="formatter")

    @pytest.fixture()
    def log(self, *, logger, formatter):
        return CompositeLogger(logger=logger, formatter=formatter)

    def test_log(self, *, log, logger, formatter):
        log.log(logging.INFO, "msg", 1, kw=2)
        log_called_with(
            logger,
            logging.INFO,
            formatter.return_value,
            1,
            stacklevel=2,
            kw=2,
        )
        formatter_called_with(
            formatter, logging.INFO, "msg", 1, kw=2, stacklevel=2
        )

    def test_log__no_formatter(self, *, log, logger):
        log.formatter = None
        log.log(logging.INFO, "msg", 1, kw=2)
        log_called_with(
            logger,
            logging.INFO,
            "msg",
            1,
            kw=2,
            stacklevel=2,
        )

    @pytest.mark.parametrize(
        "method,severity,extra",
        [
            ("debug", logging.DEBUG, {}),
            ("info", logging.INFO, {}),
            ("warn", logging.WARN, {}),
            ("warning", logging.WARN, {}),
            ("error", logging.ERROR, {}),
            ("crit", logging.CRITICAL, {}),
            ("critical", logging.CRITICAL, {}),
            ("exception", logging.ERROR, {"exc_info": 1}),
        ],
    )
    def test_severity_mixin(self, method, severity, extra, *, log, logger):
        log.formatter = None
        getattr(log, method)("msg", "arg1", kw1=3, kw2=5)
        log_called_with(
            logger,
            severity,
            "msg",
            "arg1",
            kw1=3,
            kw2=5,
            stacklevel=3,
            **extra,
        )

    def test_dev__enabled(self, log):
        log.log = Mock()
        with patch("mode.utils.logging.DEVLOG", True):
            log.dev("msg", 1, k=2)
            log_called_with(log, logging.INFO, "msg", 1, k=2, stacklevel=3)

    def test_dev__disabled(self, log):
        log.info = Mock()
        with patch("mode.utils.logging.DEVLOG", False):
            log.dev("msg", 1, k=2)
            log.info.assert_not_called()


def test_formatter():
    f = Mock()
    formatter(f)
    try:
        assert f in _formatter_registry
    finally:
        _formatter_registry.remove(f)


@pytest.mark.parametrize("MyFormatter", (DefaultFormatter, ExtensionFormatter))
class test_Formatter:
    def test_format(self, MyFormatter: Type[logging.Formatter]):
        extra = {"moo": 30, "baz": 1111, "foo": "hello"}
        record = logging.getLogger("__test_Formatter__").makeRecord(
            "name",
            logging.INFO,
            "path",
            303,
            r"msg %s",
            (1,),
            None,
            extra=extra,
        )
        _ = MyFormatter().format(record)
        keys = set(s.split("=")[0] for s in record.extra.split(", "))
        assert keys == set(extra.keys())

    def test_formatTime(self, MyFormatter):
        record = logging.LogRecord(
            "name",
            logging.INFO,
            "path",
            303,
            "msg",
            {"foo": 1, "extra": {"data": {"moo": 30, "baz": [1, 2]}}},
            exc_info=None,
        )
        s = MyFormatter().formatTime(record)
        iso = datetime.fromtimestamp(
            record.created, tz=timezone(timedelta(seconds=-time.timezone))
        ).isoformat(timespec="milliseconds")
        assert s == iso


@pytest.mark.parametrize(
    "input,expected",
    [
        ("DEBUG", logging.DEBUG),
        ("INFO", logging.INFO),
        ("WARNING", logging.WARNING),
        ("WARNING", logging.WARNING),
        ("ERROR", logging.ERROR),
        ("CRITICAL", logging.CRITICAL),
        (logging.ERROR, logging.ERROR),
    ],
)
def test_level_number(input, expected):
    assert level_number(input) == expected


@pytest.mark.parametrize(
    "input",
    (None, 10.0, 0.0),
)
def test_level_number__type(input):
    with pytest.raises(TypeError):
        level_number(input)


@pytest.mark.parametrize(
    "input,expected",
    [
        (logging.DEBUG, "DEBUG"),
        (logging.INFO, "INFO"),
        (logging.WARN, "WARNING"),
        (logging.WARNING, "WARNING"),
        (logging.ERROR, "ERROR"),
        (logging.CRITICAL, "CRITICAL"),
        ("INFO", "INFO"),
    ],
)
def test_level_name(input, expected):
    assert level_name(input) == expected


@pytest.mark.parametrize(
    "input",
    (None, 10.0, 0.0),
)
def test_level_name__type(input):
    with pytest.raises(TypeError):
        level_name(input)


class test_setup_logging:
    def test_default(self):
        with patch("mode.utils.logging._setup_logging") as _sl:
            setup_logging(log_level="INFO", log_file=None)

            _sl.assert_called_once_with(
                level=logging.INFO,
                filename=None,
                stream=sys.stdout,
                log_handlers=None,
                logging_config=None,
            )

    def test_logfile(self):
        with patch("mode.utils.logging._setup_logging") as _sl:
            setup_logging(log_level="INFO", log_file="foo.txt")

            _sl.assert_called_once_with(
                level=logging.INFO,
                filename="foo.txt",
                stream=None,
                log_handlers=None,
                logging_config=None,
            )

    def test_io(self):
        logfile = Mock(spec=IO)
        with patch("mode.utils.logging._setup_logging") as _sl:
            setup_logging(log_level="INFO", log_file=logfile)

            _sl.assert_called_once_with(
                level=logging.INFO,
                filename=None,
                stream=logfile,
                log_handlers=None,
                logging_config=None,
            )

    def test_io_no_tty(self):
        logfile = Mock(spec=IO)
        logfile.isatty.side_effect = AttributeError()
        with patch("mode.utils.logging._setup_logging") as _sl:
            setup_logging(log_level="INFO", log_file=logfile)

            _sl.assert_called_once_with(
                level=logging.INFO,
                filename=None,
                stream=logfile,
                log_handlers=None,
                logging_config=None,
            )


class test__setup_logging:
    def setup_method(self, method):
        self.extension_formatter_patch = patch(
            "mode.utils.logging.ExtensionFormatter"
        )
        self.extension_formatter = self.extension_formatter_patch.start()
        self.colorlog_patch = patch("mode.utils.logging.colorlog")
        self.colorlog = self.colorlog_patch.start()
        self.logging_patch = patch("mode.utils.logging.logging")
        self.logging = self.logging_patch.start()

    def teardown_method(self):
        self.extension_formatter_patch.stop()
        self.colorlog_patch.stop()
        self.logging_patch.stop()

    def test_get_logger(self):
        assert get_logger(__name__)
        assert get_logger(__name__).handlers == get_logger(__name__).handlers

    def test_setup_logging_helper_both_filename_and_stream(self):
        with pytest.raises(AssertionError):
            _setup_logging(filename="TEMP", stream=Mock())

    def test_setup_logging_helper_with_filename(self):
        _setup_logging(filename="TEMP")
        self.logging.config.dictConfig.assert_called_once_with(ANY)

    def test_setup_logging_helper_with_stream_no_handlers(self):
        _setup_logging(stream=Mock())
        self.logging.config.dictConfig.assert_called_once_with(ANY)

    def test_setup_logging_helper_with_stream(self):
        mock_handler = Mock()
        _setup_logging(
            filename=None,
            stream=Mock(),
            log_handlers=[mock_handler],
        )
        self.logging.config.dictConfig.assert_called_once_with(ANY)
        self.logging.root.addHandler.assert_called_once_with(mock_handler)

    def test_setup_logging_helper_with_merge_config(self):
        _setup_logging(
            filename=None,
            stream=Mock(),
            logging_config={"merge": True, "foo": 1},
        )
        self.logging.config.dictConfig.assert_called_once_with(ANY)

    def test_setup_logging_helper_no_merge_config(self):
        _setup_logging(
            logging_config={"merge": False, "foo": 1},
        )
        self.logging.config.dictConfig.assert_called_once_with(ANY)


class test_Logwrapped:
    @pytest.fixture()
    def obj(self):
        return Mock(name="obj")

    @pytest.fixture()
    def logger(self):
        return Mock(name="logger")

    @pytest.fixture()
    def wrapped(self, *, obj, logger):
        return Logwrapped(obj, logger, severity="INFO", ident="ident")

    def test_constructor(self, wrapped):
        assert wrapped.severity == logging.INFO

    def test_wrapper(self, wrapped, obj):
        obj.calculate.__name__ = "calculate"
        wrapped.calculate(1, 2, kw=1)
        obj.calculate.assert_called_once_with(1, 2, kw=1)

    def test_wrapper__no_ident(self, wrapped, obj):
        wrapped.ident = None
        obj.calculate.__name__ = "calculate"
        wrapped.calculate(1, 2, kw=1)
        obj.calculate.assert_called_once_with(1, 2, kw=1)

    def test_wrapper__no_args(self, wrapped, obj):
        obj.calculate.__name__ = "calculate"
        wrapped.calculate()
        obj.calculate.assert_called_once_with()

    def test_wrapper__only_kwargs(self, wrapped, obj):
        obj.calculate.__name__ = "calculate"
        wrapped.calculate(kw=3)
        obj.calculate.assert_called_once_with(kw=3)

    def test_wrapper__ignored(self, wrapped, logger, obj):
        logger.log.assert_not_called()
        obj.__enter__ = Mock()
        obj.__exit__ = Mock()
        assert wrapped.__enter__ is obj.__enter__
        assert wrapped.__exit__ is obj.__exit__

    def test_repr(self, wrapped, obj):
        assert repr(wrapped) == repr(obj)

    def test_dir(self, wrapped, obj):
        assert dir(wrapped) == dir(obj)


def test_print_task_name():
    out = io.StringIO()
    task = Mock()
    task.__wrapped__ = Mock()
    task._coro.__name__ = "foo"
    print_task_name(task, file=out)
    assert out.getvalue()

    task._coro.__name__ = None
    task.__wrapped__ = None
    print_task_name(task, file=out)
    assert out.getvalue()


class test_flight_recorder:
    @pytest.fixture()
    def logger(self):
        return Mock(name="logger")

    @pytest.fixture()
    def bb(self, *, logger):
        return flight_recorder(logger, timeout=30.0)

    def test_wrap_debug(self, bb):
        obj = Mock()
        bb.wrap = Mock()
        bb.wrap_debug(obj)
        bb.wrap.assert_called_once_with(logging.DEBUG, obj)

    def test_wrap_info(self, bb):
        obj = Mock()
        bb.wrap = Mock()
        bb.wrap_info(obj)
        bb.wrap.assert_called_once_with(logging.INFO, obj)

    def test_wrap_warn(self, bb):
        obj = Mock()
        bb.wrap = Mock()
        bb.wrap_warn(obj)
        bb.wrap.assert_called_once_with(logging.WARN, obj)

    def test_wrap_error(self, bb):
        obj = Mock()
        bb.wrap = Mock()
        bb.wrap_error(obj)
        bb.wrap.assert_called_once_with(logging.ERROR, obj)

    def test_wrap(self, bb):
        obj = Mock()
        with patch("mode.utils.logging.Logwrapped") as Logwrapped:
            ret = bb.wrap(logging.ERROR, obj)
            assert ret is Logwrapped.return_value
            Logwrapped.assert_called_once_with(
                logger=bb,
                severity=logging.ERROR,
                obj=obj,
            )

    def test_activate(self, bb):
        bb._fut = None
        bb._waiting = Mock()
        with patch("mode.utils.logging.current_task") as current_task:
            with patch("asyncio.ensure_future") as ensure_future:
                bb.activate()
                assert bb.started_at_date
                assert bb.enabled_by is current_task.return_value
                ensure_future.assert_called_once_with(
                    bb._waiting.return_value,
                    loop=bb.loop,
                )
                assert bb._fut is ensure_future.return_value

    def test_activate__already_activated(self, bb):
        bb._fut = Mock()
        with pytest.raises(RuntimeError):
            bb.activate()

    def test_cancel(self, bb):
        bb._fut = None
        bb._logs = [1, 2, 3]
        bb.cancel()
        assert bb._logs == []
        assert bb._fut is None

        fut = bb._fut = Mock()
        bb.cancel()
        assert bb._fut is None
        fut.cancel.assert_called_once_with()  # type: ignore

    def test_log__active(self, bb, logger):
        bb._fut = Mock()
        bb._buffer_log = Mock()
        bb.log(logging.DEBUG, "msg %r %(foo)s", 1, foo="bar")
        bb._buffer_log.assert_called_once_with(
            logging.DEBUG,
            "msg %r %(foo)s",
            (1,),
            {"foo": "bar"},
        )

    def test_log__inactive(self, bb, logger):
        bb._fut = None
        bb._buffer_log = Mock()
        bb.log(logging.DEBUG, "msg %r %(foo)s", 1, foo="bar")
        log_called_with(
            logger,
            logging.DEBUG,
            "msg %r %(foo)s",
            1,
            foo="bar",
            stacklevel=2,
        )

    def test__buffer_log(self, bb):
        with patch("time.asctime") as asctime:
            bb._buffer_log(
                logging.ERROR, "msg %r %(foo)s", (1,), {"foo": "bar"}
            )
            assert bb._logs[-1] == LogMessage(
                logging.ERROR,
                "msg %r %(foo)s",
                asctime(),
                (1,),
                {"foo": "bar"},
            )

    @pytest.mark.asyncio
    async def test__waiting__cancelled(self, bb):
        assert not bb._logs
        bb._buffer_log(logging.ERROR, "msg %r %(foo)s", (1,), {"foo": "bar"})
        with patch("asyncio.sleep", AsyncMock()) as sleep:
            sleep.coro.side_effect = asyncio.CancelledError()
            await bb._waiting()
            sleep.assert_called_once_with(bb.timeout)
            assert bb._logs

    @pytest.mark.asyncio
    async def test__waiting__has_logs(self, bb):
        assert not bb._logs
        bb._buffer_log(logging.ERROR, "msg %r %(foo)s", (1,), {"foo": "bar"})
        assert bb._logs
        with patch("asyncio.sleep", AsyncMock()):
            await bb._waiting()

    @pytest.mark.asyncio
    async def test__waiting__no_logs(self, bb):
        assert not bb._logs
        with patch("asyncio.sleep", AsyncMock()):
            await bb._waiting()

    @pytest.mark.asyncio
    async def test__waiting__enabled_by(self, bb):
        assert not bb._logs
        bb.enabled_by = Mock()
        with patch("asyncio.sleep", AsyncMock()):
            with patch("mode.utils.logging.format_task_stack") as fts:
                await bb._waiting()
            fts.assert_called_once_with(bb.enabled_by)

    @pytest.mark.asyncio
    async def test__waiting__raises(self, bb):
        assert not bb._logs
        bb.logger.warning = Mock(side_effect=KeyError())
        with patch("asyncio.sleep", AsyncMock()):
            with pytest.raises(KeyError):
                await bb._waiting()

    def test_repr(self, bb):
        assert repr(bb)

    def test_context(self, bb):
        bb.activate = Mock()
        bb.cancel = Mock()
        with bb:
            bb.activate.assert_called_once_with()
        bb.cancel.assert_called_once_with()


class test_FileLogProxy:
    def test_constructor__defaults(self):
        logger = get_logger("foo")
        logger.level = None
        assert logger.level is None
        f = FileLogProxy(logger)
        assert f.severity == logging.WARN

    def test_constructor__severity_from_logger(self):
        logger = get_logger("foo")
        logger.level = logging.DEBUG
        f = FileLogProxy(logger)
        assert f.severity == logging.DEBUG

    def test_constructor__explicit_severity(self):
        logger = get_logger("foo")
        logger.level = logging.DEBUG
        f = FileLogProxy(logger, severity=logging.ERROR)
        assert f.severity == logging.ERROR

    def test__safewrap_handler(self):
        f = FileLogProxy(get_logger("foo"))
        handler = Mock()
        f._safewrap_handler(handler)

        with patch("traceback.print_exc") as print_exc:
            record = Mock()
            handler.handleError(record)
            print_exc.assert_called_once_with(None, sys.__stderr__)
            print_exc.side_effect = IOError()
            handler.handleError(record)

    def test_write(self):
        logger = Mock(handlers=[])
        f = FileLogProxy(logger)
        f._threadlocal.recurse_protection = True
        f.write("foo")
        logger.log.assert_not_called()

        f._threadlocal.recurse_protection = False
        f.write("")
        f.write("               ")
        f.close()
        f.write("msg")
        logger.log.assert_not_called()

        f._closed = False
        f.write(" msg ")
        logger.log.assert_called_once_with(f.severity, "msg")

        f.writelines(["foo", "bar"])
        logger.log.assert_has_calls(
            [
                call(f.severity, "msg"),
                call(f.severity, "foo"),
                call(f.severity, "bar"),
            ]
        )

    def test_flush(self):
        FileLogProxy(get_logger("foo")).flush()

    def test_isatty(self):
        assert not FileLogProxy(get_logger("foo")).isatty()


def test_redirect_stdouts():
    prev_stdout = sys.stdout
    prev_stderr = sys.stderr
    with redirect_stdouts(stdout=False, stderr=False):
        assert sys.stdout == prev_stdout
        assert sys.stderr == prev_stderr
    prev_stdout = sys.stdout  # pytest keep changing this
    prev_stderr = sys.stderr
    with redirect_stdouts(stdout=True, stderr=False):
        assert isinstance(sys.stdout, FileLogProxy)
        assert sys.stderr == prev_stderr
    prev_stdout = sys.stdout
    prev_stderr = sys.stderr
    with redirect_stdouts(stdout=False, stderr=True):
        assert sys.stdout == prev_stdout
        assert isinstance(sys.stderr, FileLogProxy)
    prev_stdout = sys.stdout
    prev_stderr = sys.stderr
    with redirect_stdouts(stdout=True, stderr=True):
        assert isinstance(sys.stdout, FileLogProxy)
        assert isinstance(sys.stderr, FileLogProxy)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "extra_context",
    [
        {},
        {"foo": "bar"},
    ],
)
async def test_on_timeout(extra_context):
    logger = Mock()
    assert isinstance(on_timeout, _FlightRecorderProxy)

    # Test no errors when there's no active flight recorder
    _assert_log_severities(on_timeout)

    with patch("time.asctime") as asctime:
        asctime.return_value = "TIME"
        # Test logging to active flight recorder (with nesting)
        with flight_recorder(logger, timeout=300) as fl1:
            fl1.extra_context.update(extra_context)
            assert current_flight_recorder() is fl1
            _assert_recorder_exercised(on_timeout, fl1)

            with flight_recorder(logger, timeout=30) as fl2:
                for k, v in fl1.extra_context.items():
                    assert fl2.extra_context[k] == v
                assert current_flight_recorder() is fl2
                _assert_recorder_exercised(on_timeout, fl2)
                _assert_recorder_flush_logs(logger, fl2)

            assert current_flight_recorder() is fl1
            _assert_recorder_flush_logs(logger, fl1)
            _assert_recorder_exercised(on_timeout, fl1)
            _assert_recorder_flush_logs(logger, fl1)


def _assert_log_severities(logger):
    logger.debug("DEBUG %d %(a)s", 1, a="A")
    logger.info("INFO %d %(b)s", 2, b="B")
    logger.warning("WARNING %d %(c)s", 3, c="C")
    logger.error("ERROR %d %(d)s", 4, d="D")
    logger.critical("CRITICAL %d %(e)s", 5, e="E")


def _log_kwargs(kwargs):
    kwargs.setdefault("stacklevel", 3)
    return kwargs


EXPECTED_LOG_MESSAGES = [
    LogMessage(
        logging.DEBUG, "DEBUG %d %(a)s", "TIME", (1,), _log_kwargs({"a": "A"})
    ),
    LogMessage(
        logging.INFO, "INFO %d %(b)s", "TIME", (2,), _log_kwargs({"b": "B"})
    ),
    LogMessage(
        logging.WARNING,
        "WARNING %d %(c)s",
        "TIME",
        (3,),
        _log_kwargs({"c": "C"}),
    ),
    LogMessage(
        logging.ERROR, "ERROR %d %(d)s", "TIME", (4,), _log_kwargs({"d": "D"})
    ),
    LogMessage(
        logging.CRITICAL,
        "CRITICAL %d %(e)s",
        "TIME",
        (5,),
        _log_kwargs({"e": "E"}),
    ),
]


def _assert_recorder_exercised(logger, fl):
    _assert_log_severities(logger)
    assert fl._logs == EXPECTED_LOG_MESSAGES


def _assert_recorder_flush_logs(logger, fl):
    fl.flush_logs(ident="IDENT")

    def _get_call(sev, msg, datestr, args, kwargs):
        kw = deepcopy(kwargs)
        if fl.extra_context:
            extra = kw.setdefault("extra", {})
            data = extra.setdefault("data", {})
            data.update(fl.extra_context)
        return call(sev, f"[%s] (%s) {msg}", "IDENT", datestr, *args, **kw)

    logger.log.assert_has_calls(
        _get_call(sev, msg, datestr, args, kwargs)
        for sev, msg, datestr, args, kwargs in EXPECTED_LOG_MESSAGES
    )
