from multiprocessing import Process
from datetime import datetime

from vnpy.trader.database import BarOverview
from vnpy.trader.datafeed import get_datafeed
from vnpy.trader.database import get_database
from vnpy.trader.object import BarData, HistoryRequest
from vnpy.trader.constant import Exchange, Interval

import re

# 交易所映射关系
EXCHANGE_XT2VT = {
    "SH": Exchange.SSE,
    "SZ": Exchange.SZSE,
    "BJ": Exchange.BSE,
    "SF": Exchange.SHFE,
    "IF": Exchange.CFFEX,
    "INE": Exchange.INE,
    "DF": Exchange.DCE,
    "ZF": Exchange.CZCE,
    "GF": Exchange.GFEX
}

# 开始查询时间
START_TIME = datetime(2018, 1, 1)


def update_history_data() -> None:
    """更新历史合约信息"""
    # 在子进程中加载xtquant
    from xtquant.xtdata import download_history_data

    # 初始化数据服务
    datafeed = get_datafeed()
    datafeed.init()

    # 下载历史合约信息
    download_history_data("", "historycontract")

    print("xtquant历史合约信息下载完成")


def update_bar_data(
    sector_name: str,
    interval: Interval = Interval.MINUTE
) -> None:
    """更新K线数据"""
    # 在子进程中加载xtquant
    from xtquant.xtdata import (
        get_stock_list_in_sector,
        get_instrument_detail
    )

    # 初始化数据服务
    datafeed = get_datafeed()
    datafeed.init()

    # 连接数据库
    database = get_database()

    # 获取当前时间戳
    now: datetime = datetime.now()

    # 获取本地已有数据汇总
    data: list[BarOverview] = database.get_bar_overview()

    overviews: dict[str, BarOverview] = {}
    for o in data:
        vt_symbol: str = f"{o.symbol}.{o.exchange.value}"
        overviews[vt_symbol] = o

    # 查询交易所历史合约代码
    xt_symbols: list[str] = get_stock_list_in_sector(sector_name)

    # 遍历列表查询合约信息
    for xt_symbol in xt_symbols:
        # 查询合约信息
        data: dict = get_instrument_detail(xt_symbol, True)

        # 获取合约到期时间
        expiry: datetime = None
        if data["ExpireDate"]:
            expiry = datetime.strptime(data["ExpireDate"], "%Y%m%d")

        # 拆分迅投研代码
        symbol, xt_exchange = xt_symbol.split(".")

        symbol_main = re.split(r'(\d+)', symbol)[0]

        # 生成本地代码
        exchange: Exchange = EXCHANGE_XT2VT[xt_exchange]
        vt_symbol: str = f"{symbol_main}+'JQ00'.{exchange.value}" or f"{symbol_main}+'00'.{exchange.value}"

        # 查询数据汇总
        overview: BarOverview = overviews.get(vt_symbol, None)

        # 如果已经到期，则跳过
        if overview and expiry and expiry < now:
            continue

        # 实现增量查询
        start: datetime = START_TIME
        if overview:
            start = overview.end

        # 执行数据查询和更新入库
        req: HistoryRequest = HistoryRequest(
            symbol=symbol,
            exchange=exchange,
            start=start,
            end=now,
            interval=interval
        )

        bars: list[BarData] = datafeed.query_bar_history(req)

        if bars:
            database.save_bar_data(bars)

            start_dt: datetime = bars[0].datetime
            end_dt: datetime = bars[-1].datetime
            msg: str = f"{vt_symbol}数据更新成功，{start_dt} - {end_dt}"
            print(msg)


if __name__ == "__main__":
    # 使用子进程更新历史合约信息
    process: Process = Process(target=update_history_data)
    process.start()
    process.join()      # 等待子进程执行完成

    # 更新历史数据
    update_bar_data("上期所")
    update_bar_data("过期上期所")