#!/usr/bin/env python3
"""TRAE (ByteDance, CN edition) daily check-in client.

Endpoints (reverse-engineered from the TRAE desktop client, see community
projects wb-trae-checkin / Trae-AutoCheckin):

  base https://api.trae.cn
  status  POST /trae/api/v2/ug/checkin_credits/status   body {"req_source":1}
  claim   POST /trae/api/v2/ug/checkin_credits/claim    body {"req_source":1}

Auth: Authorization: Cloud-IDE-JWT <accessToken> (a JWT issued by the client).
The server also expects desktop-client headers including a 16-digit device id.

Result codes:
  0     success (claim returns data.points, usually 200)
  9095  already checked in today
  9074  device id remembered by the server -> rotate to a fresh device id
  1001/1002/401/403  session invalid -> re-login required
"""
import base64
import hashlib
import json
import time

from common import http_post_json, log

API_BASE = "https://api.trae.cn"
TRAE_VERSION = "1.107.1"
EP_STATUS = "/trae/api/v2/ug/checkin_credits/status"
EP_CLAIM = "/trae/api/v2/ug/checkin_credits/claim"

AUTH_FAIL_CODES = {1001, 1002, 401, 403}


def _b64url_decode(seg):
    seg += "=" * (-len(seg) % 4)
    return base64.urlsafe_b64decode(seg.encode("ascii"))


def jwt_payload(token):
    """Decode a JWT payload without verifying the signature."""
    try:
        parts = token.split(".")
        if len(parts) < 2:
            return {}
        return json.loads(_b64url_decode(parts[1]).decode("utf-8"))
    except Exception:
        return {}


def user_region(account):
    auth = account.get("auth", {})
    if auth.get("region"):
        return auth["region"]
    payload = jwt_payload(auth.get("accessToken", ""))
    region = (payload.get("userRegion") or {}).get("region")
    return region or "CN"


# ---------------------------------------------------------------------------
# Device id (16 digits): stable per account, rotated on 9074
# ---------------------------------------------------------------------------

def stable_device_id(account_id):
    digest = hashlib.sha256(("agentcheckin:%s" % account_id).encode("utf-8")).hexdigest()
    n = int(digest[:13], 16)
    return str(n).ljust(16, "0")[:16]


def rotate_device_id():
    return (str(int(time.time() * 1000))[-13:] + str(int(time.time() * 1000) % 1000).zfill(3))[:16]


def device_id(account):
    did = account.get("device_id")
    if not did:
        did = stable_device_id(account.get("id", "default"))
        account["device_id"] = did
    return did


def headers(account):
    auth = account.get("auth", {})
    return {
        "Authorization": "Cloud-IDE-JWT %s" % auth.get("accessToken", ""),
        "X-User-Region": user_region(account),
        "User-Agent": "Trae/%s" % TRAE_VERSION,
        "Content-Type": "application/json",
        "x-device-id": device_id(account),
        "x-device-type": "Windows",
        "x-os-version": "10.0.19045",
        "x-app-version": TRAE_VERSION,
    }


# ---------------------------------------------------------------------------
# Check-in API
# ---------------------------------------------------------------------------

def query_status(account):
    """Return (checked, enable, credits, error). checked True/False/None."""
    data, status, err = http_post_json(API_BASE + EP_STATUS, body={"req_source": 1},
                                       headers=headers(account), timeout=20)
    if err or data is None:
        if status in (401, 403):
            return None, None, None, "no_session"
        return None, None, None, err or "status query failed"
    if isinstance(data, dict):
        code = data.get("code")
        if code in AUTH_FAIL_CODES:
            return None, None, None, "no_session"
        if code == 0:
            return (
                bool(data.get("checked_in")),
                data.get("enable") is not False,
                data.get("credits"),
                None,
            )
    return None, None, None, None


def claim(account):
    """Claim today's credits. Returns (status, message, points).

    status: ok / already / fail / no_session / retry_later
    """
    hdrs = headers(account)
    data, status, err = http_post_json(API_BASE + EP_CLAIM, body={"req_source": 1},
                                       headers=hdrs, timeout=20)
    if status in (401, 403):
        return "no_session", "登录态已失效，请重新获取 Token", 0
    if data is None:
        return "fail", "请求失败（HTTP %s：%s）" % (status, err), 0
    if not isinstance(data, dict):
        return "fail", "非 JSON 返回（HTTP %s）" % status, 0
    code = data.get("code")
    if code == 0:
        d = data.get("data") if isinstance(data.get("data"), dict) else {}
        return "ok", "签到成功", d.get("points") or 200
    if code == 9095:
        return "already", "今天已签到", 0
    if code in AUTH_FAIL_CODES:
        return "no_session", "登录态已失效（code=%s），请重新获取 Token" % code, 0
    if code == 9074:
        return "retry_later", "触发风控（9074），需更换设备号重试", 0
    return "fail", "领取返回 code=%s %s" % (code, str(data.get("message") or data.get("msg") or "")[:80]), 0


def checkin_once(account):
    """Full single-account check-in: status -> claim -> verify.

    Returns dict {status, message, points, credits}.
    """
    checked, enable, credits, err = query_status(account)
    if err == "no_session":
        return {"status": "no_session", "message": "登录态已失效，请重新获取 Token", "points": 0}
    if checked is True:
        return {"status": "already", "message": "今日已签到",
                "points": 0, "credits": credits}
    if enable is False:
        return {"status": "inactive", "message": "当前账号/时段未开放签到活动", "points": 0}

    rotated = False
    for attempt in range(2):
        status, message, points = claim(account)
        if status == "ok":
            # verify server-side
            for _ in range(3):
                time.sleep(2)
                v_checked, _, v_credits, _ = query_status(account)
                if v_checked is True:
                    return {"status": "ok", "message": "签到成功（已回查确认）",
                            "points": points, "credits": v_credits}
                if v_checked is False:
                    break
            return {"status": "ok_unverified",
                    "message": "领取成功（+%s）但回查未确认" % points,
                    "points": points, "credits": credits}
        if status == "already":
            return {"status": "already", "message": "今天已签到", "points": 0,
                    "credits": credits}
        if status == "no_session":
            return {"status": "no_session", "message": message, "points": 0}
        if status == "retry_later" and not rotated:
            account["device_id"] = rotate_device_id()
            rotated = True
            log("trae account %s: 9074 -> rotated device id" % account.get("id"))
            time.sleep(3)
            continue
        return {"status": "fail", "message": message, "points": 0}
    return {"status": "retry_later", "message": "风控限流（9074），换设备号后仍失败，下个时段自动补签",
            "points": 0}


def account_display(account):
    return account.get("name") or (account.get("auth") or {}).get("nickname") or "TRAE 账号"