#!/usr/bin/env python3
"""A股公开行情取数 CLI —— 多源自动切换（东方财富 / 腾讯 / 新浪），免 API Key、免登录。

子命令：
  kline     单只标的 K 线（日/周/月，前复权或后复权；分钟线视数据源支持情况）
  snapshot  全市场快照（涨跌幅/换手/PE/PB/市值，含北交所）
  quote     单只实时行情
  sectors   行业板块列表（依赖东方财富，接口不可用时降级提示）
  members   板块成分股（同上）

设计要点：
  - 三家接口按优先级自动降级，任一家被限流/断连不影响整体可用性
  - 输出单位统一：成交量=手、成交额=元、总市值=元、流通市值=元
  - 结果缓存于 <skill>/cache/，盘中重复调用不会重复打接口
  - 带 --out 时写 CSV（UTF-8-SIG，Excel 可直接打开）
"""
from __future__ import annotations

import argparse
import contextlib
import hashlib
import io
import json
import re
import sys
import time
from pathlib import Path

import pandas as pd
import requests

ROOT = Path(__file__).resolve().parent.parent
CACHE = ROOT / "cache"
CACHE.mkdir(parents=True, exist_ok=True)

UA = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                  "(KHTML, like Gecko) Chrome/131.0 Safari/537.36",
    "Referer": "https://finance.sina.com.cn/",
}
EM_KLINE = "https://push2his.eastmoney.com/api/qt/stock/kline/get"
EM_LIST = "https://push2.eastmoney.com/api/qt/clist/get"
TX_KLINE = "https://web.ifzq.gtimg.cn/appstock/app/fqkline/get"
TX_QUOTE = "https://qt.gtimg.cn/q="
SINA_KLINE = "https://quotes.sina.cn/cn/api/json_v2.php/CN_MarketDataService.getKLineData"
SINA_LIST = ("https://vip.stock.finance.sina.com.cn/quotes_service/api/json_v2.php/"
             "Market_Center.getHQNodeData")

KLT = {"1m": 1, "5m": 5, "15m": 15, "30m": 30, "60m": 60,
       "day": 101, "week": 102, "month": 103}
TX_PERIOD = {"day": "day", "week": "week", "month": "month",
             "5m": "m5", "15m": "m15", "30m": "m30", "60m": "m60"}
FQT = {"none": 0, "qfq": 1, "hfq": 2}

FS_ALL = "m:0+t:6,m:0+t:80,m:1+t:2,m:1+t:23"
FS_BJ = "m:0+t:81+s:2048"

SNAP_FIELDS = {
    "f12": "code", "f14": "name", "f2": "price", "f3": "chg_pct", "f4": "chg",
    "f5": "volume", "f6": "amount", "f7": "amplitude", "f8": "turnover",
    "f9": "pe_dyn", "f10": "vol_ratio", "f15": "high", "f16": "low",
    "f17": "open", "f18": "pre_close", "f20": "mcap", "f21": "float_mcap",
    "f23": "pb", "f26": "list_date", "f37": "roe", "f41": "rev_yoy",
    "f45": "net_profit", "f46": "profit_yoy", "f62": "main_inflow",
    "f100": "industry", "f115": "pe_ttm",
}

_TRIED: list[str] = []


def log(msg: str):
    print(f"[aq] {msg}", file=sys.stderr)


def _get(url: str, params: dict | None = None, retries: int = 2,
         timeout: int = 15, encoding: str | None = None):
    last = None
    for i in range(retries):
        try:
            r = requests.get(url, params=params, headers=UA, timeout=timeout)
            r.raise_for_status()
            if encoding:
                r.encoding = encoding
            return r
        except Exception as e:  # noqa: BLE001
            last = e
            time.sleep(0.4 * (i + 1))
    raise RuntimeError(str(last))


def _cache_path(key: str) -> Path:
    return CACHE / (hashlib.md5(key.encode("utf-8")).hexdigest()[:16] + ".json")


