qihuo/market.py

"""
行情拉取：默认新浪（免费，普通用户可用）。
同花顺 iFinD HTTP 仅面向机构用户，需单独申请 token，可选开启。
"""
import os
import time
import json
import logging
from typing import Optional

import requests

logger = logging.getLogger(__name__)

THS_TOKEN_URL = "https://quantapi.51ifind.com/api/v1/get_access_token"
THS_QUOTE_URL = "https://quantapi.51ifind.com/api/v1/real_time_quotation"

# iFinD HTTP 期货交易所后缀
THS_EX_SUFFIX = {
    "SHFE": "SHFE",
    "DCE": "DCE",
    "CZCE": "CZCE",
    "CFFEX": "CFFEX",
    "INE": "INE",
}

_token_cache: dict = {"token": "", "expires": 0.0, "refresh": ""}


def _quote_source() -> str:
    return os.getenv("QUOTE_SOURCE", "sina").strip().lower()


def _has_ths_token() -> bool:
    return bool(_get_refresh_token())


def get_quote_source_label() -> str:
    """界面展示用行情源说明。"""
    source = _quote_source()
    if source == "sina":
        return "新浪（免费）"
    if source == "ths":
        return "同花顺 iFinD" if _has_ths_token() else "同花顺（未配置 token，无法使用）"
    if _has_ths_token():
        return "同花顺优先，失败回退新浪"
    return "新浪（免费）"


def _sina_headers() -> dict:
    return {"Referer": "https://finance.sina.com.cn"}


def _fetch_sina_raw(sina_code: str) -> Optional[dict]:
    try:
        url = f"https://hq.sinajs.cn/list={sina_code}"
        resp = requests.get(url, headers=_sina_headers(), timeout=5)
        resp.encoding = "gbk"
        if '"' not in resp.text:
            return None
        body = resp.text.split('"')[1]
        if not body:
            return None
        parts = body.split(",")
        if len(parts) < 9:
            return None
        price = float(parts[8])
        volume = float(parts[14]) if len(parts) > 14 and parts[14] else 0
        return {"name": parts[0], "price": price, "volume": volume}
    except Exception as exc:
        logger.debug("sina fetch failed %s: %s", sina_code, exc)
        return None


def get_sina_price(sina_code: str) -> Optional[float]:
    raw = _fetch_sina_raw(sina_code)
    return raw["price"] if raw else None


_runtime_refresh_token: str = ""


def set_ths_refresh_token(token: str):
    global _runtime_refresh_token
    _runtime_refresh_token = (token or "").strip()


def _get_refresh_token() -> str:
    if _runtime_refresh_token:
        return _runtime_refresh_token
    return os.getenv("THS_REFRESH_TOKEN", "").strip()


def _get_ths_access_token(refresh_token: str) -> Optional[str]:
    if not refresh_token:
        return None
    now = time.time()
    if (
        _token_cache["token"]
        and _token_cache["refresh"] == refresh_token
        and now < _token_cache["expires"]
    ):
        return _token_cache["token"]
    try:
        resp = requests.post(
            THS_TOKEN_URL,
            headers={"Content-Type": "application/json", "refresh_token": refresh_token},
            timeout=10,
        )
        data = resp.json()
        if data.get("errorcode") != 0:
            logger.warning("THS token error: %s", data.get("errmsg"))
            return None
        access = data["data"]["access_token"]
        _token_cache.update({
            "token": access,
            "refresh": refresh_token,
            "expires": now + 3600 * 6,
        })
        return access
    except Exception as exc:
        logger.warning("THS token request failed: %s", exc)
        return None


def _parse_ths_quote(data: dict) -> Optional[float]:
    """从同花顺实时行情响应解析最新价。"""
    try:
        tables = data.get("tables") or []
        for table in tables:
            t = table.get("table") or {}
            for key in ("latest", "new", "close", "trade", "last"):
                val = t.get(key)
                if val is None:
                    continue
                if isinstance(val, list) and val:
                    return float(val[0])
                if isinstance(val, (int, float, str)) and str(val):
                    return float(val)
        # 部分响应嵌套在 data 字段
        if "data" in data and isinstance(data["data"], dict):
            return _parse_ths_quote(data["data"])
    except Exception as exc:
        logger.debug("parse ths quote failed: %s", exc)
    return None


def get_ths_price(ths_full_code: str, refresh_token: str = "") -> Optional[float]:
    """ths_full_code 如 ag2608.SHFE、IF2606.CFFEX"""
    token = refresh_token or _get_refresh_token()
    access = _get_ths_access_token(token)
    if not access:
        return None
    try:
        resp = requests.post(
            THS_QUOTE_URL,
            headers={"Content-Type": "application/json", "access_token": access},
            json={"codes": ths_full_code, "indicators": "latest"},
            timeout=10,
        )
        data = resp.json()
        if data.get("errorcode") != 0:
            logger.warning("THS quote error %s: %s", ths_full_code, data.get("errmsg"))
            return None
        return _parse_ths_quote(data)
    except Exception as exc:
        logger.warning("THS quote failed %s: %s", ths_full_code, exc)
        return None


def get_price(market_code: str, sina_fallback: str = "") -> Optional[float]:
    """
    统一取价入口。
    sina_fallback: 新浪代码 nf_AG2608（普通用户默认使用）
    market_code: 同花顺完整代码 ag2608.SHFE（仅机构 token 可用时）
    """
    source = _quote_source()

    # 仅在有 token 且配置为 ths/auto 时才尝试同花顺
    use_ths = source == "ths" or (source == "auto" and _has_ths_token())
    if use_ths and market_code and "." in market_code:
        price = get_ths_price(market_code)
        if price is not None:
            return price
        if source == "ths":
            return None

    if sina_fallback:
        return get_sina_price(sina_fallback)

    if market_code.startswith("nf_") or market_code.startswith("CFF_RE_"):
        return get_sina_price(market_code)

    return None


def fetch_raw_for_volume(sina_code: str) -> Optional[dict]:
    """主力合约扫描用（成交量），走新浪。"""
    return _fetch_sina_raw(sina_code)