import tushare as ts
import pymysql
import pandas as pd
from datetime import datetime
import time
import argparse

from stock_config import DB_CONFIG, TUSHARE_TOKEN

# ===== 初始化 =====
ts.set_token(TUSHARE_TOKEN)
pro = ts.pro_api()

conn = None
cursor = None


def get_connection():
    global conn, cursor
    if conn is None or not conn.open:
        conn = pymysql.connect(**DB_CONFIG)
        cursor = conn.cursor()
    return conn, cursor


# ===== 建表 =====
def create_tables():
    db_conn, db_cursor = get_connection()
    sql = """
    CREATE TABLE IF NOT EXISTS stock_daily (
        id BIGINT AUTO_INCREMENT PRIMARY KEY,
        ts_code VARCHAR(20) NOT NULL,
        symbol VARCHAR(10) NOT NULL,
        trade_date DATE NOT NULL,
        open DECIMAL(16,4),
        high DECIMAL(16,4),
        low DECIMAL(16,4),
        close DECIMAL(16,4),
        pre_close DECIMAL(16,4),
        change_val DECIMAL(16,4),
        pct_chg DECIMAL(10,4),
        vol BIGINT,
        amount BIGINT,
        UNIQUE KEY uniq_code_date (ts_code, trade_date)
    );
    """
    db_cursor.execute(sql)
    db_conn.commit()
    print("表检查/创建完成")


# ===== 日期转换 =====
def str_to_date(date_str):
    return datetime.strptime(date_str, "%Y%m%d").date()


def mysql_value(value):
    if pd.isna(value):
        return None
    if hasattr(value, "item"):
        return value.item()
    return value


# ===== 获取交易日 =====
def get_trade_dates(start_date, end_date):
    _, db_cursor = get_connection()
    sql = """
    SELECT cal_date 
    FROM trade_calendar
    WHERE cal_date BETWEEN %s AND %s
      AND is_open = 1
    ORDER BY cal_date
    """
    db_cursor.execute(sql, (str_to_date(start_date), str_to_date(end_date)))
    results = db_cursor.fetchall()

    return [row[0].strftime("%Y%m%d") for row in results]


# ===== 插入数据 =====
def insert_data(df):
    db_conn, db_cursor = get_connection()
    if df is None or df.empty:
        return 0, 0

    sql = """
    INSERT IGNORE INTO stock_daily 
    (ts_code, symbol, trade_date, open, high, low, close, pre_close, change_val, pct_chg, vol, amount)
    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
    """

    data = []
    for _, row in df.iterrows():
        ts_code = row['ts_code']
        symbol = ts_code.split('.')[0]  # ✅ 只改这里

        data.append((
            ts_code,
            symbol,
            str_to_date(row['trade_date']),  # 保持 DATE 类型
            mysql_value(row['open']),
            mysql_value(row['high']),
            mysql_value(row['low']),
            mysql_value(row['close']),
            mysql_value(row['pre_close']),
            mysql_value(row['change']),
            mysql_value(row['pct_chg']),
            mysql_value(row['vol']),
            mysql_value(row['amount'])
        ))

    db_cursor.executemany(sql, data)
    db_conn.commit()

    inserted = db_cursor.rowcount
    skipped = len(data) - inserted

    return inserted, skipped


# ===== 主逻辑 =====
def run(start_date, end_date):
    create_tables()

    trade_dates = get_trade_dates(start_date, end_date)

    print(f"需要处理 {len(trade_dates)} 个交易日")

    total_inserted = 0
    total_skipped = 0

    for trade_date in trade_dates:
        print(f"\n处理交易日: {trade_date}")

        try:
            df = pro.daily(trade_date=trade_date)

            inserted, skipped = insert_data(df)

            total_inserted += inserted
            total_skipped += skipped

            print(f"新增: {inserted}, 跳过: {skipped}")

            time.sleep(0.2)

        except Exception as e:
            print(f"失败: {trade_date}, 错误: {e}")

    print("\n====== 汇总 ======")
    print(f"总新增: {total_inserted}")
    print(f"总跳过: {total_skipped}")
    return True


# ===== 输入参数 =====
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="同步日行情")
    parser.add_argument("--start-date", default="20180102", help="开始日期，格式 YYYYMMDD")
    parser.add_argument("--end-date", default="20260430", help="结束日期，格式 YYYYMMDD")
    args = parser.parse_args()

    run(args.start_date, args.end_date)