#!/usr/bin/env python3

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

import json
import os
import shlex
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
from .cli import FlexibleColorOption, debug, error, msg, verbose, warn

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


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

    :param ctx: Click context injected by :func:`click.pass_context`.
    """
    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()


@click.pass_context
def run_command(
    ctx: click.Context, command: list[str], accepted_exit_codes: tuple[int, ...] = (0,)
) -> subprocess.CompletedProcess[str]:
    """Run a 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.

    :param ctx: Click context injected by :func:`click.pass_context`.
    :param command: Command and its arguments.
    :param accepted_exit_codes: Exit codes considered successful.
    :return: Completed process with stdout and stderr captured.
    """
    command_string = shlex.join(command)
    debug("Running command", command_string, err=True)
    result = subprocess.run(command, capture_output=True, encoding="utf-8", check=False)
    if result.returncode not in accepted_exit_codes:
        error(
            f"Command ({command_string}) failed with exit code:", str(result.returncode)
        )
        if result.stderr:
            error("stderr:", result.stderr.strip())
        ctx.exit(result.returncode)
    return result


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

    Exits the application with status code 1 if the input is not valid JSON and
    reports the parsing error using :func:`error`.

    :param ctx: Click context injected by :func:`click.pass_context`.
    :param raw: Raw JSON string to parse.
    :return: 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 output.

    :param repositories: Repository objects returned by ``borgmatic info`` or
        ``borgmatic list``.
    :return: Archive counts and identifying metadata for each repository.
    """
    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``. The former provides per-repository statistics and
    metrics for the latest archive, while the latter provides the complete
    archive list needed to determine archive counts.

    When ``all_archives`` is true, only ``borgmatic info`` is run. Its output
    already contains the complete archive list, so archive counts can be
    derived directly from it.

    :param borgmatic_bin: Path to the borgmatic executable.
    :param borgmatic_configs: Paths to borgmatic configuration files or
        directories.
    :param all_archives: Whether to fetch metrics for all archives instead of
        only the latest archive.
    :return: Raw ``borgmatic info`` output and archive counts with identifying
        metadata for each repository.
    """
    verbose("Fetching repositories", err=True)
    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", err=True)
        return repositories, extract_archive_counts(repositories)

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

    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 Prometheus metrics for a repository.

    Sets cache statistics, encryption information, the last-modified timestamp,
    and metrics for the latest archive. When ``all_archives`` is true, also
    collects metrics for every individual archive.

    :param repository: Repository information returned by borgmatic.
    :param all_archives: Whether to collect metrics for every archive.
    """
    repository_metadata = repository["repository"]

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

    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"], err=True
        )
        return

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

    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 Prometheus metrics for an archive.

    :param archive: Archive information returned by borgmatic.
    :param repository_labels: Labels inherited from the parent repository.
    """
    msg("Processing archive:", archive["name"], f"({archive['id']})", err=True)

    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", "-?")}, params=[]
)
@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",
    help="Path to the borgmatic executable.",
)
@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",
    help="Directory where the Prometheus collector textfile will be written.",
)
@click.option(
    "--all-archives",
    "all_archives_flag",
    is_flag=True,
    default=False,
    help=(
        "Collect metrics for all archives instead of only the latest archive. "
        "This can significantly increase runtime and Prometheus storage usage."
    ),
)
@click.option(
    "-q",
    "--quiet",
    "quiet_flag",
    is_flag=True,
    default=False,
    help="Suppress all non-error output, including verbose output.",
)
@click.option(
    "-v",
    "--verbose",
    "verbose_flag",
    is_flag=True,
    default=False,
    help="Enable verbose output.",
)
@click.option("--color", cls=FlexibleColorOption)
@click.no_color_option()
@click.option(
    "--version",
    callback=print_version,
    default=False,
    expose_value=False,
    help="Show version information 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,
) -> None:
    """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 for each repository. Use
    ``--all-archives`` to collect metrics for every archive, at the cost of a
    significantly longer runtime and increased 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:", err=True)
    debug("  borgmatic_bin:", str(borgmatic_bin), err=True)
    debug(
        "  borgmatic_configs:",
        ", ".join(str(config) for config in borgmatic_configs)
        or "<borgmatic defaults>",
        err=True,
    )
    debug("  textfile_collector_dir:", str(textfile_collector_dir), err=True)
    debug("  all_archives:", str(all_archives_flag), err=True)
    debug("  verbose:", str(verbose_flag), err=True)
    debug("  quiet:", str(quiet_flag), err=True)

    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()), err=True)


if __name__ == "__main__":
    prometheus_borgmatic_exporter()