#!/usr/bin/env python3
from __future__ import annotations

import argparse
import getpass
import json
import shlex
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

import yaml

REPO_DIR = Path(__file__).resolve().parents[1]


class BaoRunner:
    def __init__(
        self,
        *,
        kubectl: str,
        namespace: str,
        release: str,
        dry_run: bool,
        use_token_helper: bool,
        token: str | None,
    ) -> None:
        self.kubectl_parts = shlex.split(kubectl)
        self.namespace = namespace
        self.pod = f"{release}-0"
        self.dry_run = dry_run
        self.use_token_helper = use_token_helper
        self.token = token

    def run(
        self,
        args: list[str],
        *,
        input_text: str | None = None,
        check: bool = True,
        quiet: bool = False,
    ) -> subprocess.CompletedProcess[str]:
        if self.dry_run:
            print("DRY-RUN: bao " + shlex.join(args))
            return subprocess.CompletedProcess(args, 0, "", "")

        if self.use_token_helper:
            cmd = (
                self.kubectl_parts
                + ["exec", "-i", "-n", self.namespace, self.pod, "--", "bao"]
                + args
            )
            proc_input = input_text
        else:
            if not self.token:
                raise RuntimeError(
                    "OpenBao token is required unless --use-token-helper is set"
                )
            cmd = (
                self.kubectl_parts
                + [
                    "exec",
                    "-i",
                    "-n",
                    self.namespace,
                    self.pod,
                    "--",
                    "sh",
                    "-c",
                    'read -r BAO_TOKEN; export BAO_TOKEN; export VAULT_TOKEN="$BAO_TOKEN"; exec bao "$@"',
                    "sh",
                ]
                + args
            )
            proc_input = self.token + "\n" + (input_text or "")

        result = subprocess.run(cmd, input=proc_input, capture_output=True, text=True)
        if check and result.returncode != 0:
            if result.stdout and not quiet:
                print(result.stdout, end="")
            if result.stderr:
                print(result.stderr, file=sys.stderr, end="")
            raise SystemExit(result.returncode)
        if not quiet and result.stdout:
            print(result.stdout, end="")
        if not quiet and result.stderr:
            print(result.stderr, file=sys.stderr, end="")
        return result


def run_with_token(
    *,
    kubectl: str,
    namespace: str,
    release: str,
    token: str,
    args: list[str],
    check: bool,
) -> subprocess.CompletedProcess[str]:
    kubectl_parts = shlex.split(kubectl)
    cmd = (
        kubectl_parts
        + [
            "exec",
            "-i",
            "-n",
            namespace,
            f"{release}-0",
            "--",
            "sh",
            "-c",
            'read -r BAO_TOKEN; export BAO_TOKEN; export VAULT_TOKEN="$BAO_TOKEN"; exec bao "$@"',
            "sh",
        ]
        + args
    )
    return subprocess.run(
        cmd, input=token + "\n", capture_output=True, text=True, check=False
    )


def read_token(
    token_file: str | None, dry_run: bool, use_token_helper: bool
) -> str | None:
    if dry_run or use_token_helper:
        return None
    if token_file:
        path = Path(token_file).expanduser()
        if not path.exists():
            raise SystemExit(f"ERROR: OPENBAO_TOKEN_FILE does not exist: {path}")
        lines = path.read_text(encoding="utf-8").splitlines()
        token = lines[0].strip() if lines else ""
        if not token:
            raise SystemExit(f"ERROR: OPENBAO_TOKEN_FILE is empty: {path}")
        return token
    token = getpass.getpass("OpenBao token: ")
    if not token:
        raise SystemExit("ERROR: empty OpenBao token")
    return token


def load_catalog(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle)
    if not isinstance(data, dict):
        raise SystemExit(f"ERROR: catalog root must be an object: {path}")
    return data


def selected_grants(
    catalog: dict[str, Any], grant_id: str | None
) -> list[dict[str, Any]]:
    grants = catalog.get("grants") or []
    if not isinstance(grants, list):
        raise SystemExit("ERROR: catalog grants must be a list")
    selected = [
        grant for grant in grants if not grant_id or grant.get("id") == grant_id
    ]
    if grant_id and not selected:
        raise SystemExit(f"ERROR: grant not found in catalog: {grant_id}")
    return selected


def verify_static(runner: BaoRunner, grant: dict[str, Any]) -> None:
    openbao = grant["openbao"]
    runner.run(["read", f"auth/token/roles/{openbao['token_role']}"])
    runner.run(["policy", "read", openbao["issuer_policy"]])
    runner.run(["policy", "read", openbao["policies"][0]])
    print(f"OK: static token-grant config readable for {grant['id']}")


