#!/usr/bin/env python3

# SPDX-FileCopyrightText: 2026 Dennis Fink <me+coding@dennisfink.me>
#
# SPDX-License-Identifier: BSD-3-Clause


import json
import os
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Any, cast

import click_extra as click
from prometheus_client import write_to_textfile

from . import metrics, types

###############################################################################
# SCRIPT METADATA
#
# These variables describe the script and are used for the --version output
# and for informational messages.
###############################################################################

VERSION = "1.0.0"
DESCRIPTION = "Collect borgmatic repository metrics and export them for Prometheus."
DATE_OF_CREATION = "2021-06-12"
DATE_OF_REVISION = "2026-05-01"
AUTHOR = "Dennis Fink <me+coding@dennisfink.me>"
LICENSE = "BSD-3-Clause"


###############################################################################
# FUNCTIONS
#
# This section defines the helper functions used throughout the script. They
# handle formatted output, error reporting, version information, utilities
# functions and runtime behaviour such as displaying usage instructions.
###############################################################################


def emit(prefix: str, color: str, *message: str, enabled: bool = True) -> None:
    """Render a styled prefix and message to stdout.

    Internal helper used by error(), msg(), warn(), verbose(), and debug()
    to avoid duplicating the Click styling and echo logic. The first message
    argument is rendered in bold, subsequent arguments are appended unstyled.

    Args:
        prefix: The prefix string shown before the message (e.g. "==>", "==> ERROR:").
        color: A Click-compatible color name (e.g. "red", "green") applied to the prefix.
        *message: One or more message parts to display. At least one is required.
        enabled: If False, the function returns immediately without printing.
                 Defaults to True.

    Raises:
        TypeError: If no message arguments are provided.
    """
    if not message:
        raise TypeError("emit() missing 1 required positional argument: 'message'")

    if not enabled:
        return

    click.echo(
        " ".join(
            [
                click.style(prefix, fg=color, bold=True),
                click.style(message[0], bold=True),
                *message[1:],
            ]
        )
    )


def error(*message: str) -> None:
    """Print a formatted error message to stdout in red.

    The first argument is rendered in bold, subsequent arguments are appended
    unstyled. Always prints regardless of quiet or verbose flags.

    Args:
        *message: One or more message parts to display.
    """
    emit("==> ERROR:", "red", *message)


@click.pass_obj
def msg(obj: dict[str, bool], *message: str) -> None:
    """Print a formatted informational message to stdout in green.

    Suppressed when the quiet flag is set. The first argument is rendered
    in bold, subsequent arguments are appended unstyled.

    Args:
        *message: One or more message parts to display.
    """
    emit("==>", "green", *message, enabled=not obj.get("quiet", False))


@click.pass_obj
def warn(obj: dict[str, bool], *message: str) -> None:
    """Print a formatted warning message to stdout in yellow.

    Suppressed when the quiet flag is set. The first argument is rendered
    in bold, subsequent arguments are appended unstyled.

    Args:
        *message: One or more message parts to display.
    """
    emit("==>", "yellow", *message, enabled=not obj.get("quiet", False))


@click.pass_obj
def verbose(obj: dict[str, bool], *message: str) -> None:
    """Print a formatted verbose message to stdout in blue.

    Only prints when the verbose flag is set and the quiet flag is not.
    The first argument is rendered in bold, subsequent arguments are
    appended unstyled.

    Args:
        *message: One or more message parts to display.
    """
    emit(
        "==>",
        "blue",
        *message,
        enabled=obj.get("verbose", False) and not obj.get("quiet", False),
    )


@click.pass_obj
def debug(obj: dict[str, bool], *message: str) -> None:
    """Print a formatted debug message to stdout in magenta.

    Only prints when the DEBUG environment variable is set to a truthy
    value (1, true, yes). The first argument is rendered in bold,
    subsequent arguments are appended unstyled.

    Args:
        *message: One or more message parts to display.
    """
    emit("==>", "magenta", *message, enabled=obj.get("debug", False))


