A_Share_DP/src/ashare_dp/data/sources/index.py

"""Index daily K-line backfill via Sina API.

Fetches 7 major A-share indices and stores them in the same
Hive-partitioned Parquet format as stock K-lines.

Uses akshare.stock_zh_index_daily() which works via Sina
(verified accessible outside China).
"""

from __future__ import annotations

from datetime import date, datetime

import pandas as pd
from loguru import logger

from ashare_dp.core.models import INDEX_SINA_SYMBOLS, Freq
from ashare_dp.data.store.repository import KLineRepository


def backfill_indices(
    repo: KLineRepository | None = None,
    start_date: str = "19900101",
    end_date: str | None = None,
) -> dict[str, int]:
    """Fetch and store index daily K-line history.

    Args:
        repo: KLineRepository instance. Created if None.
        start_date: Start date in YYYYMMDD format.
        end_date: End date in YYYYMMDD format (default: today).

    Returns:
        Dict mapping ts_code → number of bars stored.
    """
    if repo is None:
        repo = KLineRepository()

    if end_date is None:
        end_date = date.today().strftime("%Y%m%d")

    import akshare as ak

    results: dict[str, int] = {}

    for sina_sym, ts_code in INDEX_SINA_SYMBOLS.items():
        try:
            logger.info(f"Fetching index {sina_sym} ({ts_code})...")
            df = ak.stock_zh_index_daily(symbol=sina_sym)
            if df is None or df.empty:
                logger.warning(f"No data for {sina_sym}")
                results[ts_code] = 0
                continue

            # Normalize columns to match K-line standard schema
            df = _normalize_index_df(df, ts_code)

            # Filter by date range
            df["trade_date"] = pd.to_datetime(df["trade_date"])
            df = df[(df["trade_date"] >= start_date) & (df["trade_date"] <= end_date)]

            if df.empty:
                logger.warning(f"No data in range for {sina_sym}")
                results[ts_code] = 0
                continue

            # Store
            repo.write_klines(df, Freq.d1)
            results[ts_code] = len(df)
            logger.info(f"  {sina_sym}: stored {len(df)} bars, "
                       f"{df['trade_date'].min().date()} ~ {df['trade_date'].max().date()}")

        except Exception as e:
            logger.error(f"Failed to backfill index {sina_sym}: {e}")
            results[ts_code] = 0

    total = sum(results.values())
    logger.info(f"Index backfill complete: {total} total bars across {len(results)} indices")
    return results


def _normalize_index_df(df: pd.DataFrame, ts_code: str) -> pd.DataFrame:
    """Normalize index DataFrame to match stock K-line standard schema.

    Sina index columns: date, open, high, low, close, volume
    Missing: amount (not available for indices from Sina)
    """
    df = df.copy()

    # Standardize column names
    col_map = {
        "date": "trade_time",
        "open": "open",
        "high": "high",
        "low": "low",
        "close": "close",
        "volume": "volume",
    }
    df = df.rename(columns=col_map)

    # Add required columns
    df["ts_code"] = ts_code
    df["trade_date"] = pd.to_datetime(df["trade_time"])
    df["freq"] = Freq.d1.value

    # Index data from Sina doesn't have amount
    if "amount" not in df.columns:
        df["amount"] = 0.0

    # Ensure volume is int
    df["volume"] = df["volume"].fillna(0).astype("int64")

    # Select and order standard columns
    std_cols = ["ts_code", "trade_time", "trade_date", "open", "high", "low", "close", "volume", "amount", "freq"]
    df = df[[c for c in std_cols if c in df.columns]]

    return df