def _cached(key: str, ttl: int, fn):
    p = _cache_path(key)
    if p.exists() and (time.time() - p.stat().st_mtime) < ttl:
        return json.loads(p.read_text("utf-8"))
    data = fn()
    p.write_text(json.dumps(data, ensure_ascii=False, default=str), "utf-8")
    return data


# --------------------------------------------------------------------- 代码换算
def norm(code: str) -> tuple[str, str]:
    """返回 (纯6位代码, 市场前缀 sh/sz/bj)。"""
    c = re.sub(r"[^0-9]", "", str(code))
    if len(c) != 6:
        raise ValueError(f"非法 A 股代码: {code}")
    if c[0] in "69":
        return c, "sh"
    if c[0] in "03":
        return c, "sz"
    if c[0] == "4" or c.startswith(("82", "83", "87", "88", "92")):
        return c, "bj"
    return c, "sz"


def em_secid(code: str, kind: str = "stock") -> str:
    if "." in str(code):
        return str(code)
    c, mk = norm(code)
    if kind == "index":
        return ("1." if c.startswith("000") else "0.") + c
    return ("1." if mk == "sh" else "0.") + c


def tx_symbol(code: str, kind: str = "stock") -> str:
    if kind == "index":
        c = re.sub(r"[^0-9]", "", str(code))
        return ("sh" if c.startswith("000") else "sz") + c
    c, mk = norm(code)
    return mk + c


def sina_symbol(code: str, kind: str = "stock") -> str:
    return tx_symbol(code, kind)


# ------------------------------------------------------------------------ K 线
def _kline_em(code, period, adjust, limit, kind) -> pd.DataFrame:
    params = {
        "secid": em_secid(code, kind), "fields1": "f1,f2,f3,f4,f5,f6",
        "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
        "klt": KLT[period], "fqt": FQT[adjust],
        "beg": "0", "end": "20500101", "lmt": str(limit),
    }
    js = _get(EM_KLINE, params).json()
    data = js.get("data") or {}
    rows = data.get("klines") or []
    if not rows:
        raise RuntimeError("东财无数据")
    df = pd.DataFrame([r.split(",") for r in rows],
                      columns=["date", "open", "close", "high", "low", "volume",
                               "amount", "amplitude", "chg_pct", "chg", "turnover"])
    for c in df.columns[1:]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df.insert(0, "code", data.get("code", code))
    df.insert(1, "name", data.get("name", ""))
    df["src"] = "eastmoney"
    return df


def _kline_tx(code, period, adjust, limit, kind) -> pd.DataFrame:
    if period not in TX_PERIOD:
        raise RuntimeError("腾讯不支持该周期")
    sym = tx_symbol(code, kind)
    fq = "" if adjust == "none" else adjust
    js = _get(TX_KLINE, {"param": f"{sym},{TX_PERIOD[period]},,,{min(int(limit), 640)},{fq}"}).json()
    d = (js.get("data") or {}).get(sym) or {}
    key = None
    for k in (f"{fq}{TX_PERIOD[period]}", f"{fq}day", TX_PERIOD[period], "day"):
        if k in d:
            key = k
            break
    if not key:
        cands = [k for k, v in d.items() if isinstance(v, list) and v]
        if not cands:
            raise RuntimeError("腾讯无数据")
        key = cands[0]
    rows = d[key]
    df = pd.DataFrame([r[:6] for r in rows],
                      columns=["date", "open", "close", "high", "low", "volume"])
    for c in df.columns[1:]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df = df[df["date"].astype(str).str.match(r"^\d{4}-\d{2}-\d{2}")]
    df.insert(0, "code", re.sub(r"[^0-9]", "", str(code)))
    df.insert(1, "name", "")
    df["amount"] = pd.NA
    df["src"] = "tencent"
    return df.reset_index(drop=True)


_TF_LOCAL = __import__("threading").local()


def _tf_client():
    """TickFlow 免费服务（无需注册）：仅历史日/周/月K，含成交额，盘中不更新。

    客户端按线程隔离——SDK 实例不是线程安全的，共用一个会被内部锁串行化，
    多线程等于退化成单线程，批量取数会慢十倍。
    """
    c = getattr(_TF_LOCAL, "client", None)
    if c is None:
        with contextlib.redirect_stdout(io.StringIO()):
            from tickflow import TickFlow
            c = TickFlow.free()
        _TF_LOCAL.client = c
    return c


