import time
from collections import deque
from datetime import date, datetime

import pandas as pd
import tushare as ts
from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError

from stock_config import TUSHARE_TOKEN, get_engine


TABLE_NAME = "limit_list_d"
LIMIT_TYPES = ("U", "D", "Z")
TUSHARE_START_DATE = date(2020, 1, 1)
MAX_CALLS_PER_MINUTE = 200
FIELDS = (
    "trade_date,ts_code,industry,name,close,pct_chg,amount,limit_amount,"
    "float_mv,total_mv,turnover_ratio,fd_amount,first_time,last_time,"
    "open_times,up_stat,limit_times,limit"
)


class MinuteRateLimiter:
    def __init__(self, max_calls):
        self.max_calls = max_calls
        self.call_times = deque()

    def wait(self):
        now = time.monotonic()
        while self.call_times and now - self.call_times[0] >= 60:
            self.call_times.popleft()
        if len(self.call_times) >= self.max_calls:
            sleep_seconds = 60 - (now - self.call_times[0]) + 0.1
            print(f"⏳ Tushare 限流等待 {sleep_seconds:.1f} 秒。")
            time.sleep(sleep_seconds)
            now = time.monotonic()
            while self.call_times and now - self.call_times[0] >= 60:
                self.call_times.popleft()
        self.call_times.append(time.monotonic())


def parse_date(value):
    value = str(value).strip()
    for fmt in ("%Y%m%d", "%Y-%m-%d"):
        try:
            return datetime.strptime(value, fmt).date()
        except ValueError:
            pass
    raise ValueError(f"日期格式错误: {value}，应为 YYYYMMDD 或 YYYY-MM-DD")


def to_tushare_date(value):
    return value.strftime("%Y%m%d")


