#!/usr/bin/env python3
"""实时OKX账户查询: positions + balance + 关键ticker 三连查

用法: python3 check_account.py [symbols...]
     python3 check_account.py                    # 查所有持仓+USDT
     python3 check_account.py ETH BTC            # 查所有持仓+指定ticker
"""
import json, time, hmac, hashlib, base64, sys, os, requests

def okx(p, params=None, t=10):
    creds = open(os.path.expanduser('~/.bashrc')).read()
    k = s = pw = None
    for line in creds.split('\n'):
        if line.startswith('export OKX_API_KEY='): k = line.split('=', 1)[1].strip().strip('"').strip("'")
        elif line.startswith('export OKX_SECRET='): s = line.split('=', 1)[1].strip().strip('"').strip("'")
        elif line.startswith('export OKX_PASSPHRASE='): pw = line.split('=', 1)[1].strip().strip('"').strip("'")
    ts = time.strftime('%Y-%m-%dT%H:%M:%S.000Z', time.gmtime())
    msg = ts + 'GET' + p + (json.dumps(params) if params else '')
    sig = base64.b64encode(hmac.new(s.encode(), msg.encode(), hashlib.sha256).digest()).decode()
    h = {'OK-ACCESS-KEY': k, 'OK-ACCESS-SIGN': sig, 'OK-ACCESS-TIMESTAMP': ts, 'OK-ACCESS-PASSPHRASE': pw, 'Content-Type': 'application/json'}
    px = {'http': 'http://127.0.0.1:7890', 'https': 'http://127.0.0.1:7890'}
    return requests.get(f'https://www.okx.com{p}', params=params or {}, headers=h, proxies=px, timeout=t).json()

def main():
    extra_symbols = sys.argv[1:]

    # 1. positions
    pos_resp = okx('/api/v5/account/positions')
    positions = []
    for p in pos_resp.get('data', []):
        pos_size = float(p.get('pos', '0') or 0)
        if pos_size != 0:
            positions.append({
                'instId': p['instId'],
                'side': 'long' if pos_size > 0 else 'short',
                'contracts': abs(pos_size),
                'avgPx': float(p.get('avgPx', '0') or 0),
                'markPx': float(p.get('markPx', '0') or 0),
                'upl': float(p.get('upl', '0') or 0),
                'lever': p.get('lever'),
                'liqPx': float(p.get('liqPx', '0') or 0),
                'margin': p.get('margin', ''),
            })

    # 2. balance
    bal_resp = okx('/api/v5/account/balance')
    usdt = {}
    for d in bal_resp['data'][0].get('details', []):
        if d['ccy'] == 'USDT':
            usdt = {
                'availBal': float(d.get('availBal', '0') or 0),
                'frozenBal': float(d.get('frozenBal', '0') or 0),
                'eq': float(d.get('eq', '0') or 0),
            }
            break

    # 3. tickers for held symbols + extras
    tickers = {}
    target_insts = list(set([p['instId'] for p in positions] + extra_symbols))
    for inst in target_insts:
        try:
            tk = okx('/api/v5/market/ticker', {'instId': inst})
            if tk.get('data'):
                tickers[inst] = float(tk['data'][0]['last'])
        except Exception:
            pass

    # 4. 输出
    result = {
        'ts': int(time.time()),
        'usdt': usdt,
        'positions': positions,
        'tickers': tickers,
        'has_position': len(positions) > 0,
    }
    print(json.dumps(result, ensure_ascii=False, indent=2))

if __name__ == '__main__':
    main()