"""Test Logging tracking utilities"""
import collections
from pathlib import Path
import tempfile
import logging
import atexit
from pyarrow import fs

from mldock.platform_helpers.mldock.tracking.logging import ExperimentTracker


class TestExperimentTracker:
    """Test Experiment Tracker functionality"""

    @staticmethod
    def __setup_tracker(experiment_logger, fs_base_path):
        """setup tracker"""
        experiment_tracker = ExperimentTracker(
            experiment_logger=experiment_logger,
            base_path="test_tracker",
            experiment_name="test_experiment",
            fs_base_path=fs_base_path,
            file_system=fs.LocalFileSystem(),
        )
        # temporarily unregister write artifacts from atexit callbacks to test it
        atexit.unregister(experiment_tracker.write_artifacts)
        return experiment_tracker

    def test_logs_written_to_fs_base_path_success(self, caplog):
        """
        Test logs are created and copied to fs_base_path on write_artifacts
        """

        tmp_dir = tempfile.TemporaryDirectory()
        test_dir = Path(tmp_dir.name, "tracker_output").as_posix()
        logger = logging.getLogger(tmp_dir.name)
        logger.setLevel(logging.DEBUG)

        experiment_tracker = self.__setup_tracker(logger, test_dir)
        experiment_tracker.log(r"first log")

        experiment_tracker.write_artifacts()

        files = [path.name for path in Path(test_dir).glob("**/*") if path.is_file()]

        # cleanup
        tmp_dir.cleanup()

        assert collections.Counter(files) == collections.Counter(
            ["logs.txt", "manifest.json"]
        ), "Failure: Did not create log.txt and manifest.json"

    def test_logs_contain_info_debug_error_success(self, caplog):
        """
        Test logs are created and copied to fs_base_path on write_artifacts
        """

        logger = logging.getLogger("tracker_output")
        logger.setLevel(logging.DEBUG)

        experiment_tracker = self.__setup_tracker(logger, None)
        experiment_tracker.log(r"first log")
        experiment_tracker.log(r"another debug log", level=logging.DEBUG)
        experiment_tracker.log(r"this is a log")
        experiment_tracker.log(r"an error log", level=logging.ERROR)

        for record in caplog.records:
            assert record.levelname in ["INFO", "DEBUG", "ERROR"]

    def test_logs_contain_no_debug_success(self, caplog):
        """
        Test logs are created and copied to fs_base_path on write_artifacts
        """

        logger = logging.getLogger("tracker_output")
        logger.setLevel(logging.INFO)

        experiment_tracker = self.__setup_tracker(logger, None)
        experiment_tracker.log(r"first log")
        experiment_tracker.log(r"another debug log", level=logging.DEBUG)
        experiment_tracker.log(r"this is a log")
        experiment_tracker.log(r"an error log", level=logging.ERROR)

        for record in caplog.records:
            # these should still be logged
            assert record.levelname in ["INFO", "ERROR"]
        for record in caplog.records:
            # these should not be logged
            assert record.levelname not in ["DEBUG"]

    def test_logs_metrics_successful(self, caplog):
        """
        Test logs are created and copied to fs_base_path on write_artifacts
        """

        logger = logging.getLogger("tracker_output")
        logger.setLevel(logging.DEBUG)

        experiment_tracker = self.__setup_tracker(logger, None)
        experiment_tracker.log(r"first log")
        experiment_tracker.log_param("estimators", 5)
        experiment_tracker.log_metric("mae", 0.5)
        experiment_tracker.log(r"another debug log", level=logging.DEBUG)
        experiment_tracker.log_metric("mse", 0.3)
        experiment_tracker.log(r"this is a log")
        experiment_tracker.log_metrics({"accuracy": 0.9, "roc": 0.9})
        experiment_tracker.log(r"an error log", level=logging.ERROR)

        metric_logs = []
        for record in caplog.records:
            if "metric:" in record.message:
                metric_logs.append(record.message)

        assert len(metric_logs) == 4, "Failure. Metric logs did not match expected"

    def test_logs_params_successful(self, caplog):
        """
        Test logs are created and copied to fs_base_path on write_artifacts
        """

        logger = logging.getLogger("tracker_output")
        logger.setLevel(logging.DEBUG)

        experiment_tracker = self.__setup_tracker(logger, None)
        experiment_tracker.log(r"first log")
        experiment_tracker.log_metric("mae", 0.5)
        experiment_tracker.log(r"another debug log", level=logging.DEBUG)
        experiment_tracker.log(r"this is a log")
        experiment_tracker.log(r"an error log", level=logging.ERROR)
        experiment_tracker.log_param("estimators", 5)
        experiment_tracker.log_params({"factors": 10, "lr": 0.01})

        param_logs = []
        for record in caplog.records:
            if "param:" in record.message:
                param_logs.append(record.message)

        assert len(param_logs) == 3, "Failure. Param logs did not match expected"
