#!/usr/bin/env python3
"""把令牌 1 改为永不过期、default 分组、无限额度，并重新启用。

更新接口默认不改 status。过期状态要在过期时间改成 -1 之后，
再单独调用 /api/token/?status_only=true。
"""
import json
import sqlite3
import urllib.request

BASE = "http://127.0.0.1:3000"


def req(method, path, body=None, headers=None, timeout=120):
    data = None if body is None else json.dumps(body).encode()
    request = urllib.request.Request(BASE + path, data=data, method=method)
    request.add_header("Content-Type", "application/json")
    for key, value in (headers or {}).items():
        request.add_header(key, value)
    try:
        with urllib.request.urlopen(request, timeout=timeout) as resp:
            raw = resp.read().decode()
            return resp.status, json.loads(raw) if raw else {}
    except urllib.error.HTTPError as exc:
        raw = exc.read().decode()
        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError:
            parsed = {"raw": raw[:500]}
        return exc.code, parsed


_, login = req(
    "POST",
    "/api/user/login",
    {"username": "admin", "password": "L4o4986784$"},
)
auth = {
    "Authorization": "Bearer " + login["data"]["access_token"],
    "New-Api-User": "1",
}
status, updated = req(
    "PUT",
    "/api/token/",
    {
        "id": 1,
        "name": "api-key-v1",
        "status": 1,
        "expired_time": -1,
        "remain_quota": 500000,
        "unlimited_quota": True,
        "model_limits_enabled": False,
        "model_limits": "",
        "allow_ips": "",
        "group": "default",
    },
    auth,
)
print("update", status, json.dumps(updated, ensure_ascii=False)[:400])
status, enabled = req(
    "PUT",
    "/api/token/?status_only=true",
    {"id": 1, "status": 1},
    auth,
)
print("enable", status, json.dumps(enabled, ensure_ascii=False)[:400])

con = sqlite3.connect("/opt/new-api/data/one-api.db")
key = con.execute("SELECT key FROM tokens WHERE id=1").fetchone()[0]
status, chat = req(
    "POST",
    "/v1/chat/completions",
    {
        "model": "qwen3.6-27b",
        "messages": [{"role": "user", "content": "只回复：ok"}],
        "max_tokens": 16,
    },
    {"Authorization": "Bearer sk-" + key},
)
print("chat", status)
if "error" in chat:
    print(json.dumps(chat["error"], ensure_ascii=False)[:400])
else:
    choice = (chat.get("choices") or [{}])[0]
    message = choice.get("message") or {}
    print("finish", choice.get("finish_reason"))
    print("content", repr((message.get("content") or "")[:120]))