#!/usr/bin/env python3
"""把邀请奖励设为 500000（页面上的 $1），并补记当时已经注册但没发奖的 zhenjiang。

只在 wanyan.aff_count 仍为 0 时补记，避免重复加奖。
"""
import json
import sqlite3
import urllib.request

BASE = "http://127.0.0.1:3000"
REWARD = 500000


def req(method, path, body=None, headers=None, timeout=30):
    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)
    with urllib.request.urlopen(request, timeout=timeout) as resp:
        return json.loads(resp.read().decode())


login = req(
    "POST",
    "/api/user/login",
    {"username": "admin", "password": "L4o4986784$"},
)
auth = {
    "Authorization": "Bearer " + login["data"]["access_token"],
    "New-Api-User": "1",
}
for key in ("QuotaForInviter", "QuotaForInvitee"):
    result = req("PUT", "/api/option/", {"key": key, "value": str(REWARD)}, auth)
    print(key, result.get("success"), result.get("message"))

con = sqlite3.connect("/opt/new-api/data/one-api.db")
con.execute("BEGIN")
inviter = con.execute(
    "SELECT aff_count FROM users WHERE id=2 AND username='wanyan'"
).fetchone()
invitee = con.execute(
    "SELECT inviter_id FROM users WHERE id=3 AND username='zhenjiang'"
).fetchone()
if inviter and inviter[0] == 0 and invitee and invitee[0] == 2:
    con.execute(
        "UPDATE users SET aff_count=aff_count+1, aff_quota=aff_quota+?, aff_history=aff_history+? WHERE id=2",
        (REWARD, REWARD),
    )
    con.execute("UPDATE users SET quota=quota+? WHERE id=3", (REWARD,))
    print("backfilled zhenjiang -> wanyan")
else:
    print("skip backfill", inviter, invitee)
con.commit()