import pymysql
import logging
import sys
import time
import config
import argparse

# 配置不同市场的表名管理
tables_mapping = {
    'sp500': {
        'none_his_kline': 'sp500_his_kline_none',
        'hfq_his_kline': 'sp500_hfq_his_202410',
        'rehab_table': 'futu_rehab'
    },
    'hs300': {
        'none_his_kline': 'hs300_his_kline_none',
        'hfq_his_kline': 'hs300_hfq_his_202410',
        'rehab_table': 'futu_rehab'
    }
}

# 日志配置
config.setup_logging("./log/stat_hfq_kline.log")
logger = logging.getLogger()

# MySQL数据库连接
def connect_to_db():
    try:
        return pymysql.connect(
            **config.db_config,
            cursorclass=pymysql.cursors.DictCursor
        )
    except pymysql.MySQLError as e:
        logger.error(f"Error connecting to the database: {e}", exc_info=True)
        return None

# 从指定市场表中读取code和code_name字段
def fetch_codes_from_market_table(market):
    db = connect_to_db()
    if db is None:
        logger.error("Failed to connect to database.")
        return []
    
    try:
        with db.cursor() as cursor:
            cursor.execute(f"SELECT code, code_name FROM {market} limit 1")
            codes = cursor.fetchall()
        return codes
    finally:
        db.close()

# 读取复权因子数据
def fetch_rehab_data(db, code):
    try:
        with db.cursor() as cursor:
            cursor.execute(f"SELECT ex_div_date, backward_adj_factorA, backward_adj_factorB FROM futu_rehab WHERE code = %s ORDER BY ex_div_date DESC", (code,))
            return cursor.fetchall()
    except pymysql.MySQLError as e:
        logger.error(f"Error fetching rehab data for {code}: {e}", exc_info=True)
        return []

# 读取不复权的股票价格数据
def fetch_kline_none_data(db, table_name, code):
    try:
        with db.cursor() as cursor:
            cursor.execute(f"SELECT code, time_key, close FROM {table_name} WHERE code = %s ORDER BY time_key ASC", (code,))
            return cursor.fetchall()
    except pymysql.MySQLError as e:
        logger.error(f"Error fetching kline none data for {code}: {e}", exc_info=True)
        return []

# 插入后复权价格到 hfq 表
def insert_hfq_data(db, hfq_data, hfq_table):
    try:
        with db.cursor() as cursor:
            insert_query = f"""
            INSERT INTO {hfq_table} (code, name, time_key, open, close)
            VALUES (%s, %s, %s, %s, %s)
            ON DUPLICATE KEY UPDATE close = VALUES(close)
            """
            cursor.executemany(insert_query, hfq_data)
        db.commit()
    except pymysql.MySQLError as e:
        logger.error(f"Error inserting hfq data: {e}", exc_info=True)

# 计算后复权价格
def calculate_hfq_price(market):
    db = connect_to_db()
    if db is None:
        return
    
    # 获取表名
    table_names = tables_mapping[market]
    none_his_kline_table = table_names['none_his_kline']
    hfq_his_kline_table = table_names['hfq_his_kline']

    # 获取股票代码
    codes = fetch_codes_from_market_table(market)
    
    for row in codes:
        code = row['code']
        name = row['code_name']
        
        # 如果是 sp500 市场，拼接 'US.' + code
        rehab_code = code
        if market == 'sp500':
            rehab_code = 'US.' + code
        
        logger.info(f"Processing {code} ({name})...")
        
        # 获取复权因子数据
        rehab_res = fetch_rehab_data(db, rehab_code)
        if not rehab_res:
            logger.warning(f"No rehab data found for {code}")
            continue
        
        # 获取不复权的价格数据
        kline_none = fetch_kline_none_data(db, none_his_kline_table, code)
        if not kline_none:
            logger.warning(f"No kline none data found for {code}")
            continue
        
        hfq_data = []
        # 遍历kline_none，计算后复权价格
        for kline_row in kline_none:
            close_price = kline_row['close']
            time_key = kline_row['time_key']
            
            # 将 time_key 转换为 date 格式
            time_key_date = time_key.date()

            # 遍历 rehab_res，找到相应的复权因子
            for rehab_row in rehab_res:
                if rehab_row['ex_div_date'] <= time_key_date:
                    hfq_close = (close_price * rehab_row['backward_adj_factorA']) + rehab_row['backward_adj_factorB']
                    close_price = hfq_close
            
            # 保存计算后的复权价格
            hfq_data.append((code, name, time_key, 0.0, hfq_close))
        
        # 插入后复权价格数据
        insert_hfq_data(db, hfq_data, hfq_his_kline_table)
        logger.info(f"Inserted HFQ data for {code} ({name})")

        time.sleep(1)
    
    db.close()

if __name__ == "__main__":
    # 命令行参数处理
    parser = argparse.ArgumentParser(description='Calculate HFQ Prices for Market')
    parser.add_argument('--market', type=str, default='hs300', help='Market to process (sp500 or hs300)')
    args = parser.parse_args()
    
    # 调用主函数
    calculate_hfq_price(args.market)