xiaobai-review/xiaobai-datahub/datahub/adapters/tencent.py

from __future__ import annotations

import time
import urllib.request
from datetime import datetime
from typing import Any

from datahub.adapters.base import AdapterError, MarketAdapter
from datahub.numbers import finite_number, round4

TENCENT_INDEX_URL = "https://qt.gtimg.cn/q=sh000001,sz399001,sz399006"
TENCENT_QUOTE_URL = "https://qt.gtimg.cn/q="
TENCENT_QUOTE_BATCH = 80
BROWSER_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36"
)


class TencentAdapter(MarketAdapter):
    name = "tencent"

    def __init__(self, timeout: int = 8) -> None:
        self.timeout = timeout

    def probe(self) -> dict[str, Any]:
        started = time.perf_counter()
        try:
            rows = self.fetch_indices()
            state = "ok" if len(rows) == 3 else "empty"
        except AdapterError as exc:
            return {
                "provider": self.name,
                "configured": True,
                "state": "error",
                "message": str(exc),
                "latency_ms": round((time.perf_counter() - started) * 1000),
            }
        return {
            "provider": self.name,
            "configured": True,
            "state": state,
            "latency_ms": round((time.perf_counter() - started) * 1000),
        }

    def fetch(self, dataset: str, params: dict[str, Any]) -> list[dict[str, Any]]:
        if dataset in {"indexes_quotes", "index_quotes"}:
            return self.fetch_indices()
        if dataset in {"quotes", "quotes_latest"}:
            codes = params.get("codes") or []
            if isinstance(codes, str):
                codes = [item.strip() for item in codes.split(",") if item.strip()]
            return self.fetch_quotes(list(codes))
        if dataset in {"quotes_market", "market_quotes"}:
            codes = params.get("codes") or []
            if isinstance(codes, str):
                codes = [item.strip() for item in codes.split(",") if item.strip()]
            return self.fetch_quotes(list(codes))
        raise AdapterError(f"{self.name} unsupported dataset: {dataset}")

    def normalize(self, dataset: str, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        return list(rows)

    def fetch_indices(self) -> list[dict[str, Any]]:
        request = urllib.request.Request(
            TENCENT_INDEX_URL,
            headers={"User-Agent": BROWSER_UA, "Referer": "https://gu.qq.com/"},
            method="GET",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                raw = response.read().decode("gb18030", errors="ignore")
        except Exception as exc:
            raise AdapterError(f"tencent request failed: {exc}") from exc
        result = []
        for line in raw.splitlines():
            if '="' not in line:
                continue
            fields = line.split('="', 1)[1].rsplit('";', 1)[0].split("~")
            if len(fields) < 38:
                continue
            code = fields[2]
            if code not in {"000001", "399001", "399006"}:
                continue
            try:
                quote_time = datetime.strptime(fields[30], "%Y%m%d%H%M%S").astimezone()
            except ValueError as exc:
                raise AdapterError(f"Tencent invalid quote time for {code}") from exc
            ts_code = "000001.SH" if code == "000001" else f"{code}.SZ"
            result.append(
                {
                    "ts_code": ts_code,
                    "code": code,
                    "name": fields[1] or code,
                    "price": round4(finite_number(fields[3])),
                    "pct_chg": round4(finite_number(fields[32])),
                    "change_amount": round4(finite_number(fields[31])),
                    "open": round4(finite_number(fields[5])),
                    "high": round4(finite_number(fields[33])),
                    "low": round4(finite_number(fields[34])),
                    "previous_close": round4(finite_number(fields[4])),
                    "amount": round4(finite_number(fields[37]) * 10000),
                    "quote_time_epoch": int(quote_time.timestamp()),
                    "quote_time": quote_time.isoformat(timespec="seconds"),
                    "source": "tencent_qt",
                }
            )
        if len(result) != 3:
            raise AdapterError(f"Tencent returned {len(result)}/3 indices")
        return result

    def fetch_quotes(self, codes: list[str]) -> list[dict[str, Any]]:
        symbols: list[str] = []
        seen: set[str] = set()
        for raw in codes:
            symbol = _tencent_symbol(str(raw or ""))
            if not symbol or symbol in seen:
                continue
            seen.add(symbol)
            symbols.append(symbol)
        if not symbols:
            return []
        result: list[dict[str, Any]] = []
        errors: list[str] = []
        for index in range(0, len(symbols), TENCENT_QUOTE_BATCH):
            batch = symbols[index:index + TENCENT_QUOTE_BATCH]
            try:
                raw = self._get_text(f"{TENCENT_QUOTE_URL}{','.join(batch)}")
            except AdapterError as exc:
                errors.append(str(exc))
                continue
            for line in raw.splitlines():
                quote = _parse_tencent_stock_quote(line)
                if quote:
                    result.append(quote)
        if not result:
            detail = f"；{'; '.join(errors[:3])}" if errors else ""
            raise AdapterError(f"Tencent quotes empty{detail}")
        return result

    def _get_text(self, url: str) -> str:
        request = urllib.request.Request(
            url,
            headers={"User-Agent": BROWSER_UA, "Referer": "https://gu.qq.com/"},
            method="GET",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                return response.read().decode("gb18030", errors="ignore")
        except Exception as exc:
            raise AdapterError(f"tencent request failed: {exc}") from exc


def _tencent_symbol(code: str) -> str:
    raw = str(code or "").strip().upper()
    if not raw:
        return ""
    symbol = raw.split(".")[0]
    if not symbol.isdigit() or len(symbol) != 6:
        return ""
    if raw.endswith(".SH") or symbol.startswith(("5", "6", "9")):
        return f"sh{symbol}"
    if raw.endswith(".BJ") or symbol.startswith(("4", "8")):
        return f"bj{symbol}"
    return f"sz{symbol}"


def _parse_tencent_stock_quote(line: str) -> dict[str, Any] | None:
    if '="' not in line:
        return None
    prefix, payload = line.split('="', 1)
    fields = payload.rsplit('";', 1)[0].split("~")
    if len(fields) < 38:
        return None
    symbol = str(fields[2] or "")
    if not symbol.isdigit() or len(symbol) != 6:
        return None
    close = round4(finite_number(fields[3]))
    previous = round4(finite_number(fields[4]))
    if not close or not previous or close <= 0 or previous <= 0:
        return None
    marker = prefix.lower()
    if "sh" in marker:
        ts_code = f"{symbol}.SH"
    elif "bj" in marker:
        ts_code = f"{symbol}.BJ"
    else:
        ts_code = f"{symbol}.SZ"
    quote_stamp = ""
    quote_date = ""
    epoch = 0
    try:
        parsed = datetime.strptime(fields[30], "%Y%m%d%H%M%S")
        quote_date = parsed.strftime("%Y%m%d")
        epoch = int(parsed.timestamp())
        quote_stamp = parsed.astimezone().isoformat(timespec="seconds")
    except ValueError:
        pass
    return {
        "ts_code": ts_code,
        "name": fields[1] or symbol,
        "price": close,
        "close": close,
        "pct_chg": round4(finite_number(fields[32])),
        "change_amount": round4(finite_number(fields[31])),
        "open": round4(finite_number(fields[5])),
        "high": round4(finite_number(fields[33])),
        "low": round4(finite_number(fields[34])),
        "pre_close": previous,
        "previous_close": previous,
        "volume": round4(finite_number(fields[6]) * 100),
        "vol": round4(finite_number(fields[6]) * 100),
        "amount": round4(finite_number(fields[37]) * 10000),
        "quote_date": quote_date,
        "quote_time_epoch": epoch,
        "quote_time": quote_stamp,
        "source": "tencent_qt",
    }