def print_version(
    ctx: click.Context,
    param: click.Parameter | None,
    value: bool,
) -> None:
    """Print script metadata (name, version, description, author, dates) and exit.

    Args:
        ctx: The Click context associated with the currently running command.
    """
    if not value or ctx.resilient_parsing:
        return

    click.echo(click.style("Scriptname:", fg="red", bold=True) + f" {ctx.info_name}")
    click.echo(click.style("Version:", fg="green", bold=True) + f" {VERSION}")
    click.echo(click.style("Description:", fg="yellow", bold=True) + f" {DESCRIPTION}")
    click.echo(click.style("Author:", fg="blue", bold=True) + f" {AUTHOR}")
    click.echo(
        click.style("Date of creation:", fg="magenta", bold=True)
        + f" {DATE_OF_CREATION}"
    )
    click.echo(
        click.style("Date of revision:", fg="cyan", bold=True) + f" {DATE_OF_REVISION}"
    )
    click.echo(click.style("License:", fg="red", bold=True) + f" {LICENSE}")

    click.echo("""Copyright (c) 2026 Dennis Fink <me+coding@dennisfink.me>.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its contributors
may be used to endorse or promote products derived from this software without
specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.""")

    ctx.exit()


###############################################################################
# MAIN EXECUTION
###############################################################################


@click.pass_context
def run_command(
    ctx: click.Context, command: list[str]
) -> subprocess.CompletedProcess[str]:
    """Run a shell command and return the result.

    Logs the command at debug level before execution. Exits the application
    with the command's return code if it fails.

    Args:
        ctx: The Click context, injected by @click.pass_context.
        command: The command and its arguments as a list of strings.

    Returns:
        The completed process result with stdout and stderr captured.
    """
    debug("Running command", " ".join(command))
    try:
        return subprocess.run(
            command,
            capture_output=True,
            encoding="utf-8",
            check=True,
        )
    except subprocess.CalledProcessError as e:
        error(f"Command ({' '.join(command)}) failed with exit code {e.returncode}")
        if e.stderr:
            error("stderr:", e.stderr.strip())
        ctx.exit(e.returncode)


@click.pass_context
def load_json(ctx: click.Context, raw: str) -> Any:
    """Parse a JSON string and return the result.

    Exits the application with code 1 if the input is not valid JSON,
    printing the parse error via the error() helper.

    Args:
        ctx: The Click context, injected by @click.pass_context.
        raw: The raw JSON string to parse.

    Returns:
        The parsed JSON value.
    """
    try:
        return json.loads(raw)
    except json.JSONDecodeError as e:
        error("Failed to parse JSON:", str(e))
        ctx.exit(1)


def extract_archive_counts(
    repositories: list[types.BorgInfoReturnType] | list[types.BorgListReturnType],
) -> list[types.ArchiveCount]:
    """Extract archive counts and repository metadata from borgmatic info output.

    Args:
        repositories: List of repository objects as returned by borgmatic info
                      or borgmatic list.

    Returns:
        A list of ArchiveCount dicts, one per repository, containing the
        repository id, label, location, and total number of archives.
    """
    return [
        {
            "id": repo["repository"]["id"],
            "label": repo["repository"]["label"],
            "location": repo["repository"]["location"],
            "count": len(repo["archives"]),
        }
        for repo in repositories
    ]


def fetch_repositories(
    borgmatic_bin: Path, borgmatic_configs: tuple[Path, ...], all_archives: bool
) -> tuple[list[types.BorgInfoReturnType], list[types.ArchiveCount]]:
    """Fetch repository and archive data from borgmatic.

    When all_archives is False, runs `borgmatic info --archive latest` and
    `borgmatic list` in complementary intent: info provides per-repository
    stats and the latest archive metrics, while list provides the total archive
    count — since info with --archive latest only returns a single archive,
    making len() unreliable for counting.

    When all_archives is True, only `borgmatic info` is run. The full archive
    list is available directly from the info output, so no separate list call
    is needed and archive counts are derived from it.

    Args:
        borgmatic_bin: Path to the borgmatic executable.
        borgmatic_config: Path to the borgmatic configuration file.
        all_archives: If True, fetch and return metrics for all archives.
                      This is significantly slower for large repositories.

    Returns:
        A tuple of:
        - repositories: Raw borgmatic info output, one entry per repository.
        - archive_counts: Archive count and identifying metadata per repository,
                          always populated regardless of the all_archives flag.
    """
    verbose("Fetching repositories")
    borgmatic_cmd = [str(borgmatic_bin)]

    for config in borgmatic_configs:
        borgmatic_cmd.extend(["--config", str(config)])

    info_command = [*borgmatic_cmd, "info", "--json"]

    if all_archives:
        borgmatic_info = run_command(info_command)
        repositories = cast(
            list[types.BorgInfoReturnType],
            load_json(borgmatic_info.stdout),
        )
        debug(f"Fetched {len(repositories)} repositories")
        return repositories, extract_archive_counts(repositories)

    info_command += ["--archive", "latest"]
    borgmatic_info = run_command(info_command)
    repositories = cast(
        list[types.BorgInfoReturnType], load_json(borgmatic_info.stdout)
    )
    debug(f"Fetched {len(repositories)} repositories")

    borgmatic_list = run_command([*borgmatic_cmd, "list", "--json"])
    archives = cast(list[types.BorgListReturnType], load_json(borgmatic_list.stdout))
    return repositories, extract_archive_counts(archives)


