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

import argparse
import getpass
import hashlib
import json
import os
import re
import shlex
import subprocess
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

import yaml

REPO_DIR = Path(__file__).resolve().parents[1]
DEFAULT_CATALOG = REPO_DIR / "credential-grants/catalog.yaml"
DEFAULT_LEASE_DIR = REPO_DIR / ".local/credential-leases"
TOKEN_MARKERS = re.compile(
    r"\bhv[bcms]\.[A-Za-z0-9._-]+|\b(?:VAULT_TOKEN|BAO_TOKEN)=[^\s]+"
)
ENV_ASSIGNMENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=")
UNSAFE_TOKEN_ENV = {"VAULT_TOKEN", "BAO_TOKEN", "OPENBAO_ROOT_TOKEN"}
UNSAFE_VERBOSE_VALUES = {"debug", "trace"}


def emit_json(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, indent=2, sort_keys=True))


def fail(message: str) -> None:
    raise SystemExit(f"ERROR: {message}")


def redact(text: str, extra_secrets: list[str] | None = None) -> str:
    redacted = TOKEN_MARKERS.sub("[REDACTED]", text)
    for secret in extra_secrets or []:
        if secret:
            redacted = redacted.replace(secret, "[REDACTED]")
    return redacted


def guard_verbose_env() -> None:
    for name in ("BAO_LOG_LEVEL", "VAULT_LOG_LEVEL"):
        value = os.environ.get(name, "").lower()
        if value in UNSAFE_VERBOSE_VALUES:
            fail(
                f"refusing to run with {name}={value}; it may print secret-bearing traces"
            )


def read_issuer_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():
            fail(f"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:
            fail(f"OPENBAO_TOKEN_FILE is empty: {path}")
        return token
    token = getpass.getpass("OpenBao issuer token: ")
    if not token:
        fail("empty OpenBao issuer token")
    return token


def ttl_seconds(value: str) -> int:
    match = re.match(r"^([1-9][0-9]*)([smhd])$", value)
    if not match:
        fail(f"TTL must match <positive integer><s|m|h|d>: {value!r}")
    amount = int(match.group(1))
    multiplier = {"s": 1, "m": 60, "h": 3600, "d": 86400}[match.group(2)]
    return amount * multiplier


def expires_at(ttl: str) -> str:
    return (
        datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds(ttl))
    ).isoformat()


def resolve_repo_path(path: Path) -> Path:
    path = path.expanduser()
    if path.is_absolute():
        return path
    return (REPO_DIR / path).resolve()


def load_catalog(path: Path) -> dict[str, Any]:
    data = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        fail(f"catalog root must be an object: {path}")
    return data


def get_grant(catalog: dict[str, Any], grant_id: str) -> dict[str, Any]:
    grants = catalog.get("grants") or []
    if not isinstance(grants, list):
        fail("catalog grants must be a list")
    for grant in grants:
        if isinstance(grant, dict) and grant.get("id") == grant_id:
            if grant.get("credential_type") != "openbao-token":
                fail(
                    f"unsupported credential_type for {grant_id}: {grant.get('credential_type')}"
                )
            return grant
    fail(f"grant not found in catalog: {grant_id}")


def validate_issue_request(
    grant: dict[str, Any], ttl: str, purpose: str, delivery: str | None = None
) -> None:
    if not purpose.strip():
        fail("--purpose is required and must not be empty")
    ttl_cfg = grant.get("ttl") or {}
    max_ttl = ttl_cfg.get("max")
    if not isinstance(max_ttl, str):
        fail(f"grant {grant.get('id')} has no ttl.max")
    if ttl_seconds(ttl) > ttl_seconds(max_ttl):
        fail(f"requested TTL {ttl} exceeds grant max TTL {max_ttl}")
    if delivery is not None:
        allowed = set((grant.get("delivery") or {}).get("allowed") or [])
        if delivery not in allowed:
            fail(
                f"delivery mode {delivery!r} is not allowed for grant {grant.get('id')}"
            )


def safe_handle(handle: str) -> str:
    digest = hashlib.sha256(handle.encode("utf-8")).hexdigest()[:16]
    cleaned = re.sub(r"[^A-Za-z0-9_.-]", "_", handle)[:64]
    return f"{cleaned}-{digest}"


