market_sync/app/tasks/task_stock_node.py

"""同步任务：股票-指数/行业/概念 映射（mairui `/hszg` 三接口）。

数据源:
  - GET /hszg/list/{licence}        节点树(1464 节点,A 股+港股+基金...)
  - GET /hszg/gg/{node_code}/{licence} 节点→成分股

设计要点:
  1) 每周六 11:30 触发(mairui 11:00 更新后留 30 min buffer)
  2) 覆盖范围:A 股(type1=0) 白名单 type2: 0/1/2/3/4/6/7
     - 0 = 申万一级行业 (31)
     - 1 = 申万二级行业 (131)
     - 2 = 热门概念 (698)
     - 3 = 概念板块 (214)
     - 4 = 地域板块 (31)
     - 6 = 板块分类 沪/深/北 (16)
     - 7 = 指数成分 (43)
     跳过 type2 = 5/8/9/10(证监会行业与 baostock 重复;风险警示/大盘指数/次新股 是单节点)
  3) 用 /hszg/gg 反查路径(1464 调用)比 /hszg/zg 正向(5206 调用)快 5 倍
  4) ThreadPoolExecutor=5 保守起步(钻石 100 RPS 留 buffer)
  5) 单节点失败 warn 跳过,不中断整体
  6) 全量替换写库(set 语义,idempotent)
  7) stock_code 统一 hermes 格式(SH600519),与项目 2026-07-01 约定一致
"""
from __future__ import annotations

import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

from app.core.db import ops as db_ops
from app.core.datasource.base import registry as ds_registry
from app.core.datasource.registry import is_source_ready
from app.core.datasource.utils import to_hermes
from app.core.sync.base import SyncTask
from app.core.sync.registry import mark_sync_blocked
from app.core.utils.logging import get_logger

logger = get_logger("sync.stock_node")


# type1=0 (A 股) 子类映射
TYPE2_TO_CATEGORY = {
    0: ("A股-申万行业", "industry"),
    1: ("A股-申万二级", "industry_sub"),
    2: ("A股-热门概念", "concept"),
    3: ("A股-概念板块", "concept"),
    4: ("A股-地域板块", "region"),
    5: ("A股-证监会行业", "industry_regulator"),  # 跳过,与 baostock 重复
    6: ("A股-分类", "class"),
    7: ("A股-指数成分", "index"),
    8: ("A股-风险警示", "warning"),
    9: ("A股-大盘指数", "index_main"),
    10: ("A股-次新股", "subnew"),
}

# 默认同步的 type2 集合
INCLUDED_TYPE2 = {0, 1, 2, 3, 4, 6, 7}