def _kline_tf(code, period, adjust, limit, kind) -> pd.DataFrame:
    p = {"day": "1d", "week": "1w", "month": "1M"}.get(period)
    if not p:
        raise RuntimeError("TickFlow 免费服务不含该周期")
    if adjust == "hfq":
        raise RuntimeError("TickFlow 免费服务仅提供前复权")
    c, mk = norm(code)
    sym = f"{c}.{mk.upper()}"
    with contextlib.redirect_stdout(io.StringIO()):
        raw = _tf_client().klines.get(sym, period=p, count=min(int(limit), 10000),
                                      as_dataframe=True)
    if raw is None or len(raw) == 0:
        raise RuntimeError("TickFlow 无数据")
    df = raw.rename(columns={"trade_date": "date"}).copy()
    df["date"] = df["date"].astype(str).str[:10]
    for col in ["open", "high", "low", "close", "volume", "amount"]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    name = str(raw["name"].iloc[0]) if "name" in raw.columns else ""
    return pd.DataFrame({
        "code": c, "name": name, "date": df["date"],
        "open": df["open"], "close": df["close"],
        "high": df["high"], "low": df["low"],
        "volume": df["volume"], "amount": df["amount"], "src": "tickflow",
    })


def _kline_sina(code, period, adjust, limit, kind) -> pd.DataFrame:
    scale = {"day": 240, "week": 1200, "month": 7200}.get(period)
    if not scale:
        raise RuntimeError("新浪仅支持日/周/月")
    js = _get(SINA_KLINE, {"symbol": sina_symbol(code, kind), "scale": scale,
                           "ma": "no", "datalen": limit}).json()
    if not js:
        raise RuntimeError("新浪无数据")
    df = pd.DataFrame(js)
    df = df.rename(columns={"day": "date"})
    for c in ["open", "high", "low", "close", "volume"]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df["volume"] = df["volume"] / 100.0          # 股 -> 手
    df.insert(0, "code", re.sub(r"[^0-9]", "", str(code)))
    df.insert(1, "name", "")
    df["amount"] = pd.NA
    df["src"] = "sina"
    return df[["code", "name", "date", "open", "close", "high", "low",
               "volume", "amount", "src"]]


def fetch_kline(code: str, period: str = "day", adjust: str = "qfq",
                limit: int = 500, kind: str = "stock", source: str = "auto") -> pd.DataFrame:
    if source != "auto":
        chain = [source]
    elif period in ("day", "week", "month"):
        # 样本短（≤640 根）优先「盘中最新」走腾讯；样本长优先「够深且干净」走 TickFlow。
        # 东财一律垫底：它的前复权长历史会返回 0/负价与离谱振幅，不可信。
        chain = (["tx", "tf", "em", "sina"] if limit <= 640
                 else ["tf", "tx", "em", "sina"])
    else:
        chain = ["em", "tx"]
    fns = {"em": _kline_em, "tx": _kline_tx, "sina": _kline_sina, "tf": _kline_tf}
    errs = []
    for s in chain:
        try:
            df = fns[s](code, period, adjust, limit, kind)
            if df is None or not len(df):
                continue
            # 数据体检：前复权长历史可能出现 0/负价，这类行必须剔除，否则回测必废
            bad = (df["close"] <= 0) | (df["open"] <= 0)
            if bad.any():
                log(f"警告：{s} 源剔除 {int(bad.sum())} 行非正价格（前复权长历史伪值）")
                df = df[~bad]
            if len(df):
                _TRIED.append(s)
                return df.tail(limit).reset_index(drop=True)
        except Exception as e:  # noqa: BLE001
            errs.append(f"{s}: {e}")
    raise RuntimeError("所有 K 线源均失败 -> " + " | ".join(errs))


KLINE_CACHE = ROOT / "cache" / "kline"
KLINE_TTL = 1800


