#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
OKX 账户只读查询: 余额 + 当前持仓
- 仅使用 GET 只读接口, 绝不写入/下单
- 密钥从同目录 okx_keys.json 读取
- 支持代理: 命令行 --proxy 或环境变量 OKX_PROXY, 如 http://127.0.0.1:1081 或 socks5://127.0.0.1:1082
- 支持自定义入口: --host 默认 https://www.okx.com
- 支持 DNS 劫持绕过: --dns-ip 指定连接 IP (如 TUN fake-ip 198.18.0.0),
  仅影响 TCP 连接地址, Host/SNI/证书校验仍用真实域名
"""
import json, hmac, hashlib, base64, datetime, sys, os, argparse
import subprocess

HERE = os.path.dirname(os.path.abspath(__file__))
KEYS_FILE = os.path.join(HERE, "okx_keys.json")


def load_keys():
    with open(KEYS_FILE, "r", encoding="utf-8") as f:
        return json.load(f)


def sign(secret, timestamp, method, path, body=""):
    pre = timestamp + method.upper() + path + body
    mac = hmac.new(secret.encode("utf-8"), pre.encode("utf-8"), hashlib.sha256)
    return base64.b64encode(mac.digest()).decode()


def build_opener(proxy):
    if proxy:
        handler = urllib.request.ProxyHandler({"http": proxy, "https": proxy})
        return urllib.request.build_opener(handler)
    return urllib.request.build_opener()


def patch_dns(dns_ip):
    """把指定域名的 DNS 解析替换为 dns_ip (绕过 DNS 污染, TUN fake-ip 场景)"""
    import socket
    real_getaddrinfo = socket.getaddrinfo

    def patched(host, port, *args, **kwargs):
        if host == "www.okx.com":
            return real_getaddrinfo(dns_ip, port, *args, **kwargs)
        return real_getaddrinfo(host, port, *args, **kwargs)

    socket.getaddrinfo = patched


def okx_get_curl(path, keys, host, dns_ip=None, proxy=None):
    """用 curl 发请求 (绕过 Cloudflare 对 Python-urllib 的 TLS/UA 拦截)"""
    ts = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
    sig = sign(keys["secret"], ts, "GET", path, "")
    url = host + path
    cmd = ["curl", "-s", "--max-time", "15"]
    if dns_ip:
        domain = host.split("//")[1].split("/")[0]
        cmd += ["--resolve", f"{domain}:443:{dns_ip}"]
    if proxy:
        cmd += ["-x", proxy]
    cmd += [
        "-H", f"OK-ACCESS-KEY:{keys['api_key']}",
        "-H", f"OK-ACCESS-SIGN:{sig}",
        "-H", f"OK-ACCESS-TIMESTAMP:{ts}",
        "-H", f"OK-ACCESS-PASSPHRASE:{keys['passphrase']}",
        "-H", "Content-Type:application/json",
        "-H", "User-Agent:curl/8.5.0",
        url,
    ]
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=20).stdout
        return json.loads(out)
    except Exception as e:
        return {"code": "-1", "msg": str(e)}


def okx_get(path, keys, host, opener):
    ts = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
    sig = sign(keys["secret"], ts, "GET", path, "")
    req = urllib.request.Request(host + path)
    req.add_header("OK-ACCESS-KEY", keys["api_key"])
    req.add_header("OK-ACCESS-SIGN", sig)
    req.add_header("OK-ACCESS-TIMESTAMP", ts)
    req.add_header("OK-ACCESS-PASSPHRASE", keys["passphrase"])
    req.add_header("Content-Type", "application/json")
    try:
        with opener.open(req, timeout=15) as r:
            return json.loads(r.read().decode())
    except Exception as e:
        return {"code": "-1", "msg": str(e)}


def fmt(n, d=4):
    try:
        return f"{float(n):,.{d}f}"
    except Exception:
        return str(n)


def main():
    ap = argparse.ArgumentParser(description="OKX 账户只读查询")
    ap.add_argument("--proxy", default=os.environ.get("OKX_PROXY"),
                    help="代理地址, 如 http://127.0.0.1:1081 或 socks5://127.0.0.1:1082")
    ap.add_argument("--host", default="https://www.okx.com",
                    help="API 入口, 默认 https://www.okx.com (可换 aws.okx.com)")
    ap.add_argument("--dns-ip", default=os.environ.get("OKX_DNS_IP"),
                    help="绕过 DNS 污染: 强制用该 IP 连接 www.okx.com (如 TUN fake-ip 198.18.0.0)")
    args = ap.parse_args()

    keys = load_keys()

    bal = okx_get_curl("/api/v5/account/balance", keys, args.host, args.dns_ip, args.proxy)
    if bal.get("code") != "0":
        print("余额查询失败:", bal.get("msg", bal))
        return

    pos = okx_get_curl("/api/v5/account/positions", keys, args.host, args.dns_ip, args.proxy)
    if pos.get("code") != "0":
        print("持仓查询失败:", pos.get("msg", pos))
        return

    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    d = bal["data"][0]

    # 从 details 汇总各币种余额
    details_bal = d.get("details", [])
    rows = []
    total_avail = total_frozen = total_upl_bal = 0.0
    for c in details_bal:
        ccy = c.get("ccy", "")
        avail = float(c.get("availEq") or c.get("availBal") or 0)
        frozen = float(c.get("frozenBal") or 0)
        ord_froz = float(c.get("ordFrozen") or 0)
        upl_c = float(c.get("upl") or 0)
        if ccy:
            rows.append((ccy, avail, frozen, ord_froz, upl_c))
        total_avail += avail
        total_frozen += frozen
        total_upl_bal += upl_c

    print("=" * 78)
    print(f"OKX 账户快照  {now}")
    print(f"总权益 totalEq : {fmt(d.get('totalEq'), 2)} USDT")
    print(f"可用余额合计   : {total_avail:,.2f} USDT")
    print(f"冻结资金合计   : {total_frozen:,.2f} USDT (其中挂单冻结 {sum(r[3] for r in rows):,.2f})")
    print(f"未实现盈亏 UPL : {total_upl_bal:+,.2f} USDT")
    print("-" * 78)
    print(f"{'币种':<8}{'可用余额':>16}{'冻结':>16}{'挂单冻结':>16}")
    for ccy, avail, frozen, ord_froz, _ in rows:
        if avail or frozen:
            print(f"{ccy:<8}{avail:>16,.6f}{frozen:>16,.6f}{ord_froz:>16,.6f}")
    print("=" * 78)

    details = pos.get("data", [])
    if not details:
        print("当前无持仓")
        return

    print(f"{'合约':<14}{'方向':<6}{'数量':>12}{'开仓均价':>14}{'标记价':>14}{'保证金':>12}{'杠杆':>6}{'强平价':>14}{'浮动盈亏':>14}")
    print("-" * 78)
    total_upl = 0.0
    for p in details:
        inst = p.get("instId", "")
        pos_side = p.get("posSide", "")
        sz = float(p.get("pos", 0) or 0)
        avg_px = float(p.get("avgPx", 0) or 0)
        mark_px = float(p.get("markPx", 0) or 0)
        margin = float(p.get("margin", 0) or 0)
        lev = p.get("lever", "")
        liq = float(p.get("liqPx", 0) or 0)
        upl = float(p.get("upl", 0) or 0)
        total_upl += upl
        side_cn = {"long": "多", "short": "空", "net": "净"}.get(pos_side, pos_side)
        print(f"{inst:<14}{side_cn:<6}{fmt(sz):>12}{fmt(avg_px):>14}{fmt(mark_px):>14}"
              f"{fmt(margin, 2):>12}{lev:>6}{fmt(liq):>14}{fmt(upl, 2):>14}")
    print("-" * 78)
    print(f"持仓合计浮动盈亏: {total_upl:+,.2f} USDT")


if __name__ == "__main__":
    main()