"""
Fetch adjusted daily OHLCV for CSI 300 / CSI 500 ever-members from Yahoo
(through the market-data proxy). Delisted names Yahoo lacks stay missing
(listed in data/research/cn_missing.txt): Eastmoney and Tencent throttle bulk
requests, and Tencent's back-adjusted series damps returns (limit-down days
show -7% instead of -10%), so neither is used.

    uv run python -m research.lab.fetch_cn_yahoo [--hs300-only]

Writes data/research/cn_ohlcv.pkl (fields open/high/low/close/volume/preclose/
pctchg/tradestatus + members + index + per-code "source").
"""

import argparse
import glob
import os
import pickle

import pandas as pd

from research.lab.fetch_us import _proxy_env, fetch

OUT = os.path.join("data", "research", "cn_ohlcv.pkl")


def to_yahoo(code: str) -> str:
    ex, num = code.split(".")
    return f"{num}.{'SS' if ex == 'sh' else 'SZ'}"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--hs300-only", action="store_true")
    ap.add_argument("--start", default="2007-01-01")
    args = ap.parse_args()
    member_file = sorted(glob.glob(os.path.join("data", "research", "cn_bars", "members_*.pkl")))[-1]
    members = pd.read_pickle(member_file)
    hs300 = sorted({c for lst in members["hs300"].values() for c in lst})
    codes = list(hs300)
    if not args.hs300_only:
        codes += sorted({c for lst in members["zz500"].values() for c in lst} - set(hs300))
    ymap = {to_yahoo(c): c for c in codes}
    _proxy_env()
    raw = fetch(list(ymap) + ["000300.SS", "000905.SS"], args.start)
    fields = {k: raw[k].rename(columns=ymap) for k in ("open", "high", "low", "close", "volume")}
    source = {c: "yahoo" for c in fields["close"].columns if c in codes}

    idx = pd.DataFrame({"sh.000300": raw["close"].get("000300.SS"),
                        "sh.000905": raw["close"].get("000905.SS")})
    keep = [c for c in fields["close"].columns if c in source]
    fields = {k: v[keep].sort_index() for k, v in fields.items()}
    # Yahoo prints flat bars (volume 0, O=H=L=C) on suspended days: treat as no trade
    close = fields["close"]
    traded = (fields["volume"] > 0) & close.notna()
    fields["tradestatus"] = traded.astype(float)
    fields["preclose"] = close.where(traded).ffill().shift(1)
    fields["pctchg"] = (close / fields["preclose"] - 1) * 100
    missing = sorted(set(codes) - set(source))
    panel = {"fields": fields, "members": members, "index": idx, "source": source,
             "missing": missing}
    with open(OUT, "wb") as fh:
        pickle.dump(panel, fh)
    print(f"Saved {OUT}: {len(keep)} codes, {len(missing)} missing, "
          f"{close.index[0].date()} → {close.index[-1].date()}")
    with open(os.path.join("data", "research", "cn_missing.txt"), "w") as fh:
        fh.write("\n".join(missing))


if __name__ == "__main__":
    main()