def kline_cached(code: str, period: str = "day", adjust: str = "qfq",
                 limit: int = 300, ttl: int = KLINE_TTL):
    """带磁盘缓存的日K。盘中每轮只重抓快照，K 线在 TTL 内复用，省跨境 RTT。"""
    KLINE_CACHE.mkdir(parents=True, exist_ok=True)
    f = KLINE_CACHE / f"{code}.csv"
    if f.exists() and (time.time() - f.stat().st_mtime) < ttl:
        try:
            df = pd.read_csv(f)
            if len(df) >= 130:
                return df
        except Exception:  # noqa: BLE001
            pass
    try:
        k = fetch_kline(code, period, adjust, limit)
    except Exception:  # noqa: BLE001
        try:
            k = fetch_kline(code, period, adjust, limit, source="tf")
        except Exception:  # noqa: BLE001
            return None
    try:
        k.to_csv(f, index=False, encoding="utf-8-sig")
    except Exception:  # noqa: BLE001
        pass
    return k


# -------------------------------------------------------------------- 全市场快照
def _snap_em(scope: str) -> list[dict]:
    fs = FS_BJ if scope == "bj" else FS_ALL
    out, pn = [], 1
    while True:
        params = {"pn": pn, "pz": 200, "po": 1, "np": 1, "fltt": 2, "invt": 2,
                  "fid": "f3", "fs": fs, "fields": ",".join(SNAP_FIELDS),
                  "ut": "bd1d9ddb04089700cf9c27f6f7426281"}
        js = _get(EM_LIST, params).json()
        d = js.get("data") or {}
        diff = d.get("diff") or []
        if isinstance(diff, dict):
            diff = list(diff.values())
        if not diff:
            break
        out.extend(diff)
        if len(out) >= d.get("total", 0):
            break
        pn += 1
        time.sleep(0.05)
    if not out:
        raise RuntimeError("东财快照为空")
    df = pd.DataFrame(out).rename(columns=SNAP_FIELDS)
    df["code"] = df["code"].astype(str).str.zfill(6)
    df["src"] = "eastmoney"
    return df.to_dict("records")


def _snap_sina(scope: str) -> list[dict]:
    """新浪全市场快照。串行 56 页要 70 秒以上，改成 4 路并发拉分页。"""
    from concurrent.futures import ThreadPoolExecutor

    node = "hs_bjs" if scope == "bj" else "hs_a"

    def page(i: int):
        params = {"page": i, "num": 100, "sort": "symbol", "asc": 1,
                  "node": node, "symbol": "", "_s_r_a": "page"}
        try:
            r = _get(SINA_LIST, params, timeout=20)
            rows = r.json()
            return i, (rows if isinstance(rows, list) else [])
        except Exception:  # noqa: BLE001
            return i, []

    batches: list[list[dict]] = []
    empty_streak = 0
    page_no = 1
    with ThreadPoolExecutor(max_workers=4) as ex:
        while empty_streak < 2 and page_no < 160:
            res = list(ex.map(page, range(page_no, page_no + 4)))
            hit = False
            for _, rows in res:
                if rows:
                    batches.append(rows)
                    hit = True
            empty_streak = 0 if hit else empty_streak + 1
            page_no += 4

    seen, out = set(), []
    for rows in batches:
        for x in rows:
            code = x.get("code")
            if not code or code in seen:
                continue
            seen.add(code)
            try:
                vol = float(x.get("volume") or 0) / 100.0        # 股 -> 手
                amt = float(x.get("amount") or 0)
                pe = float(x.get("per") or 0)
                pbv = float(x.get("pb") or 0)
                out.append({
                    "code": code, "name": x.get("name"),
                    "price": float(x.get("trade") or 0),
                    "chg_pct": float(x.get("changepercent") or 0),
                    "chg": float(x.get("pricechange") or 0),
                    "volume": vol, "amount": amt,
                    "open": float(x.get("open") or 0),
                    "high": float(x.get("high") or 0),
                    "low": float(x.get("low") or 0),
                    "pre_close": float(x.get("settlement") or 0),
                    "turnover": float(x.get("turnoverratio") or 0),
                    "pe_ttm": pe if pe else None,
                    "pb": pbv if pbv else None,
                    "mcap": float(x.get("mktcap") or 0) * 1e4,       # 万元 -> 元
                    "float_mcap": float(x.get("nmc") or 0) * 1e4,
                    "src": "sina",
                })
            except Exception:  # noqa: BLE001
                continue
    if not out:
        raise RuntimeError("新浪快照为空")
    if len(out) < 4000:
        log(f"警告：新浪快照只取到 {len(out)} 只，可能被限流")
    return out


