import pymysql
import pandas as pd
from futu import *
import logging
from config import db_config  # 引用config.py中的数据库配置

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s - %(message)s')
logger = logging.getLogger()

# 数据库连接函数
def connect_to_db():
    return pymysql.connect(**db_config)

# 从sp300表中获取所有股票代码
def fetch_sp300_codes(connection):
    query = "SELECT code FROM hs300"
    return pd.read_sql(query, connection)

# 获取市场快照并保存到 CSV 文件
def get_market_snapshot_and_save_to_csv(stock_codes, output_file):
    quote_ctx = OpenQuoteContext(host='127.0.0.1', port=11111)
    ret, data = quote_ctx.get_market_snapshot(stock_codes)
    
    if ret == RET_OK:
        logger.info(f"Successfully fetched market snapshot for {len(stock_codes)} codes.")
        # 将数据写入CSV文件
        data.to_csv(output_file, index=False)
        logger.info(f"Snapshot data saved to {output_file}")
    else:
        logger.error(f"Error fetching market snapshot: {data}")
    
    quote_ctx.close()

# 主函数
def main():
    try:
        # 连接数据库
        connection = connect_to_db()

        # 从 sp300 表中获取所有的股票代码
        sp300_codes_df = fetch_sp300_codes(connection)
        
        # 提取股票代码列表
        stock_codes = sp300_codes_df['code'].tolist()

        if not stock_codes:
            logger.warning("No stock codes found in sp300 table.")
            return
        
        # 获取市场快照并保存到 CSV 文件
        output_file = "market_snapshot.csv"
        get_market_snapshot_and_save_to_csv(stock_codes, output_file)

    except Exception as e:
        logger.error(f"An error occurred: {e}")
    finally:
        if connection:
            connection.close()

if __name__ == "__main__":
    main()