#!/usr/bin/env python3
"""Package telemetry collector assets for offline/air-gapped installs.

Outputs a tarball containing the collector configuration, Compose overlay,
Helm defaults, and operator README. A SHA-256 checksum sidecar is emitted, and
optional Cosign signing can be enabled with --sign.
"""
from __future__ import annotations

import argparse
import hashlib
import os
import subprocess
import sys
import tarfile
from pathlib import Path
from typing import Iterable

REPO_ROOT = Path(__file__).resolve().parents[3]
DEFAULT_OUTPUT = REPO_ROOT / "out" / "telemetry" / "telemetry-offline-bundle.tar.gz"
BUNDLE_CONTENTS: tuple[Path, ...] = (
    Path("deploy/telemetry/README.md"),
    Path("deploy/telemetry/otel-collector-config.yaml"),
    Path("deploy/telemetry/storage/README.md"),
    Path("deploy/telemetry/storage/prometheus.yaml"),
    Path("deploy/telemetry/storage/tempo.yaml"),
    Path("deploy/telemetry/storage/loki.yaml"),
    Path("deploy/telemetry/storage/tenants/tempo-overrides.yaml"),
    Path("deploy/telemetry/storage/tenants/loki-overrides.yaml"),
    Path("deploy/helm/stellaops/files/otel-collector-config.yaml"),
    Path("deploy/helm/stellaops/values.yaml"),
    Path("deploy/helm/stellaops/templates/otel-collector.yaml"),
    Path("deploy/compose/docker-compose.telemetry.yaml"),
    Path("deploy/compose/docker-compose.telemetry-storage.yaml"),
    Path("docs/ops/telemetry-collector.md"),
    Path("docs/ops/telemetry-storage.md"),
)


def compute_sha256(path: Path) -> str:
    sha = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            sha.update(chunk)
    return sha.hexdigest()


def validate_files(paths: Iterable[Path]) -> None:
    missing = [str(p) for p in paths if not (REPO_ROOT / p).exists()]
    if missing:
        raise FileNotFoundError(f"Missing bundle artefacts: {', '.join(missing)}")


def create_bundle(output_path: Path) -> Path:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with tarfile.open(output_path, "w:gz") as tar:
        for rel_path in BUNDLE_CONTENTS:
            abs_path = REPO_ROOT / rel_path
            tar.add(abs_path, arcname=str(rel_path))
    return output_path


def write_checksum(bundle_path: Path) -> Path:
    digest = compute_sha256(bundle_path)
    sha_path = bundle_path.with_suffix(bundle_path.suffix + ".sha256")
    sha_path.write_text(f"{digest}  {bundle_path.name}\n", encoding="utf-8")
    return sha_path


def cosign_sign(bundle_path: Path, key_ref: str | None, identity_token: str | None) -> None:
    cmd = ["cosign", "sign-blob", "--yes", str(bundle_path)]
    if key_ref:
        cmd.extend(["--key", key_ref])
    env = os.environ.copy()
    if identity_token:
        env["COSIGN_IDENTITY_TOKEN"] = identity_token
    try:
        subprocess.run(cmd, check=True, env=env)
    except FileNotFoundError as exc:
        raise RuntimeError("cosign not found on PATH; install cosign or omit --sign") from exc
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(f"cosign sign-blob failed: {exc}") from exc


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_OUTPUT,
        help=f"Output bundle path (default: {DEFAULT_OUTPUT})",
    )
    parser.add_argument(
        "--sign",
        action="store_true",
        help="Sign the bundle using cosign (requires cosign on PATH)",
    )
    parser.add_argument(
        "--cosign-key",
        type=str,
        default=os.environ.get("COSIGN_KEY_REF"),
        help="Cosign key reference (file:..., azurekms://..., etc.)",
    )
    parser.add_argument(
        "--identity-token",
        type=str,
        default=os.environ.get("COSIGN_IDENTITY_TOKEN"),
        help="OIDC identity token for keyless signing",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    validate_files(BUNDLE_CONTENTS)

    bundle_path = args.output.resolve()
    print(f"[*] Creating telemetry bundle at {bundle_path}")
    create_bundle(bundle_path)
    sha_path = write_checksum(bundle_path)
    print(f"[✓] SHA-256 written to {sha_path}")

    if args.sign:
        print("[*] Signing bundle with cosign")
        cosign_sign(bundle_path, args.cosign_key, args.identity_token)
        sig_path = bundle_path.with_suffix(bundle_path.suffix + ".sig")
        if sig_path.exists():
            print(f"[✓] Cosign signature written to {sig_path}")
        else:
            print("[!] Cosign completed but signature file not found (ensure cosign version >= 2.2)")

    return 0


if __name__ == "__main__":
    sys.exit(main())