def fetch_snapshot(scope: str = "all", ttl: int = 60, source: str = "auto") -> pd.DataFrame:
    chain = ([source] if source != "auto" else ["em", "sina"])
    def _run(s):
        return _cached(f"snap|{s}|{scope}", ttl,
                       lambda: _snap_em(scope) if s == "em" else _snap_sina(scope))
    errs = []
    for s in chain:
        try:
            recs = _run(s)
            if recs:
                df = pd.DataFrame(recs)
                for c in df.columns:
                    if c not in ("code", "name", "industry", "src", "list_date"):
                        df[c] = pd.to_numeric(df[c], errors="coerce")
                _TRIED.append(f"snap:{s}")
                return df
        except Exception as e:  # noqa: BLE001
            errs.append(f"{s}: {e}")
    raise RuntimeError("所有快照源均失败 -> " + " | ".join(errs))


# ------------------------------------------------------------------------- 实时
_TX_IDX = {"name": 1, "code": 2, "price": 3, "pre_close": 4, "open": 5,
           "volume": 6, "time": 30, "chg": 31, "chg_pct": 32, "high": 33,
           "low": 34, "amount_wan": 37, "turnover": 38, "pe": 39,
           "amplitude": 43, "float_mcap_yi": 44, "mcap_yi": 45, "pb": 46,
           "vol_ratio": 49}


def fetch_quote(code: str, kind: str = "stock") -> dict:
    sym = tx_symbol(code, kind)
    r = _get(TX_QUOTE + sym, encoding="gbk")
    m = re.search(r'="([^"]*)"', r.text)
    if not m:
        raise RuntimeError("腾讯行情解析失败")
    parts = m.group(1).split("~")
    out = {"src": "tencent", "symbol": sym}
    for k, i in _TX_IDX.items():
        if i < len(parts):
            v = parts[i]
            if k in ("name", "time", "code"):
                out[k] = v
            else:
                try:
                    out[k] = float(v)
                except (TypeError, ValueError):
                    out[k] = None
    if out.get("amount_wan") is not None:
        out["amount"] = out["amount_wan"] * 1e4
    if out.get("mcap_yi") is not None:
        out["mcap"] = out["mcap_yi"] * 1e8
    if out.get("float_mcap_yi") is not None:
        out["float_mcap"] = out["float_mcap_yi"] * 1e8
    return out


# ------------------------------------------------------------------------- 板块
def fetch_sectors(kind: str = "industry", ttl: int = 120) -> pd.DataFrame:
    fs = "m:90+t:2" if kind == "industry" else "m:90+t:3"
    fields = {"f12": "code", "f14": "name", "f3": "chg_pct", "f6": "amount",
              "f8": "turnover", "f62": "main_inflow", "f104": "up_count",
              "f105": "down_count", "f128": "leader", "f136": "leader_chg"}

    def _pull():
        params = {"pn": 1, "pz": 500, "po": 1, "np": 1, "fltt": 2, "invt": 2,
                  "fid": "f3", "fs": fs, "fields": ",".join(fields),
                  "ut": "bd1d9ddb04089700cf9c27f6f7426281"}
        js = _get(EM_LIST, params).json()
        d = js.get("data") or {}
        diff = d.get("diff") or []
        if isinstance(diff, dict):
            diff = list(diff.values())
        return diff

    recs = _cached(f"sector|{kind}", ttl, _pull)
    if not recs:
        raise RuntimeError("板块接口不可用（东财 push2 可能被限流），可改用 snapshot 后按行业代码自行聚合")
    df = pd.DataFrame(recs).rename(columns=fields)
    for c in df.columns:
        if c not in ("code", "name", "leader"):
            df[c] = pd.to_numeric(df[c], errors="coerce")
    return df


