import tushare as ts
from sqlalchemy import create_engine, text
from datetime import datetime
import sys

# =======================
# 1. 日期校验（YYYYMMDD）
# =======================
def validate_dates(start_str, end_str):
    start_str = start_str.strip()
    end_str = end_str.strip()

    if len(start_str) != 8 or len(end_str) != 8:
        print(f"❌ 日期格式错误，应为 YYYYMMDD: {start_str}, {end_str}")
        sys.exit(1)

    if not (start_str.isdigit() and end_str.isdigit()):
        print("❌ 日期必须为纯数字")
        sys.exit(1)

    try:
        start = datetime.strptime(start_str, "%Y%m%d")
        end = datetime.strptime(end_str, "%Y%m%d")
    except Exception as e:
        print(f"❌ 日期解析失败: {e}")
        sys.exit(1)

    if start > end:
        print("❌ 开始时间不能大于结束时间")
        sys.exit(1)

    return start_str, end_str


# =======================
# 2. 创建数据库连接（解决 % 编码问题）
# =======================
def get_engine():
    engine = create_engine(
        "mysql+pymysql://yangfan@8.159.129.156:10836/level?charset=utf8mb4",
        connect_args={
            "password": "aA%8023321088"
        },
        pool_pre_ping=True
    )
    return engine


# =======================
# 3. 确保表存在
# =======================
def ensure_table(engine):
    sql = """
    CREATE TABLE IF NOT EXISTS trade_calendar (
        cal_date DATE NOT NULL,
        exchange VARCHAR(10) NOT NULL,
        is_open TINYINT,
        pretrade_date DATE,
        PRIMARY KEY (cal_date, exchange)
    ) CHARSET=utf8mb4;
    """

    # 防止隐藏乱码字符
    sql = sql.encode("utf-8", "ignore").decode("utf-8")

    with engine.begin() as conn:
        conn.execute(text(sql))

    print("✅ 表检查完成")


# =======================
# 4. 主逻辑
# =======================
def run(start_date, end_date, token):
    start_date, end_date = validate_dates(start_date, end_date)

    ts.set_token(token)
    pro = ts.pro_api()

    engine = get_engine()

    ensure_table(engine)

    print(f"📅 获取交易日历: {start_date} ~ {end_date}")

    try:
        df = pro.trade_cal(
            exchange="",
            start_date=start_date,
            end_date=end_date
        )
    except Exception as e:
        print(f"❌ 获取数据失败: {e}")
        return

    if df is None or df.empty:
        print("⚠️ 无数据")
        return

    records = []
    for _, row in df.iterrows():
        try:
            records.append({
                "cal_date": datetime.strptime(row["cal_date"], "%Y%m%d").date(),
                "exchange": row["exchange"],
                "is_open": int(row["is_open"]),
                "pretrade_date": (
                    datetime.strptime(row["pretrade_date"], "%Y%m%d").date()
                    if row["pretrade_date"] else None
                )
            })
        except Exception as e:
            print(f"⚠️ 数据解析异常: {row} -> {e}")
            continue

    sql = text("""
    INSERT INTO trade_calendar (cal_date, exchange, is_open, pretrade_date)
    VALUES (:cal_date, :exchange, :is_open, :pretrade_date)
    ON DUPLICATE KEY UPDATE
        is_open = VALUES(is_open),
        pretrade_date = VALUES(pretrade_date)
    """)

    with engine.begin() as conn:
        conn.execute(sql, records)

    print(f"✅ 写入完成，共 {len(records)} 条")


# =======================
# 5. 程序入口（只改这里）
# =======================
if __name__ == "__main__":
    START_DATE = "20240101"
    END_DATE   = "20260423"

    TUSHARE_TOKEN = "0fad3cf498757089e2630028455d5bbe1637475788bcdaa1f2175e93"

    run(START_DATE, END_DATE, TUSHARE_TOKEN)