def collect_repository_metrics(
    repository: types.BorgInfoReturnType, all_archives: bool
) -> None:
    """Collect and set all Prometheus metrics for a single repository.

    Sets cache stats, encryption info, last modified timestamp, and latest
    archive metrics. When all_archives is True, also collects metrics for
    every individual archive by calling _collect_archive_metrics.

    Args:
        repository: Repository info as returned by borgmatic.
        all_archives: If True, collect metrics for all archives, not just
                      the latest.
    """
    repository_metadata = repository["repository"]

    msg(
        "Processing repository:",
        repository_metadata["label"],
        f"({repository_metadata['id']})",
        repository_metadata["location"],
    )

    repository_labels: types.RepositoryLabels = {
        "repository_id": repository_metadata["id"],
        "repository_label": repository_metadata["label"],
        "location": repository_metadata["location"],
    }

    repository_last_modified = datetime.fromisoformat(
        repository_metadata["last_modified"]
    )
    metrics.LAST_MODIFIED.labels(**repository_labels).set(
        repository_last_modified.timestamp()
    )

    metrics.ENCRYPTION.labels(**repository_labels).info(
        {"mode": repository["encryption"]["mode"]}
    )

    repository_stats = repository["cache"]["stats"]
    for metric, key in (
        (metrics.TOTAL_CHUNKS, "total_chunks"),
        (metrics.TOTAL_COMPRESSED_SIZE, "total_csize"),
        (metrics.TOTAL_ORIGINAL_SIZE, "total_size"),
        (metrics.TOTAL_UNIQUE_CHUNKS, "total_unique_chunks"),
        (metrics.TOTAL_DEDUPLICATED_COMPRESSED_SIZE, "unique_csize"),
        (metrics.TOTAL_DEDUPLICATED_SIZE, "unique_size"),
    ):
        metric.labels(**repository_labels).set(repository_stats[key])

    if not repository["archives"]:
        warn("No archives found for repository:", repository_metadata["label"])
        return

    latest_archive = max(
        repository["archives"],
        key=lambda archive: archive["start"],
    )
    debug(f"Latest archive {latest_archive['name']}:", latest_archive["start"])

    metrics.LATEST_ARCHIVE_DURATION.labels(**repository_labels).set(
        latest_archive["duration"]
    )

    latest_archive_stats = latest_archive["stats"]
    for metric, key in (
        (metrics.LATEST_ARCHIVE_COMPRESSED_SIZE, "compressed_size"),
        (metrics.LATEST_ARCHIVE_DEDUPLICATED_SIZE, "deduplicated_size"),
        (metrics.LATEST_ARCHIVE_NUMBER_FILES, "nfiles"),
        (metrics.LATEST_ARCHIVE_ORIGINAL_SIZE, "original_size"),
    ):
        metric.labels(**repository_labels).set(latest_archive_stats[key])

    if all_archives:
        for archive in repository["archives"]:
            collect_archive_metrics(archive, repository_labels)


def collect_archive_metrics(
    archive: types.BorgArchive,
    repository_labels: types.RepositoryLabels,
) -> None:
    """Collect and set Prometheus metrics for a single archive.

    Args:
        archive: Archive info as returned by borgmatic.
        repository_labels: The parent repository's label dict, used to
                           populate the shared label dimensions on each metric.
    """
    msg("Processing archive:", archive["name"], f"({archive['id']})")

    archive_labels = {
        **repository_labels,
        "archive_id": archive["id"],
        "archive_name": archive["name"],
    }

    metrics.ARCHIVE_DURATION.labels(**archive_labels).set(archive["duration"])

    archive_stats = archive["stats"]
    for metric, key in (
        (metrics.ARCHIVE_COMPRESSED_SIZE, "compressed_size"),
        (metrics.ARCHIVE_DEDUPLICATED_SIZE, "deduplicated_size"),
        (metrics.ARCHIVE_NUMBER_FILES, "nfiles"),
        (metrics.ARCHIVE_ORIGINAL_SIZE, "original_size"),
    ):
        metric.labels(**archive_labels).set(archive_stats[key])