def fetch_members(board: str, ttl: int = 300) -> pd.DataFrame:
    def _pull():
        params = {"pn": 1, "pz": 500, "po": 1, "np": 1, "fltt": 2, "invt": 2,
                  "fid": "f3", "fs": f"b:{board}", "fields": ",".join(SNAP_FIELDS),
                  "ut": "bd1d9ddb04089700cf9c27f6f7426281"}
        js = _get(EM_LIST, params).json()
        d = js.get("data") or {}
        diff = d.get("diff") or []
        return list(diff.values()) if isinstance(diff, dict) else diff

    recs = _cached(f"members|{board}", ttl, _pull)
    if not recs:
        raise RuntimeError(f"板块 {board} 无成分数据")
    df = pd.DataFrame(recs).rename(columns=SNAP_FIELDS)
    df["code"] = df["code"].astype(str).str.zfill(6)
    return df


# --------------------------------------------------------------------------- CLI
def _emit(df: pd.DataFrame, out: str | None, head: int):
    if out:
        df.to_csv(out, index=False, encoding="utf-8-sig")
        print(json.dumps({"saved": out, "rows": int(len(df)),
                          "columns": list(df.columns)}, ensure_ascii=False))
        return
    with pd.option_context("display.width", 200, "display.max_columns", 50):
        print(df.head(head).to_string(index=False, max_colwidth=16))
    if len(df) > head:
        print(f"... 共 {len(df)} 行（--out 落盘看全量）")


def main(argv=None):
    ap = argparse.ArgumentParser(description="A股公开数据 CLI（多源自动降级）")
    sub = ap.add_subparsers(dest="cmd", required=True)

    k = sub.add_parser("kline")
    k.add_argument("code")
    k.add_argument("--period", default="day", choices=list(KLT))
    k.add_argument("--adjust", default="qfq", choices=list(FQT))
    k.add_argument("--limit", type=int, default=500)
    k.add_argument("--kind", default="stock", choices=["stock", "index"])
    k.add_argument("--source", default="auto", choices=["auto", "em", "tx", "sina", "tf"])
    k.add_argument("--out")
    k.add_argument("--head", type=int, default=10)

    s = sub.add_parser("snapshot")
    s.add_argument("--scope", default="all", choices=["all", "bj"])
    s.add_argument("--ttl", type=int, default=60)
    s.add_argument("--source", default="auto", choices=["auto", "em", "sina"])
    s.add_argument("--out")
    s.add_argument("--head", type=int, default=15)

    q = sub.add_parser("quote")
    q.add_argument("code")
    q.add_argument("--kind", default="stock", choices=["stock", "index"])

    sec = sub.add_parser("sectors")
    sec.add_argument("--kind", default="industry", choices=["industry", "concept"])
    sec.add_argument("--out")
    sec.add_argument("--head", type=int, default=20)

    m = sub.add_parser("members")
    m.add_argument("board")
    m.add_argument("--out")
    m.add_argument("--head", type=int, default=20)

    a = ap.parse_args(argv)

    if a.cmd == "kline":
        _emit(fetch_kline(a.code, a.period, a.adjust, a.limit, a.kind, a.source), a.out, a.head)
    elif a.cmd == "snapshot":
        _emit(fetch_snapshot(a.scope, a.ttl, a.source), a.out, a.head)
    elif a.cmd == "quote":
        print(json.dumps(fetch_quote(a.code, a.kind), ensure_ascii=False, indent=2))
    elif a.cmd == "sectors":
        _emit(fetch_sectors(a.kind), a.out, a.head)
    elif a.cmd == "members":
        _emit(fetch_members(a.board), a.out, a.head)


if __name__ == "__main__":
    sys.exit(main())