"""CSI 300 (HS300), CSI 500 (ZZ500), and full A-share universe helpers."""
import logging
from datetime import date, timedelta

import baostock as bs
import pandas as pd

logger = logging.getLogger(__name__)

# A-share code patterns (baostock dotted form): SH main/STAR (sh.6xxxxx),
# SZ main/SME (sz.0xxxxx), ChiNext (sz.3xxxxx). Excludes indices and B-shares.
_ASHARE_RE = r"^sh\.6\d{5}$|^sz\.[03]\d{5}$"
_SZ_INDEX_RE = r"^sz\.399"


def get_hs300_stocks() -> pd.DataFrame:
    """Fetch the current CSI 300 constituents from baostock.

    Returns:
        DataFrame with columns ``code`` (e.g. ``sh600000``), ``name``, ``date``.
    """
    bs.login()
    try:
        rs = bs.query_hs300_stocks()
        stocks = []
        while rs.next():
            stocks.append(rs.get_row_data())
    finally:
        bs.logout()
    df = pd.DataFrame(stocks, columns=["code", "name", "date"])
    df["code"] = df["code"].str.replace(".", "", regex=False)
    return df


def get_zz500_stocks() -> pd.DataFrame:
    """Fetch the current CSI 500 (ZZ500) constituents from baostock.

    Returns:
        DataFrame with columns ``code`` (e.g. ``sh600006``), ``name``, ``date``.
    """
    bs.login()
    try:
        rs = bs.query_zz500_stocks()
        stocks = []
        while rs.next():
            stocks.append(rs.get_row_data())
    finally:
        bs.logout()
    df = pd.DataFrame(stocks, columns=["code", "name", "date"])
    df["code"] = df["code"].str.replace(".", "", regex=False)
    return df


def get_all_stocks(day: str = "") -> pd.DataFrame:
    """Fetch every listed A-share from baostock's all-stock snapshot.

    Queries ``query_all_stock`` for a single trading day and keeps only A-shares
    (SH main/STAR, SZ main/SME/ChiNext), dropping indices and B-shares. If the
    given day is a non-trading day baostock returns nothing, so we walk back up
    to 10 days to land on the most recent trading day.

    Args:
        day: ``YYYY-MM-DD`` snapshot day; defaults to today (walks back to the
            last trading day).

    Returns:
        DataFrame with columns ``code`` (e.g. ``sh600000``), ``name``.
    """
    start = date.fromisoformat(day) if day else date.today()
    bs.login()
    try:
        rows: list = []
        fields: list = []
        for back in range(11):
            probe = (start - timedelta(days=back)).isoformat()
            rs = bs.query_all_stock(day=probe)
            fields = rs.fields
            while rs.next():
                rows.append(rs.get_row_data())
            if rows:
                logger.info("query_all_stock: %d rows on %s", len(rows), probe)
                break
    finally:
        bs.logout()

    df = pd.DataFrame(rows, columns=fields)
    code = df["code"]
    keep = code.str.match(_ASHARE_RE) & ~code.str.match(_SZ_INDEX_RE)
    df = df[keep].copy()
    df["code"] = df["code"].str.replace(".", "", regex=False)
    df = df.rename(columns={"code_name": "name"})
    return df[["code", "name"]].reset_index(drop=True)