#!/usr/bin/env python3
"""
One-shot: create a dedicated NewAPI sk- token for the downstream
"蜂群程序" (swarm bee program) and print the base URL + key.

Why direct DB INSERT instead of the /api/token POST endpoint:
  - that endpoint requires a logged-in admin session cookie we
    don't have a clean way to mint headless;
  - the token table shape is stable (new-api fork has tracked
    these columns for two years).

Safety:
  - Attaches the token to the platform owner user (whose email
    we look up first).
  - Marks UnlimitedQuota=true so this token consumes user-row
    quota (no separate billing surface to manage).
  - HideFromUserUI=false so the operator can revoke it from the
    /keys page if needed.
  - Sets a descriptive Name + Remark per the user's "备注好"
    requirement.
"""
import os
import sys
import secrets
import string

import paramiko

HOST = "20.24.50.121"
USER = "heicode"
PASS = os.environ.get("MGR_PASS", "")
if not PASS:
    print("MGR_PASS not set", file=sys.stderr)
    sys.exit(2)

DB_PASS = os.environ.get("MGR_DB_PASS", "")
DB_HOST = "heicode.postgres.database.azure.com"
DB_USER = "heicode"
DB_NAME = "heicode"

if not DB_PASS:
    print("MGR_DB_PASS not set", file=sys.stderr)
    sys.exit(2)

# new-api token keys are 48-char alphanum (no padding). Match that
# exactly so admin tools / dashboards parse it.
ALPHABET = string.ascii_letters + string.digits
RAW_KEY = "".join(secrets.choice(ALPHABET) for _ in range(48))
# Token bearer format on the wire is "sk-<48 chars>".

SQL_FIND = (
    "SELECT id, username, email FROM users "
    "WHERE email IN ('zsbgnw@gmail.com','chenchen@xinghanlab.com') "
    "   OR username IN ('chenchen','root') "
    "   OR role >= 10 "
    "ORDER BY role DESC, id ASC LIMIT 1;"
)

INSERT_SQL_TEMPLATE = (
    "INSERT INTO tokens "
    '(user_id, name, "key", status, created_time, accessed_time, '
    "expired_time, remain_quota, used_quota, unlimited_quota, "
    "model_limits_enabled, model_limits, allow_ips, "
    '"group", cross_group_retry, hide_from_user_ui) '
    "VALUES "
    "({user_id}, '蜂群程序 (Swarm Bot)', '{key}', 1, "
    "EXTRACT(EPOCH FROM NOW())::bigint, EXTRACT(EPOCH FROM NOW())::bigint, "
    "-1, 0, 0, true, false, '', '', '', false, false) "
    "RETURNING id, name, \"key\";"
)


def main() -> None:
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(HOST, 22, USER, PASS, look_for_keys=False, allow_agent=False, timeout=30)

    def psql(sql: str) -> str:
        # Pipe SQL via stdin so quoted identifiers like "group" survive
        # without shell-escape hell.
        cmd = (
            f"PGPASSWORD='{DB_PASS}' psql "
            f"'sslmode=require host={DB_HOST} user={DB_USER} dbname={DB_NAME}' "
            f"-At -F'|' -v ON_ERROR_STOP=1"
        )
        stdin, stdout, stderr = c.exec_command(cmd, timeout=60, get_pty=False)
        stdin.write(sql)
        stdin.channel.shutdown_write()
        out = stdout.read().decode("utf-8", "replace").strip()
        err = stderr.read().decode("utf-8", "replace").strip()
        if err:
            print(f"[psql stderr] {err}", file=sys.stderr)
        return out

    # Step 1: find owner user
    print(f"--- Looking up platform owner user ---", flush=True)
    owner_row = psql(SQL_FIND)
    if not owner_row:
        print("ERROR: no owner user found", file=sys.stderr)
        sys.exit(3)
    fields = owner_row.split("|")
    owner_id = int(fields[0])
    owner_username = fields[1] if len(fields) > 1 else ""
    owner_email = fields[2] if len(fields) > 2 else ""
    print(f"Owner: id={owner_id} username={owner_username} email={owner_email}")

    # Step 2: insert the token row
    print(f"\n--- Inserting dedicated 蜂群程序 token ---", flush=True)
    insert_sql = INSERT_SQL_TEMPLATE.format(user_id=owner_id, key=RAW_KEY)
    inserted = psql(insert_sql)
    if not inserted:
        print("ERROR: INSERT returned nothing", file=sys.stderr)
        sys.exit(4)
    fields = inserted.split("|")
    token_id = int(fields[0])
    token_name = fields[1]
    token_key = fields[2]
    print(f"Token row inserted: id={token_id} name={token_name}")

    # Step 3: print result
    base_url = "https://code.xinghanlab.com"
    print("\n" + "=" * 60)
    print("NEWAPI CREDENTIAL — 蜂群程序 (SWARM BOT)")
    print("=" * 60)
    print(f"Base URL (OpenAI compat):  {base_url}/v1")
    print(f"Base URL (Anthropic compat): {base_url}/v1/messages")
    print(f"Base URL (raw root):        {base_url}")
    print(f"")
    print(f"API Key (Authorization: Bearer ...):")
    print(f"  sk-{token_key}")
    print(f"")
    print(f"Token row id: {token_id}")
    print(f"Owner user id: {owner_id} ({owner_email})")
    print(f"Quota model: unlimited (consumes the owner-account quota)")
    print(f"Expires: never")
    print(f"Manage / revoke from: {base_url}/keys")
    print("=" * 60)


if __name__ == "__main__":
    main()