"""
Script Name: 
Description: 从富途获取历史K线。通过不同变量，可以获取不复权、前复权、后复权等数据。
    参考地址: https://openapi.futunn.com/futu-api-doc/quote/get-market-snapshot.html

Author: [Your Name]
Created Date: YYYY-MM-DD
Last Modified: YYYY-MM-DD
Version: 1.0

Modification History:
    - YYYY-MM-DD [Your Name]: 
    - YYYY-MM-DD [Your Name]: 
    - YYYY-MM-DD [Your Name]: 
"""

import pymysql
import pandas as pd
from futu import *
import logging
import config

# 设置日志
config.setup_logging()
logger = logging.getLogger()

# 数据库连接函数
def connect_to_db():
    return pymysql.connect(**config.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 = "./result/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()