"""
Script Name: 
Description: 从富途获取历史K线。通过不同变量，可以获取不复权、前复权、后复权等数据。
    参考地址: https://openapi.futunn.com/futu-api-doc/quote/request-history-kline.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 time
import logging
from futu import *
from datetime import datetime, timedelta
import config

config.setup_logging("./log/get_hs300_his_kline.log")

# 连接 MySQL
connection = pymysql.connect(**config.db_config)

# 复权类型，不复权
# selected_autype = AuType.NONE
# selected_table = "hs300_his_kline_none"

# 复权类型，后复权
selected_autype = AuType.HFQ
selected_table = "hs300_his_kline_hfq"

# 复权类型，默认为 AuType.QFQ ，即前复权
# selected_autype = AuType.QFQ
# selected_table = "hs300_qfq_his"

# 获取当前日期
end_date = datetime.now().strftime('%Y-%m-%d')
# 计算 start_date 为当前日期减去10年，再加一天
start_date = (datetime.now() - timedelta(days=365*10-1)).strftime('%Y-%m-%d')

# 定义插入数据的函数
def insert_data(connection, data):
    try:
        with connection.cursor() as cursor:
            for index, row in data.iterrows():
                sql = f"""
                INSERT INTO {selected_table} (code, name, time_key, open, close, high, low, pe_ratio, turnover_rate, volume, turnover, change_rate, last_close)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                name = VALUES(name),
                open = VALUES(open),
                close = VALUES(close),
                high = VALUES(high),
                low = VALUES(low),
                pe_ratio = VALUES(pe_ratio),
                turnover_rate = VALUES(turnover_rate),
                volume = VALUES(volume),
                turnover = VALUES(turnover),
                change_rate = VALUES(change_rate),
                last_close = VALUES(last_close)
                """
                cursor.execute(sql, (
                    row['code'], row['name'], row['time_key'], row['open'], row['close'],
                    row['high'], row['low'], row['pe_ratio'], row['turnover_rate'], 
                    row['volume'], row['turnover'], row['change_rate'], row['last_close']
                ))
            connection.commit()
    except pymysql.MySQLError as e:
        logging.error(f"Error occurred while inserting data: {e}")
        print(f"Error occurred while inserting data: {e}")

# 获取 hs300 表中的所有股票代码
def get_hs300_codes():
    with connection.cursor() as cursor:
        cursor.execute("SELECT code FROM hs300 ")
        return cursor.fetchall()

# 初始化 futu 行情连接
quote_ctx = OpenQuoteContext(host='127.0.0.1', port=11111)

try:
    hs300_codes = get_hs300_codes()
    for code_row in hs300_codes:
        code = code_row[0]  # 从数据库行中提取 code

        # 获取历史 K 线数据，设置分页请求
        ret, data, page_req_key = quote_ctx.request_history_kline(code, autype=selected_autype, start=start_date, end=end_date, max_count=500)
        if ret == RET_OK:
            logging.info(f"成功获取 {code} 的第一页数据，共 {len(data)} 行")
            print(f"成功获取 {code} 的第一页数据，共 {len(data)} 行")
            
            # 插入第一页数据
            insert_data(connection, data)
        else:
            logging.error(f"获取 {code} 的数据失败: {data}")
            print(f"获取 {code} 的数据失败: {data}")

        # 分页拉取
        while page_req_key is not None:
            time.sleep(1)  # 休眠 5 秒
            ret, data, page_req_key = quote_ctx.request_history_kline(code, autype=selected_autype, start=start_date, end=end_date, max_count=500, page_req_key=page_req_key)
            if ret == RET_OK:
                logging.info(f"成功获取 {code} 的分页数据，共 {len(data)} 行")
                print(f"成功获取 {code} 的分页数据，共 {len(data)} 行")

                # 插入分页数据
                insert_data(connection, data)
            else:
                logging.error(f"分页数据获取失败: {data}")
                print(f"分页数据获取失败: {data}")

        # 每次获取完一个股票的数据后，休眠 5 秒
        time.sleep(2)

finally:
    quote_ctx.close()  # 关闭 futu 连接
    connection.close()  # 关闭 MySQL 连接