import os

import akshare as ak
import pandas as pd
from loguru import logger

# index_hist_df = ak.index_zh_a_hist(
#     symbol="000001",  # 指数代码,如上证指数
#     period="daily",  # K线周期: daily(日K)
#     start_date="19700101",  # 开始日期
#     end_date="22220101",  # 结束日期
# )


def get_all_stock_index():
    index_choice = ["沪深重要指数", "上证系列指数", "深证系列指数", "中证系列指数"]
    index_df_list = []
    for source in index_choice:
        logger.info(f"正在获取 {source}...")
        index_df = ak.stock_zh_index_spot_em(symbol=source)
        index_df["symbol"] = source
        index_df_list.append(index_df)
        logger.info(f"{source}: {index_df.shape[0]}")
    df = pd.concat(index_df_list)
    return df


if __name__ == "__main__":
    # df = get_all_stock_index()
    # df.to_csv("index_all_stock.csv", index=False, encoding="utf-8-sig")

    df = pd.read_csv("index_all_stock.csv", encoding="utf-8-sig")
    for index, row in df.iterrows():
        print(f"正在获取 {index} {row['代码']} ...")
        file_path = f"data/{row['代码']}.csv"
        if os.path.exists(file_path):
            print(f"{row['代码']} 存在")
            continue
        index_hist_df = ak.index_zh_a_hist(
            symbol=row["代码"],  # 指数代码,如上证指数
            period="daily",  # K线周期: daily(日K)
            start_date="19700101",  # 开始日期
            end_date="22220101",  # 结束日期
        )
        index_hist_df.to_csv(
            f"data/{row['代码']}.csv", index=False, encoding="utf-8-sig"
        )
        # time.sleep(15)
    # print(index_hist_df)

    ...