@click.command(
    context_settings={"help_option_names": ("-h", "--help", "-?")},
    formatter_settings=click.HelpExtraFormatter.settings(
        theme=click.HelpExtraTheme.light(),
    ),
    params=[],
)
@click.color_option(
    "--color/--no-color",
    help="Force/Disable colored output.",
)
@click.option(
    "--borgmatic-bin",
    type=click.Path(
        exists=True,
        readable=True,
        executable=True,
        dir_okay=False,
        allow_dash=False,
        path_type=Path,
    ),
    default="/usr/bin/borgmatic",
)
@click.option(
    "--borgmatic-config",
    "borgmatic_configs",
    multiple=True,
    type=click.Path(
        exists=True, readable=True, dir_okay=True, allow_dash=False, path_type=Path
    ),
    help=(
        "Configuration file or directory. "
        "Can be specified multiple times. "
        "If not specified, borgmatic's default configuration discovery is used."
    ),
)
@click.option(
    "--textfile-collector-dir",
    type=click.Path(
        exists=True,
        readable=True,
        writable=True,
        file_okay=False,
        allow_dash=False,
        path_type=Path,
    ),
    default="/var/lib/prometheus/node-exporter",
)
@click.option("--all-archives", "all_archives_flag", is_flag=True, default=False)
@click.option("-q", "--quiet", "quiet_flag", is_flag=True, default=False)
@click.option("-v", "--verbose", "verbose_flag", is_flag=True, default=False)
@click.option(
    "--version",
    callback=print_version,
    default=False,
    expose_value=False,
    help="Show the version and exit.",
    is_eager=True,
    is_flag=True,
)
@click.pass_context
def prometheus_borgmatic_exporter(
    ctx: click.Context,
    borgmatic_bin: Path,
    borgmatic_configs: tuple[Path, ...],
    textfile_collector_dir: Path,
    all_archives_flag: bool,
    quiet_flag: bool,
    verbose_flag: bool,
):
    """Collect borgmatic repository metrics and export them for Prometheus.

    Queries borgmatic for repository and archive statistics and writes them
    to a .prom file for consumption by the Prometheus node exporter textfile
    collector.

    By default only the latest archive is queried per repository. Use
    --all-archives to collect metrics for every archive, at the cost of
    a significantly longer runtime.

    Using --all-archives will create one time series per archive and may
    significantly increase Prometheus storage usage.
    """

    ctx.ensure_object(dict)
    ctx.obj["debug"] = os.environ.get("DEBUG", "").lower() in ("1", "true", "yes")
    ctx.obj["quiet"] = quiet_flag
    ctx.obj["verbose"] = verbose_flag

    debug("Effective configuration:")
    debug("  borgmatic_bin:", str(borgmatic_bin))
    debug(
        "  borgmatic_configs:",
        ", ".join(str(config) for config in borgmatic_configs)
        or "<borgmatic defaults>",
    )
    debug("  textfile_collector_dir:", str(textfile_collector_dir))
    debug("  all_archives:", str(all_archives_flag))
    debug("  verbose:", str(verbose_flag))
    debug("  quiet:", str(quiet_flag))

    repositories, archive_counts = fetch_repositories(
        borgmatic_bin, borgmatic_configs, all_archives_flag
    )

    for repository in archive_counts:
        metrics.ARCHIVES.labels(
            repository_id=repository["id"],
            repository_label=repository["label"],
            location=repository["location"],
        ).set(repository["count"])

    for repository in repositories:
        collect_repository_metrics(repository, all_archives_flag)

    metrics_file = textfile_collector_dir / "borgmatic.prom"
    write_to_textfile(str(metrics_file.absolute()), metrics.REGISTRY)
    msg("Metrics written to:", str(metrics_file.absolute()))


if __name__ == "__main__":
    prometheus_borgmatic_exporter()