class SyncStockNode(SyncTask):
    """股票-节点 N×M 映射(指数/行业/概念 统一入口)。

    CLI 用法:
      python -m app.entrypoints.cli sync stock_node                        # 默认全跑
      python -m app.entrypoints.cli sync stock_node --type2 2,3            # 只跑概念
      python -m app.entrypoints.cli sync stock_node --max-workers 10      # 提高并发
    """
    dataset_id = "stock_node"

    def _run(
        self,
        *,
        trigger_source: str = "manual",
        type2: list[int] | str | None = None,
        max_workers: int = 5,
        **kwargs,
    ) -> dict[str, Any]:
        mr = ds_registry.get("datasource_mairui")
        if mr is None:
            return {"status": "error", "message": "mairui 数据源未注册"}
        ok, reason = is_source_ready(mr.key)
        if not ok:
            mark_sync_blocked(self.dataset_id, message=f"mairui 未就绪: {reason}")
            return {"status": "blocked", "message": f"mairui 未就绪: {reason}"}

        # ── 解析 type2 过滤 ──
        if type2 is None:
            included = INCLUDED_TYPE2
        elif isinstance(type2, str):
            raw = [t.strip() for t in type2.split(",") if t.strip()]
            if not raw:
                return {
                    "status": "error",
                    "message": "--type2 不能为空字符串（要限制子集请用 '0,2,3'；要全量不要传 --type2）",
                }
            bad = [t for t in raw if not t.isdigit()]
            if bad:
                valid = ", ".join(f"{k}={v[1]}" for k, v in TYPE2_TO_CATEGORY.items())
                return {
                    "status": "error",
                    "message": f"--type2 含非法 token {bad}；应为 0-10 整数（如 '2,3'）；有效映射: {valid}",
                }
            included = {int(t) for t in raw}
            unknown = {t for t in included if t not in TYPE2_TO_CATEGORY}
            if unknown:
                valid = ", ".join(f"{k}={v[1]}" for k, v in TYPE2_TO_CATEGORY.items())
                return {
                    "status": "error",
                    "message": f"--type2 含未知 mairui type2 {sorted(unknown)}；有效范围 0-10；映射: {valid}",
                }
        else:
            # 程序化调用: 接受 list[int] | set[int] | tuple,但仍校验范围
            # (2026-07-07 防御: 也覆盖 programmatic caller 传非法值的情况)
            included = set(type2)
            bad = [t for t in included if not isinstance(t, int) or t not in TYPE2_TO_CATEGORY]
            if bad:
                valid = ", ".join(f"{k}={v[1]}" for k, v in TYPE2_TO_CATEGORY.items())
                return {
                    "status": "error",
                    "message": f"type2 含非法值 {bad}；应为 0-10 整数；有效映射: {valid}",
                }
        logger.info(f"[stock_node] 同步 type2: {sorted(included)}")

        # ── Phase 1: 拉节点树(1 次调用) ──
        t0 = time.time()
        df_tree = mr.fetch_node_tree()
        if df_tree is None or df_tree.empty:
            return {"status": "error", "message": "mairui /hszg/list 返回空"}

        # type1=0 (A 股) + 白名单 type2 + isleaf=1
        # mairui 的 type2 字段是 float(如 2.0);部分叶子 type2 为 NaN,先 dropna
        df_a = df_tree[df_tree["type1"] == 0].copy()
        df_a = df_a.dropna(subset=["type2"])
        df_a["type2_int"] = df_a["type2"].astype(int)
        df_leaves = df_a[
            (df_a["type2_int"].isin(included))
            & (df_a["isleaf"] == 1)
        ].copy()
        if df_leaves.empty:
            return {"status": "warning", "message": f"无匹配叶子节点(type2={sorted(included)})"}

        df_leaves["type2_int"] = df_leaves["type2"].astype(int)
        df_leaves["category_key"] = "0:" + df_leaves["type2_int"].astype(str)
        df_leaves["category_type"] = df_leaves["type2_int"].map(
            lambda t: TYPE2_TO_CATEGORY.get(int(t), ("?", "other"))[1]
        )
        df_leaves["parent_code"] = df_leaves["pcode"]
        df_leaves["parent_name"] = df_leaves["pname"]
        df_leaves["is_leaf"] = 1

        # 写 node_categories(7 行)
        cat_counter = Counter(df_leaves["category_key"])
        category_rows = []
        for cat_key, count in cat_counter.items():
            t2 = int(cat_key.split(":")[1])
            display_name, cat_type = TYPE2_TO_CATEGORY.get(t2, (cat_key, "other"))
            category_rows.append({
                "category_key": cat_key,
                "display_name": display_name,
                "market": "A 股",
                "category_type": cat_type,
                "node_count": count,
            })
        db_ops.replace_all_node_categories(category_rows)
        logger.info(f"[stock_node] node_categories: {len(category_rows)} 行")

        # 写 nodes(约 1100 行) —— dedup 按 node_code,因为 _bulk_upsert_orm 同 chunk 内
        # 不能有重复 PK(PG ON CONFLICT 同 batch 第二次出现会报错)
        node_rows = [
            {
                "node_code": str(r["code"]),
                "node_name": str(r["name"]),
                "category_key": r["category_key"],
                "parent_code": str(r["pcode"]) if r["pcode"] else None,
                "parent_name": str(r["pname"]) if r["pname"] else None,
                "level": int(r["level"]),
                "is_leaf": 1,
                "mairui_type1": int(r["type1"]),
                "mairui_type2": int(r["type2_int"]),
            }
            for _, r in df_leaves.iterrows()
        ]
        # 去重 node_code(罕见:同一节点可能因 pname 差异重复出现,保留 first)
        seen = set()
        node_rows_dedup = []
        for nr in node_rows:
            if nr["node_code"] in seen:
                continue
            seen.add(nr["node_code"])
            node_rows_dedup.append(nr)
        if len(node_rows) != len(node_rows_dedup):
            logger.info(f"[stock_node] nodes dedup: {len(node_rows)} → {len(node_rows_dedup)}")
        db_ops.replace_all_nodes(node_rows_dedup)
        logger.info(f"[stock_node] nodes(leaves): {len(node_rows)} 行")

        # ── Phase 2: 拉每节点的成分股(gg 路径) ──
        leaf_codes = df_leaves["code"].tolist()
        all_map_rows: list[dict] = []
        ok_cnt = fail_cnt = 0
        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            futures = {
                pool.submit(self._scrape_one, mr, code): code for code in leaf_codes
            }
            for i, fut in enumerate(as_completed(futures), 1):
                node_code = futures[fut]
                try:
                    rows = fut.result()
                    if rows:
                        all_map_rows.extend(rows)
                        ok_cnt += 1
                    else:
                        fail_cnt += 1
                except Exception as e:
                    fail_cnt += 1
                    logger.warning(f"[stock_node {node_code}] {e}")
                if i % 100 == 0 or i == len(leaf_codes):
                    self._progress(
                        message=f"stock_node {i}/{len(leaf_codes)} ok={ok_cnt} fail={fail_cnt} mappings={len(all_map_rows)}",
                        current=i, total=len(leaf_codes), current_step=node_code,
                    )

        db_ops.replace_all_stock_node_map(all_map_rows)
        elapsed = round(time.time() - t0, 1)
        msg = (
            f"节点映射 {len(node_rows)}节点 ok={ok_cnt} fail={fail_cnt} "
            f"共{len(all_map_rows)}映射 {elapsed}s"
        )
        return {
            "status": "ok" if fail_cnt == 0 else "warning",
            "message": msg,
            "nodes": len(node_rows),
            "categories": len(category_rows),
            "ok": ok_cnt, "fail": fail_cnt,
            "mappings": len(all_map_rows),
            "elapsed_sec": elapsed,
        }

    def _scrape_one(self, mr, node_code: str) -> list[dict]:
        """拉单个节点的成分股 → N 条 (stock_code, node_code) 行。

        stock_code 用 hermes 格式(SH600519),与项目其他表一致。
        单只股票可能重复(节点 X、节点 Y 都属于该股)→ 不去重,反正 PK 幂等。
        """
        df = mr.fetch_node_stocks(node_code)
        if df is None or df.empty:
            return []
        rows = []
        for _, r in df.iterrows():
            dm = str(r.get("dm", "") or "").strip()
            if len(dm) != 6 or not dm.isdigit():
                continue
            rows.append({
                "stock_code": to_hermes(dm),
                "node_code": str(node_code),
            })
        return rows