#!/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
from .cli import FlexibleColorOption, debug, error, msg, verbose, warn

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"


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()


@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", "-?")}, 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 textfile collector file is 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 normal 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 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()