def issue_smoke_token(
    runner: BaoRunner,
    *,
    kubectl: str,
    namespace: str,
    release: str,
    grant: dict[str, Any],
) -> None:
    openbao = grant["openbao"]
    ttl = grant["ttl"]["default"]
    args = [
        "token",
        "create",
        f"-role={openbao['token_role']}",
        f"-ttl={ttl}",
        "-format=json",
    ]
    for policy in openbao["policies"]:
        args.append(f"-policy={policy}")
    result = runner.run(args, quiet=True, check=False)
    if result.returncode != 0:
        raise SystemExit(
            f"ERROR: token create failed (rc={result.returncode}): "
            f"{(result.stderr or result.stdout or '').strip()}"
        )
    try:
        payload = json.loads(result.stdout)
        auth = payload.get("auth") or payload.get("data") or {}
        child_token = auth["client_token"]
        accessor = auth["accessor"]
    except Exception as exc:  # noqa: BLE001
        raise SystemExit(
            f"ERROR: could not parse token create response: {exc}"
        ) from exc

    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            key_path = Path(tmpdir) / "warden-sign-smoke_ed25519"
            keygen = subprocess.run(
                ["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-f", str(key_path)],
                capture_output=True,
                text=True,
                check=False,
            )
            if keygen.returncode != 0:
                raise SystemExit(
                    "ERROR: could not generate smoke SSH key: "
                    f"{(keygen.stderr or keygen.stdout).strip()}"
                )
            public_key = key_path.with_suffix(key_path.suffix + ".pub").read_text(encoding="utf-8").strip()

        positive = run_with_token(
            kubectl=kubectl,
            namespace=namespace,
            release=release,
            token=child_token,
            args=["write", "-field=signed_key", "ssh/sign/agt-role", f"public_key={public_key}"],
            check=False,
        )
        if positive.returncode != 0 or not positive.stdout.strip():
            raise SystemExit(
                "ERROR: child token could not sign with ssh/sign/agt-role: "
                f"{(positive.stderr or positive.stdout).strip()}"
            )
        print("OK: child token can sign with ssh/sign/agt-role")

        negative = run_with_token(
            kubectl=kubectl,
            namespace=namespace,
            release=release,
            token=child_token,
            args=["policy", "read", "warden-sign"],
            check=False,
        )
        if negative.returncode == 0:
            raise SystemExit("ERROR: child token unexpectedly read policy metadata")
        print("OK: child token cannot read policy metadata")
    finally:
        runner.run(["write", "auth/token/revoke-accessor", f"accessor={accessor}"], quiet=True)
        print("OK: smoke child token revoked by accessor")


def main() -> int:
    parser = argparse.ArgumentParser(description="Verify OpenBao token grants.")
    parser.add_argument("--catalog", default="credential-grants/catalog.yaml")
    parser.add_argument("--grant", help="Limit to one grant id")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--use-token-helper",
        action="store_true",
        help="Use the OpenBao CLI token helper inside the pod",
    )
    parser.add_argument(
        "--issue-smoke-token",
        action="store_true",
        help="Mint and revoke a short-lived child token for live verification",
    )
    parser.add_argument("--namespace", default=None)
    parser.add_argument("--release", default=None)
    parser.add_argument("--kubectl", default=None)
    args = parser.parse_args()

    import os

    namespace = args.namespace or os.environ.get("OPENBAO_NAMESPACE", "openbao")
    release = args.release or os.environ.get("OPENBAO_RELEASE", "openbao")
    kubectl = args.kubectl or os.environ.get("KUBECTL", "kubectl")
    token_file = os.environ.get("OPENBAO_TOKEN_FILE")
    token = read_token(token_file, args.dry_run, args.use_token_helper)
    catalog = load_catalog(REPO_DIR / args.catalog)
    runner = BaoRunner(
        kubectl=kubectl,
        namespace=namespace,
        release=release,
        dry_run=args.dry_run,
        use_token_helper=args.use_token_helper,
        token=token,
    )

    for grant in selected_grants(catalog, args.grant):
        verify_static(runner, grant)
        if args.issue_smoke_token:
            if args.dry_run:
                print(
                    f"DRY-RUN: would mint and revoke smoke child token for {grant['id']}"
                )
            else:
                issue_smoke_token(
                    runner,
                    kubectl=kubectl,
                    namespace=namespace,
                    release=release,
                    grant=grant,
                )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
