# Copyright 2023 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""An implementation of the RemoteexecConfig proto interface."""

import datetime
import getpass
import json
import logging
import os
from pathlib import Path
import shlex
from typing import List

from chromite.lib import cros_build_lib
from chromite.lib import osutils
from chromite.utils import hostname_util


# Synced with "log_dir" and "proxy_log_dir" in ./sdk/reclient_cfgs/reproxy.cfg
_REMOTEEXEC_LOG_BASE_DIR = Path("/tmp")
_REMOTEEXEC_LOG_DIRNAME_PATTERN = "reclient-*"

_REMOTEEXEC_LOG_FILE_PATTERN = (
    "*.INFO.*",
    "reproxy_*.INFO",
    "*.rrpl",
)


class LogsArchiver:
    """Manages archiving remoteexec log files."""

    def __init__(self, build_log_dir: Path, dest_dir: Path) -> None:
        """Initializes the archiver.

        Args:
            build_log_dir: path to the build directory.
            dest_dir: path to the target directory to which logs are written.
        """
        self.remoteexec_log_dir_for_testing = None
        self._dest_base_dir = dest_dir
        self._build_log_dir = build_log_dir
        osutils.SafeMakedirs(self._dest_base_dir)

    def archive(self) -> List[str]:
        """Archives all log files.

        Returns:
            A list of compressed log file paths.
        """
        return self.archive_remoteexec_logs() + self.archive_ninja_log()

    def archive_remoteexec_logs(self) -> List[str]:
        """Archives remoteexec log files.

        Returns:
            A list of compressed log file paths.
        """

        log_dir_pattern = (
            self.remoteexec_log_dir_for_testing or _REMOTEEXEC_LOG_BASE_DIR
        )
        log_dirs = sorted(log_dir_pattern.glob(_REMOTEEXEC_LOG_DIRNAME_PATTERN))

        archived_log_files = []
        for log_dir in log_dirs:
            logging.info("Processing log dir: %s", log_dir)

            for pattern in _REMOTEEXEC_LOG_FILE_PATTERN:
                archived_log_files += self._archive_files(log_dir, pattern)

        return archived_log_files

    def _archive_files(self, directory: Path, pattern: str) -> List[Path]:
        """Archives files matched with pattern, with gzip'ing.

        Args:
            directory: directory of the files.
            pattern: matching path pattern (eg. "*.INFO").

        Returns:
            A list of compressed log file paths.
        """
        # Find files matched with the |pattern| in |directory|. Sort for
        # stabilization.
        paths = sorted(directory.glob(pattern))
        if not paths:
            logging.warning("No glob files matched with %s", pattern)

        result = []
        for path in paths:
            archived_filename = f"{path.name}.gz"
            log_label = f"{directory.name}/{archived_filename}"
            logging.info("Compressing %s", log_label)

            dest_filepath = (
                self._dest_base_dir / directory.name / archived_filename
            )
            osutils.SafeMakedirs(dest_filepath.parent)
            cros_build_lib.CompressFile(path, dest_filepath)

            osutils.SafeUnlink(path)

            result.append(log_label)

        return result

    def archive_ninja_log(self) -> List[str]:
        """Archives ninja_log file and its related metadata.

        This archives the ninja_log file generated by ninja to build Chrome.
        Also, it appends some related metadata at the end of the file following
        '# end of ninja log' marker.

        Returns:
            The list of the archived files.
        """
        ninja_log_path = os.path.join(self._build_log_dir, "ninja_log")
        if not os.path.exists(ninja_log_path):
            logging.warning("ninja_log is not found: %s", ninja_log_path)
            return []
        ninja_log_content = osutils.ReadFile(ninja_log_path)

        try:
            st = os.stat(ninja_log_path)
            ninja_log_mtime = datetime.datetime.fromtimestamp(st.st_mtime)
        except OSError:
            logging.exception("Failed to get timestamp: %s", ninja_log_path)
            return []

        ninja_log_info = self._build_ninja_info()

        # Append metadata at the end of the log content.
        ninja_log_content += "# end of ninja log\n" + json.dumps(ninja_log_info)

        # Aligned with goma_utils in chromium bot.
        pid = os.getpid()

        ninja_log_path = self._build_log_dir / (
            "ninja_log.%s.%s.%s.%d"
            % (
                getpass.getuser(),
                hostname_util.get_host_name(),
                ninja_log_mtime.strftime("%Y%m%d-%H%M%S"),
                pid,
            )
        )
        osutils.WriteFile(ninja_log_path, ninja_log_content)

        archived_filename = os.path.basename(ninja_log_path) + ".gz"
        archived_path = self._dest_base_dir / archived_filename
        cros_build_lib.CompressFile(ninja_log_path, archived_path)

        return [archived_filename]

    def _build_ninja_info(self):
        """Reads metadata for the ninja run.

        Each metadata should be written into a dedicated file in the log
        directory. Read the info, and build the dict containing metadata.

        Returns:
            A dict of the metadata.
        """

        info = {"platform": "chromeos"}

        command_path = self._build_log_dir / "ninja_command"
        if os.path.exists(command_path):
            info["cmdline"] = shlex.split(
                osutils.ReadFile(command_path).strip()
            )

        cwd_path = self._build_log_dir / "ninja_cwd"
        if os.path.exists(cwd_path):
            info["cwd"] = osutils.ReadFile(cwd_path).strip()

        exit_path = self._build_log_dir / "ninja_exit"
        if os.path.exists(exit_path):
            info["exit"] = int(osutils.ReadFile(exit_path).strip())

        env_path = self._build_log_dir / "ninja_env"
        if os.path.exists(env_path):
            # env is null-byte separated, and has a trailing null byte.
            content = osutils.ReadFile(env_path).rstrip("\0")
            info["env"] = dict(
                line.split("=", 1) for line in content.split("\0")
            )

        return info