class BaoRunner:
    def __init__(
        self,
        *,
        kubectl: str,
        namespace: str,
        release: str,
        dry_run: bool,
        use_token_helper: bool,
        issuer_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.issuer_token = issuer_token

    def run(
        self,
        args: list[str],
        *,
        input_text: str | None = None,
        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.issuer_token:
                raise RuntimeError(
                    "issuer 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; exec bao "$@"',
                    "sh",
                ]
                + args
            )
            proc_input = self.issuer_token + "\n" + (input_text or "")

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


def token_create_args(
    grant: dict[str, Any], ttl: str, wrap_ttl: str | None = None
) -> list[str]:
    openbao = grant["openbao"]
    args = ["token", "create"]
    if wrap_ttl:
        args.append(f"-wrap-ttl={wrap_ttl}")
    args.extend(
        [
            f"-role={openbao['token_role']}",
            f"-ttl={ttl}",
            "-format=json",
        ]
    )
    for policy in openbao["policies"]:
        args.append(f"-policy={policy}")
    return args


def parse_token_create(stdout: str) -> tuple[str, str]:
    try:
        payload = json.loads(stdout)
        auth = payload.get("auth") or payload.get("data") or {}
        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
    return token, accessor


def parse_wrap_create(stdout: str) -> dict[str, Any]:
    try:
        payload = json.loads(stdout)
        wrap_info = payload["wrap_info"]
    except Exception as exc:  # noqa: BLE001
        raise SystemExit(
            f"ERROR: could not parse wrapped token response: {exc}"
        ) from exc
    return wrap_info


def write_mode_0600(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.parent.chmod(0o700)
    flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    fd = os.open(path, flags, 0o600)
    with os.fdopen(fd, "w", encoding="utf-8") as handle:
        handle.write(content)
    path.chmod(0o600)


def write_local_lease(
    *,
    lease_dir: Path,
    grant: dict[str, Any],
    purpose: str,
    ttl: str,
    token: str,
    accessor: str,
) -> dict[str, Any]:
    stem = safe_handle(accessor)
    token_path = lease_dir / f"{stem}.openbao-token"
    meta_path = lease_dir / f"{stem}.json"
    write_mode_0600(token_path, token + "\n")
    metadata = {
        "grant_id": grant["id"],
        "lease_handle": accessor,
        "accessor": accessor,
        "purpose": purpose,
        "issued_at": datetime.now(timezone.utc).isoformat(),
        "expires_at": expires_at(ttl),
        "ttl": ttl,
        "delivery_mode": "local-token-file",
        "token_file": str(token_path),
        "status": "issued",
    }
    write_mode_0600(meta_path, json.dumps(metadata, indent=2, sort_keys=True) + "\n")
    return {key: value for key, value in metadata.items() if key != "token_file"} | {
        "token_file": str(token_path),
        "metadata_file": str(meta_path),
    }


def find_local_metadata(
    lease_dir: Path, handle: str
) -> tuple[Path, dict[str, Any]] | None:
    if not lease_dir.exists():
        return None
    for path in lease_dir.glob("*.json"):
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            continue
        if data.get("lease_handle") == handle or data.get("accessor") == handle:
            return path, data
    return None


def path_is_within(path: Path, root: Path) -> bool:
    try:
        path.resolve().relative_to(root.resolve())
        return True
    except ValueError:
        return False


def remove_local_lease_files(lease_dir: Path, handle: str) -> list[str]:
    found = find_local_metadata(lease_dir, handle)
    if not found:
        return []
    meta_path, metadata = found
    lease_root = lease_dir.resolve()
    removed: list[str] = []
    for key in ("token_file",):
        value = metadata.get(key)
        if isinstance(value, str):
            path = Path(value)
            if not path.is_absolute():
                path = (REPO_DIR / path).resolve()
            if path.exists() and path.is_file():
                if not path_is_within(path, lease_root):
                    fail(
                        f"refusing to remove local lease file outside {lease_root}: {path}"
                    )
                path.unlink()
                removed.append(str(path))
    if meta_path.exists():
        if not path_is_within(meta_path, lease_root):
            fail(f"refusing to remove metadata file outside {lease_root}: {meta_path}")
        meta_path.unlink()
        removed.append(str(meta_path))
    return removed


def split_env_prefix(command: list[str]) -> tuple[dict[str, str], list[str]]:
    if command and command[0] == "--":
        command = command[1:]
    extra_env: dict[str, str] = {}
    index = 0
    while index < len(command) and ENV_ASSIGNMENT.match(command[index]):
        name, value = command[index].split("=", 1)
        if name in UNSAFE_TOKEN_ENV:
            fail(
                f"refusing caller-supplied {name}; the helper owns credential injection"
            )
        if TOKEN_MARKERS.search(value):
            fail(f"refusing secret-looking value in env assignment {name}")
        extra_env[name] = value
        index += 1
    program = command[index:]
    if not program:
        fail("missing child command after --")
    return extra_env, program


def command_request(
    args: argparse.Namespace, runner: BaoRunner, grant: dict[str, Any]
) -> int:
    ttl = args.ttl or grant["ttl"]["default"]
    validate_issue_request(grant, ttl, args.purpose, args.delivery)
    if args.dry_run:
        emit_json(
            {
                "status": "dry-run",
                "command": "request",
                "grant_id": grant["id"],
                "purpose": args.purpose,
                "ttl": ttl,
                "delivery_mode": args.delivery,
            }
        )
        if args.delivery == "response-wrap":
            print(
                "DRY-RUN: bao "
                + shlex.join(token_create_args(grant, ttl, args.wrap_ttl))
            )
        else:
            print("DRY-RUN: bao " + shlex.join(token_create_args(grant, ttl)))
        return 0

    if args.delivery == "response-wrap":
        result = runner.run(token_create_args(grant, ttl, args.wrap_ttl), quiet=True)
        wrap_info = parse_wrap_create(result.stdout)
        emit_json(
            {
                "grant_id": grant["id"],
                "purpose": args.purpose,
                "ttl": ttl,
                "delivery_mode": "response-wrap",
                "wrapping_token": wrap_info.get("token"),
                "wrapping_accessor": wrap_info.get("accessor"),
                "wrapped_accessor": wrap_info.get("wrapped_accessor"),
                "wrap_ttl": wrap_info.get("ttl"),
                "status": "wrapped",
            }
        )
        return 0

    result = runner.run(token_create_args(grant, ttl), quiet=True)
    token, accessor = parse_token_create(result.stdout)
    payload = write_local_lease(
        lease_dir=args.lease_dir,
        grant=grant,
        purpose=args.purpose,
        ttl=ttl,
        token=token,
        accessor=accessor,
    )
    emit_json(payload)
    return 0


def command_exec(
    args: argparse.Namespace, runner: BaoRunner, grant: dict[str, Any]
) -> int:
    ttl = args.ttl or grant["ttl"]["default"]
    validate_issue_request(grant, ttl, args.purpose, "exec-env")
    extra_env, program = split_env_prefix(args.command)
    if args.dry_run:
        emit_json(
            {
                "status": "dry-run",
                "command": "exec",
                "grant_id": grant["id"],
                "purpose": args.purpose,
                "ttl": ttl,
                "delivery_mode": "exec-env",
                "child_env": sorted([*extra_env.keys(), "VAULT_TOKEN"]),
                "child_command": program,
            }
        )
        print("DRY-RUN: bao " + shlex.join(token_create_args(grant, ttl)))
        return 0

    result = runner.run(token_create_args(grant, ttl), quiet=True)
    token, accessor = parse_token_create(result.stdout)
    env = os.environ.copy()
    env.update(extra_env)
    env["VAULT_TOKEN"] = token
    try:
        child = subprocess.run(
            program, env=env, capture_output=True, text=True, check=False
        )
        if child.stdout:
            print(redact(child.stdout, [token]), end="")
        if child.stderr:
            print(redact(child.stderr, [token]), file=sys.stderr, end="")
        return child.returncode
    finally:
        runner.run(
            ["write", "auth/token/revoke-accessor", f"accessor={accessor}"], quiet=True
        )


def command_status(args: argparse.Namespace, runner: BaoRunner) -> int:
    local = find_local_metadata(args.lease_dir, args.lease_handle)
    if args.dry_run:
        print(
            "DRY-RUN: bao write -format=json auth/token/lookup-accessor accessor=<lease_handle>"
        )
        emit_json(
            {
                "status": "dry-run",
                "lease_handle": args.lease_handle,
                "local_metadata_found": local is not None,
            }
        )
        return 0
    result = runner.run(
        [
            "write",
            "-format=json",
            "auth/token/lookup-accessor",
            f"accessor={args.lease_handle}",
        ],
        quiet=True,
    )
    payload = json.loads(result.stdout)
    if local:
        _, metadata = local
        payload.setdefault("local", metadata | {"token_file": "[LOCAL-SECRET-FILE]"})
    print(redact(json.dumps(payload, indent=2, sort_keys=True)))
    return 0


def command_revoke(args: argparse.Namespace, runner: BaoRunner) -> int:
    if args.dry_run:
        print("DRY-RUN: bao write auth/token/revoke-accessor accessor=<lease_handle>")
        emit_json({"status": "dry-run", "lease_handle": args.lease_handle})
        return 0
    runner.run(
        ["write", "auth/token/revoke-accessor", f"accessor={args.lease_handle}"],
        quiet=True,
    )
    removed = remove_local_lease_files(args.lease_dir, args.lease_handle)
    emit_json(
        {
            "status": "revoked",
            "lease_handle": args.lease_handle,
            "removed_local_files": removed,
        }
    )
    return 0


def add_common_issue_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--grant", default="ops-warden/warden-sign")
    parser.add_argument("--purpose", required=True)
    parser.add_argument("--ttl")
    parser.add_argument("--dry-run", action="store_true")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Request, use, inspect, and revoke bounded OpenBao credential leases."
    )
    parser.add_argument("--catalog", default=str(DEFAULT_CATALOG))
    parser.add_argument(
        "--namespace", default=os.environ.get("OPENBAO_NAMESPACE", "openbao")
    )
    parser.add_argument(
        "--release", default=os.environ.get("OPENBAO_RELEASE", "openbao")
    )
    parser.add_argument("--kubectl", default=os.environ.get("KUBECTL", "kubectl"))
    parser.add_argument(
        "--issuer-token-file", default=os.environ.get("OPENBAO_TOKEN_FILE")
    )
    parser.add_argument(
        "--use-token-helper",
        action="store_true",
        help="Use the OpenBao CLI token helper inside the pod",
    )
    parser.add_argument("--lease-dir", type=Path, default=DEFAULT_LEASE_DIR)

    subparsers = parser.add_subparsers(dest="command_name", required=True)

    request = subparsers.add_parser(
        "request", help="Issue a wrapped token or local-token-file lease"
    )
    add_common_issue_args(request)
    request.add_argument(
        "--delivery",
        choices=["local-token-file", "response-wrap"],
        default="local-token-file",
    )
    request.add_argument("--wrap-ttl", default="5m")

    exec_parser = subparsers.add_parser(
        "exec", help="Run a child process with VAULT_TOKEN injected"
    )
    add_common_issue_args(exec_parser)
    exec_parser.add_argument("command", nargs=argparse.REMAINDER)

    status = subparsers.add_parser(
        "status", help="Look up a lease by non-secret accessor handle"
    )
    status.add_argument("--dry-run", action="store_true")
    status.add_argument("lease_handle")

    revoke = subparsers.add_parser(
        "revoke", help="Revoke a lease by non-secret accessor handle"
    )
    revoke.add_argument("--dry-run", action="store_true")
    revoke.add_argument("lease_handle")

    return parser


def main() -> int:
    guard_verbose_env()
    parser = build_parser()
    args = parser.parse_args()
    args.lease_dir = resolve_repo_path(args.lease_dir)
    catalog = load_catalog(resolve_repo_path(Path(args.catalog)))
    grant = None
    if args.command_name in {"request", "exec"}:
        grant = get_grant(catalog, args.grant)
    issuer_token = read_issuer_token(
        args.issuer_token_file, args.dry_run, args.use_token_helper
    )
    runner = BaoRunner(
        kubectl=args.kubectl,
        namespace=args.namespace,
        release=args.release,
        dry_run=args.dry_run,
        use_token_helper=args.use_token_helper,
        issuer_token=issuer_token,
    )

    if args.command_name == "request":
        assert grant is not None
        return command_request(args, runner, grant)
    if args.command_name == "exec":
        assert grant is not None
        return command_exec(args, runner, grant)
    if args.command_name == "status":
        return command_status(args, runner)
    if args.command_name == "revoke":
        return command_revoke(args, runner)
    parser.error(f"unknown command: {args.command_name}")
    return 2


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