#!/usr/bin/env python3
"""Manage cube-cloud API keys in DynamoDB.

Usage:
    python scripts/manage-keys.py create --profile admin --description "Vinay's key"
    python scripts/manage-keys.py list
    python scripts/manage-keys.py revoke --hash <key_hash>
"""
import argparse
import hashlib
import json
import secrets
from datetime import datetime, timezone

import boto3

REGION = "us-west-2"
TABLE = "cube-mcp-prod-api-keys"


def get_table():
    return boto3.resource("dynamodb", region_name=REGION).Table(TABLE)


def cmd_create(args):
    raw_key = f"cmcp_{secrets.token_urlsafe(32)}"
    key_hash = hashlib.sha256(raw_key.encode()).hexdigest()

    get_table().put_item(Item={
        "key_hash": key_hash,
        "profile": args.profile,
        "description": args.description or "",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "enabled": True,
    })

    print(f"API key created!")
    print(f"  Key:     {raw_key}")
    print(f"  Hash:    {key_hash}")
    print(f"  Profile: {args.profile}")
    print()
    print("Save this key — it cannot be retrieved again.")
    print()
    print("Use in Claude Desktop remote MCP server config:")
    print(f'  "headers": {{"Authorization": "Bearer {raw_key}"}}')


def cmd_list(args):
    resp = get_table().scan()
    items = resp.get("Items", [])
    if not items:
        print("No API keys found.")
        return

    print(f"{'Hash (first 12)':>14}  {'Profile':<20} {'Enabled':<8} {'Description'}")
    print("-" * 70)
    for item in sorted(items, key=lambda x: x.get("created_at", "")):
        h = item["key_hash"][:12] + "..."
        print(f"{h:>14}  {item.get('profile', ''):20} {str(item.get('enabled', True)):8} {item.get('description', '')}")


def cmd_revoke(args):
    get_table().update_item(
        Key={"key_hash": args.hash},
        UpdateExpression="SET enabled = :v",
        ExpressionAttributeValues={":v": False},
    )
    print(f"Key {args.hash[:12]}... revoked.")


def main():
    parser = argparse.ArgumentParser(description="Manage cube-cloud API keys")
    sub = parser.add_subparsers(dest="command", required=True)

    create = sub.add_parser("create", help="Create a new API key")
    create.add_argument("--profile", required=True, help="Profile name (e.g. admin)")
    create.add_argument("--description", help="Key description")

    sub.add_parser("list", help="List all API keys")

    revoke = sub.add_parser("revoke", help="Revoke an API key")
    revoke.add_argument("--hash", required=True, help="Key hash to revoke")

    args = parser.parse_args()
    {"create": cmd_create, "list": cmd_list, "revoke": cmd_revoke}[args.command](args)


if __name__ == "__main__":
    main()
