#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Tushare Pro 接口自检脚本
=========================
检查当前 token 对常用接口的可用性，避免改代码时用到不存在的接口名/无权限接口。

用法（在 quantitative_data/ 目录下运行，自动读取 .env 中的 token）:
    /opt/quant-venv/bin/python tushare_api_check.py
    /opt/quant-venv/bin/python tushare_api_check.py --token 你的token

输出说明:
    ✅ 可用   — 接口存在且有权限
    ❌ 40101  — 接口名不存在（如 daily_vip），别用这个接口名
    ❌ 40203  — 接口存在但当前 token 无权限（积分不够）
    ❌ 其他   — 见错误码提示
"""
import argparse
import json
import sys
import urllib.request

API_URL = "http://api.tushare.pro"

# 自检清单：接口名 -> 测试参数（用最小代价的参数）
CHECK_LIST = {
    "stock_basic": {},
    "trade_cal": {"exchange": "SSE", "start_date": "20260101", "end_date": "20260131"},
    "daily": {"trade_date": "20260105"},
    "daily_vip": {"trade_date": "20260105"},
    "daily_basic": {"trade_date": "20260105"},
    "daily_basic_vip": {"trade_date": "20260105"},
    "adj_factor": {"trade_date": "20260105"},
    "adj_factor_vip": {"trade_date": "20260105"},
    "moneyflow": {"trade_date": "20260105"},
    "index_daily": {"ts_code": "000001.SH", "start_date": "20260105", "end_date": "20260105"},
    "income": {"ts_code": "000001.SZ", "period": "20251231"},
    "income_vip": {"ts_code": "000001.SZ", "period": "20251231"},
    "balancesheet": {"ts_code": "000001.SZ", "period": "20251231"},
    "balancesheet_vip": {"ts_code": "000001.SZ", "period": "20251231"},
    "cashflow": {"ts_code": "000001.SZ", "period": "20251231"},
    "cashflow_vip": {"ts_code": "000001.SZ", "period": "20251231"},
    "fina_indicator": {"ts_code": "000001.SZ", "period": "20251231"},
    "fina_indicator_vip": {"ts_code": "000001.SZ", "period": "20251231"},
}


def get_token(args_token):
    if args_token:
        return args_token.strip()
    try:
        from config import TUSHARE_TOKEN
        if TUSHARE_TOKEN:
            return TUSHARE_TOKEN.strip()
    except Exception:
        pass
    print("未找到 token：请加 --token 参数，或确认 .env 里 TUSHARE_TOKEN 已配置")
    sys.exit(1)


def call_api(token, api_name, params):
    body = json.dumps({
        "api_name": api_name,
        "token": token,
        "params": params,
        "fields": "",
    }).encode("utf-8")
    req = urllib.request.Request(API_URL, data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read().decode("utf-8"))


def main():
    parser = argparse.ArgumentParser(description="Tushare 接口自检")
    parser.add_argument("--token", default=None, help="Tushare token（默认读 .env）")
    args = parser.parse_args()

    token = get_token(args.token)
    print("=" * 64)
    print(f"Tushare 接口自检  token前8位: {token[:8]}")
    print("=" * 64)

    ok, fail = 0, 0
    for api_name, params in CHECK_LIST.items():
        try:
            result = call_api(token, api_name, params)
            code = result.get("code", -1)
            msg = result.get("msg", "")
            if code == 0:
                print(f"  ✅ {api_name:24s} 可用")
                ok += 1
            else:
                print(f"  ❌ {api_name:24s} code={code} {msg}")
                fail += 1
        except Exception as e:
            print(f"  ❌ {api_name:24s} 请求异常: {e}")
            fail += 1

    print("=" * 64)
    print(f"结果: {ok} 可用, {fail} 不可用")
    if fail:
        print("提示: 40101=接口名不存在(别用这个名)  40203=无权限(积分不够)")
    sys.exit(1 if fail else 0)


if __name__ == "__main__":
    main()