def init_db(engine):
    sql = f"""
    CREATE TABLE IF NOT EXISTS `{TABLE_NAME}` (
        `id` BIGINT AUTO_INCREMENT PRIMARY KEY,
        `trade_date` DATE NOT NULL COMMENT '交易日期',
        `ts_code` VARCHAR(20) NOT NULL COMMENT '股票代码',
        `industry` VARCHAR(100) COMMENT '所属行业',
        `name` VARCHAR(50) COMMENT '股票名称',
        `close` DECIMAL(16, 4) COMMENT '收盘价',
        `pct_chg` DECIMAL(10, 4) COMMENT '涨跌幅',
        `amount` DECIMAL(24, 4) COMMENT '成交额',
        `limit_amount` DECIMAL(24, 4) COMMENT '板上成交金额',
        `float_mv` DECIMAL(24, 4) COMMENT '流通市值',
        `total_mv` DECIMAL(24, 4) COMMENT '总市值',
        `turnover_ratio` DECIMAL(10, 4) COMMENT '换手率',
        `fd_amount` DECIMAL(24, 4) COMMENT '封单金额',
        `first_time` VARCHAR(20) COMMENT '首次封板时间',
        `last_time` VARCHAR(20) COMMENT '最后封板时间',
        `open_times` INT COMMENT '打开次数',
        `up_stat` VARCHAR(50) COMMENT '涨停统计',
        `limit_times` INT COMMENT '连板数',
        `limit` VARCHAR(1) NOT NULL COMMENT 'D跌停 U涨停 Z炸板',
        UNIQUE KEY `uniq_trade_code_limit` (`trade_date`, `ts_code`, `limit`),
        KEY `idx_trade_date` (`trade_date`),
        KEY `idx_ts_code` (`ts_code`),
        KEY `idx_limit` (`limit`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    with engine.begin() as conn:
        conn.execute(text(sql))
    print(f"✅ 表 {TABLE_NAME} 检查/创建完成。")


def get_trade_dates(engine, start_date, end_date):
    try:
        query = text(
            """
            SELECT cal_date
            FROM trade_calendar
            WHERE cal_date BETWEEN :start_date AND :end_date
              AND is_open = 1
            ORDER BY cal_date
            """
        )
        df = pd.read_sql(
            query,
            engine,
            params={"start_date": start_date, "end_date": end_date},
        )
        if not df.empty:
            return [pd.to_datetime(value).date() for value in df["cal_date"]]
    except SQLAlchemyError as exc:
        print(f"⚠️ 读取 trade_calendar 失败，将按自然日尝试: {exc}")

    return [value.date() for value in pd.date_range(start_date, end_date, freq="D")]


def clean_text(value):
    if pd.isna(value):
        return None
    return str(value).strip()


def normalize_data(df):
    columns = [
        "trade_date",
        "ts_code",
        "industry",
        "name",
        "close",
        "pct_chg",
        "amount",
        "limit_amount",
        "float_mv",
        "total_mv",
        "turnover_ratio",
        "fd_amount",
        "first_time",
        "last_time",
        "open_times",
        "up_stat",
        "limit_times",
        "limit",
    ]
    for column in columns:
        if column not in df.columns:
            df[column] = None

    result = df[columns].copy()
    result["trade_date"] = pd.to_datetime(
        result["trade_date"], format="%Y%m%d", errors="coerce"
    ).dt.date

    for column in [
        "close",
        "pct_chg",
        "amount",
        "limit_amount",
        "float_mv",
        "total_mv",
        "turnover_ratio",
        "fd_amount",
        "open_times",
        "limit_times",
    ]:
        result[column] = pd.to_numeric(result[column], errors="coerce")

    for column in [
        "ts_code",
        "industry",
        "name",
        "first_time",
        "last_time",
        "up_stat",
        "limit",
    ]:
        result[column] = result[column].map(clean_text)

    result = result.dropna(subset=["trade_date", "ts_code", "limit"])
    return result.drop_duplicates(
        subset=["trade_date", "ts_code", "limit"],
        keep="first",
    )


def fetch_limit_list_d(trade_dates):
    if not trade_dates:
        return pd.DataFrame()

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

    frames = []
    rate_limiter = MinuteRateLimiter(MAX_CALLS_PER_MINUTE)
    total_calls = len(trade_dates) * len(LIMIT_TYPES)
    current_call = 0

    print(
        f"🔎 开始获取 Tushare limit_list_d：{len(trade_dates)} 个交易日，"
        f"{len(LIMIT_TYPES)} 类 limit_type，共 {total_calls} 次调用。"
    )
    for trade_date in trade_dates:
        tushare_date = to_tushare_date(trade_date)
        for limit_type in LIMIT_TYPES:
            current_call += 1
            print(f"  [{current_call}/{total_calls}] {tushare_date} limit_type={limit_type}")
            rate_limiter.wait()
            try:
                df = pro.limit_list_d(
                    trade_date=tushare_date,
                    limit_type=limit_type,
                    fields=FIELDS,
                )
            except Exception as exc:
                print(f"⚠️ 获取失败: {tushare_date} {limit_type}, 错误: {exc}")
                continue

            if df is None or df.empty:
                continue
            frames.append(normalize_data(df))

    if not frames:
        return pd.DataFrame()
    return pd.concat(frames, ignore_index=True).drop_duplicates(
        subset=["trade_date", "ts_code", "limit"],
        keep="first",
    )


def save_data(engine, df, trade_dates):
    if not trade_dates:
        return 0

    with engine.begin() as conn:
        delete_params = {f"date_{idx}": value for idx, value in enumerate(trade_dates)}
        placeholders = ", ".join(f":date_{idx}" for idx in range(len(trade_dates)))
        delete_sql = text(f"DELETE FROM `{TABLE_NAME}` WHERE trade_date IN ({placeholders})")
        delete_result = conn.execute(delete_sql, delete_params)
        print(f"🧹 已清理旧数据 {delete_result.rowcount} 条。")

        if df.empty:
            print("ℹ️ 本次没有可写入的 limit_list_d 数据。")
            return 0

        df.to_sql(
            TABLE_NAME,
            conn,
            if_exists="append",
            index=False,
            chunksize=1000,
            method="multi",
        )
        return len(df)


def run(start_date, end_date):
    start = parse_date(start_date)
    end = parse_date(end_date)
    if start > end:
        raise ValueError("开始日期不能大于结束日期")

    if end < TUSHARE_START_DATE:
        print(f"ℹ️ limit_list_d 接口从 {TUSHARE_START_DATE} 开始处理，当前区间无需同步。")
        return True
    if start < TUSHARE_START_DATE:
        print(f"ℹ️ start_date 早于 {TUSHARE_START_DATE}，已自动调整。")
        start = TUSHARE_START_DATE

    engine = get_engine()
    init_db(engine)

    trade_dates = get_trade_dates(engine, start, end)
    print(f"📅 需要处理 {len(trade_dates)} 个交易日/日期：{start} ~ {end}")
    if not trade_dates:
        return True

    df = fetch_limit_list_d(trade_dates)
    inserted = save_data(engine, df, trade_dates)
    print(f"✅ {TABLE_NAME} 同步完成，写入 {inserted} 条。")
    return True


if __name__ == "__main__":
    START_DATE = "20260507"  # 开始日期，格式 YYYYMMDD 或 YYYY-MM-DD
    END_DATE = "20260512"  # 结束日期，格式 YYYYMMDD 或 YYYY-MM-DD

    run(START_DATE, END_DATE)