data_collector.py

from __future__ import annotations

import datetime as dt
import argparse
import logging
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Optional

import pandas as pd
import tushare as ts
from sqlalchemy import create_engine, text


# 加载 .env 文件（如果存在）
_ENV_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
if os.path.exists(_ENV_FILE):
    with open(_ENV_FILE) as _f:
        for _line in _f:
            _line = _line.strip()
            if _line and not _line.startswith("#") and "=" in _line:
                _key, _val = _line.split("=", 1)
                _key = _key.strip()
                _val = _val.strip().strip('"').strip("'")
                if _key not in os.environ:
                    os.environ[_key] = _val

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DB_PATH = os.getenv("FINANCE_DB_PATH") or os.path.join(BASE_DIR, "finance_sqlite.db")
TS_TOKEN = os.getenv("TUSHARE_TOKEN") or os.getenv("TS_TOKEN")
TS_API_URL = os.getenv("TS_API_URL")
ENGINE = create_engine(f"sqlite:///{DB_PATH}", echo=False, future=True)
MAX_WORKERS = int(os.getenv("COLLECTOR_MAX_WORKERS", "1"))
API_CALLS_PER_MIN = min(int(os.getenv("TS_API_CALLS_PER_MIN", "90")), 90)
API_MIN_INTERVAL = 60.0 / max(API_CALLS_PER_MIN, 1)

TABLE_DATE_COLUMNS = {
    "income_statement": "end_date",
    "balance_sheet": "end_date",
    "cash_flow": "end_date",
    "daily_basic": "trade_date",
    "fina_indicator": "end_date",
}
DISCLOSURE_LOOKBACK_PERIODS = int(os.getenv("DISCLOSURE_LOOKBACK_PERIODS", "6"))

# ── 常量 ──
_NO_DATA_DATE = "00000000"
_DEFAULT_START_DATE = "20100101"
_QUARTER_ENDS = ("0331", "0630", "0930", "1231")
_MAX_FINANCE_WORKERS = 10
_PROGRESS_LOG_INTERVAL = 50
_STOCK_DATE_DEFAULTS = {
    "income_date": _NO_DATA_DATE, "income_ann_date": _NO_DATA_DATE,
    "balance_date": _NO_DATA_DATE, "balance_ann_date": _NO_DATA_DATE,
    "cashflow_date": _NO_DATA_DATE, "cashflow_ann_date": _NO_DATA_DATE,
    "fina_date": _NO_DATA_DATE, "fina_ann_date": _NO_DATA_DATE,
    "daily_date": _NO_DATA_DATE,
}


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
_TS_CLIENT_LOCAL = threading.local()
_API_RATE_LOCK = threading.Lock()
_API_LAST_CALL_TIME = 0.0


def get_tushare_client():
    client = getattr(_TS_CLIENT_LOCAL, "client", None)
    if client is None:
        if not TS_TOKEN:
            raise RuntimeError("未配置 TUSHARE_TOKEN，请在 .env 文件中设置或通过环境变量传入")
        ts.set_token(TS_TOKEN)
        client = ts.pro_api()
        if TS_API_URL:
            client._DataApi__http_url = TS_API_URL
        _TS_CLIENT_LOCAL.client = client
    return client


def _wait_for_api_slot() -> None:
    global _API_LAST_CALL_TIME
    with _API_RATE_LOCK:
        now = time.monotonic()
        wait_s = API_MIN_INTERVAL - (now - _API_LAST_CALL_TIME)
        if wait_s > 0:
            time.sleep(wait_s)
        _API_LAST_CALL_TIME = time.monotonic()


def _pro_query(api_name: str, **kwargs):
    _wait_for_api_slot()
    pro = get_tushare_client()
    return getattr(pro, api_name)(**kwargs)


def check_db_exists() -> bool:
    try:
        with ENGINE.connect() as conn:
            row = conn.execute(
                text("SELECT name FROM sqlite_master WHERE type='table' AND name='stock_basic'")
            ).fetchone()
        return row is not None
    except Exception:
        return False


def create_finance_tables() -> bool:
    statements = [
        """
        CREATE TABLE IF NOT EXISTS stock_basic (
            ts_code TEXT PRIMARY KEY,
            name TEXT,
            industry TEXT,
            list_date TEXT,
            list_status TEXT,
            fullname TEXT
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS income_statement (
            ts_code TEXT,
            end_date TEXT,
            ann_date TEXT,
            report_type TEXT,
            revenue REAL,
            profit_dedt REAL,
            net_profit REAL,
            eps REAL,
            roe REAL,
            PRIMARY KEY (ts_code, end_date)
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS fina_indicator (
            ts_code TEXT,
            end_date TEXT,
            ann_date TEXT,
            report_type TEXT,
            roe REAL,
            roe_yearly REAL,
            grossprofit_margin REAL,
            netprofit_margin REAL,
            debt_to_assets REAL,
            -- 银行专用指标
            npl_ratio REAL,              -- 不良贷款率
            provision_coverage REAL,     -- 拨备覆盖率
            net_interest_margin REAL,    -- 净息差(NIM)
            core_capital_adequacy REAL,  -- 核心一级资本充足率
            dividend_yield REAL,         -- 股息率
            -- 保险专用指标
            ev REAL,                     -- 内含价值
            vnb REAL,                    -- 新业务价值
            persistency_rate REAL,       -- 保单继续率
            solvency_ratio REAL,         -- 偿付能力充足率
            investment_return REAL,      -- 总投资收益率
            PRIMARY KEY (ts_code, end_date)
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS balance_sheet (
            ts_code TEXT,
            end_date TEXT,
            ann_date TEXT,
            report_type TEXT,
            total_assets REAL,
            total_liab REAL,
            total_equity REAL,
            PRIMARY KEY (ts_code, end_date)
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS cash_flow (
            ts_code TEXT,
            end_date TEXT,
            ann_date TEXT,
            report_type TEXT,
            n_cashflow_act REAL,
            n_cashflow_inv REAL,
            n_cashflow_fin REAL,
            c_pay_acq_const_fiolta REAL,
            free_cashflow REAL,
            PRIMARY KEY (ts_code, end_date)
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS daily_basic (
            ts_code TEXT,
            trade_date TEXT,
            close REAL,
            pct_chg REAL,
            total_share REAL,
            float_share REAL,
            free_share REAL,
            total_mv REAL,
            circ_mv REAL,
            vol REAL,
            amount REAL,
            turnover_rate REAL,
            pe REAL,
            pb REAL,
            PRIMARY KEY (ts_code, trade_date)
        )
        """,
    ]

    try:
        with ENGINE.begin() as conn:
            for statement in statements:
                conn.execute(text(statement))
            _ensure_table_column(conn, "income_statement", "profit_dedt", "REAL")
            _ensure_table_column(conn, "stock_basic", "list_status", "TEXT")
            _ensure_table_column(conn, "stock_basic", "fullname", "TEXT")
            _ensure_table_column(conn, "daily_basic", "turnover_rate", "REAL")
        logger.info("数据库表结构已就绪")
        return True
    except Exception as exc:
        logger.error(f"创建数据库表失败：{exc}")
        return False


def _ensure_table_column(conn, table_name: str, column_name: str, column_type: str) -> None:
    columns = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table_name})"))}
    if column_name not in columns:
        conn.execute(text(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"))


def download_stock_basic() -> bool:
    try:
        df = _pro_query("stock_basic", exchange="", list_status="L", fields="ts_code,name,industry,list_date,list_status,fullname")
        if df is None or df.empty:
            logger.warning("股票列表为空")
            return False
        df.to_sql("stock_basic", ENGINE, if_exists="replace", index=False)
        logger.info(f"股票列表已更新，共 {len(df)} 只")
        return True
    except Exception as exc:
        logger.error(f"下载股票列表失败：{exc}")
        return False


def get_latest_report_date(ts_code: str, table: str) -> str | None:
    date_column = TABLE_DATE_COLUMNS.get(table)
    if not date_column:
        return None

    try:
        with ENGINE.connect() as conn:
            row = conn.execute(
                text(f"SELECT MAX({date_column}) FROM {table} WHERE ts_code = :ts_code"),
                {"ts_code": ts_code},
            ).fetchone()
        return row[0] if row and row[0] else None
    except Exception as exc:
        logger.warning(f"读取 {ts_code} 在 {table} 的最新日期失败：{exc}")
        return None


def check_table_has_data(table: str) -> bool:
    try:
        with ENGINE.connect() as conn:
            row = conn.execute(text(f"SELECT COUNT(*) FROM {table}")).fetchone()
        return bool(row and row[0] > 0)
    except Exception as exc:
        logger.error(f"检查表 {table} 失败：{exc}")
        return False


def get_latest_trade_date() -> str:
    today = dt.datetime.now()
    if today.weekday() >= 5:
        today -= dt.timedelta(days=today.weekday() - 4)
    return today.strftime("%Y%m%d")


def get_latest_available_report_date() -> str:
    today = dt.datetime.now()
    year, month, day = today.year, today.month, today.day

    # 按A股常规财报最晚披露截止日判断“应已普遍可获得”的最新报告期
    # 年报: 4/30, 一季报: 4/30, 中报: 8/31, 三季报: 10/31
    # 1/1-4/30: 上年年报(1231) + 当年一季报(0331) → 取最晚 0331
    # 5/1-8/31: 当年一季报(0331) + 当年中报(0630) → 取最晚 0630
    # 9/1-10/31: 当年中报(0630)
    # 11/1-12/31: 当年三季报(0930)
    if (month, day) <= (4, 30):
        return f"{year}0331"
    if (month, day) <= (8, 31):
        return f"{year}0630"
    if (month, day) <= (10, 31):
        return f"{year}0630"
    return f"{year}0930"


def _recent_report_periods(limit: int = DISCLOSURE_LOOKBACK_PERIODS) -> list[str]:
    today = dt.datetime.now().date()
    quarter_ends = []
    year = today.year
    while len(quarter_ends) < limit + 2:
        quarter_ends.extend(
            [
                f"{year}1231",
                f"{year}0930",
                f"{year}0630",
                f"{year}0331",
            ]
        )
        year -= 1
    quarter_ends = sorted({period for period in quarter_ends if period <= today.strftime("%Y%m%d")}, reverse=True)
    return quarter_ends[:limit]


def should_update_finance_report(local_report_date: str | None) -> bool:
    if not local_report_date or local_report_date == _NO_DATA_DATE:
        return True
    return local_report_date < get_latest_available_report_date()


def should_update_by_local_date(table: str, local_latest: str | None) -> bool:
    if not local_latest or local_latest == _NO_DATA_DATE:
        return True
    if table == "daily_basic":
        return local_latest < get_latest_trade_date()
    return should_update_finance_report(local_latest)


def get_incremental_start_date(local_latest: str | None, default_start: str = _DEFAULT_START_DATE) -> str:
    if not local_latest or local_latest == _NO_DATA_DATE:
        return default_start
    try:
        latest_date = dt.datetime.strptime(local_latest, "%Y%m%d").date()
        return (latest_date + dt.timedelta(days=1)).strftime("%Y%m%d")
    except Exception:
        return default_start


def get_lookback_start_date(local_latest: str | None, lookback_days: int, default_start: str = _DEFAULT_START_DATE) -> str:
    if not local_latest or local_latest == _NO_DATA_DATE:
        return default_start
    try:
        latest_date = dt.datetime.strptime(local_latest, "%Y%m%d").date()
        return max(
            dt.datetime.strptime(default_start, "%Y%m%d").date(),
            latest_date - dt.timedelta(days=lookback_days),
        ).strftime("%Y%m%d")
    except Exception:
        return default_start


def should_update_data(ts_code: str, table: str) -> bool:
    local_latest = get_latest_report_date(ts_code, table)
    return should_update_by_local_date(table, local_latest)


def _latest_record_map(table: str, date_column: str, ann_column: str | None = None) -> dict[str, dict[str, str]]:
    result: dict[str, dict[str, str]] = {}
    columns = ["ts_code", date_column]
    if ann_column:
        columns.append(ann_column)
    order_cols = [f"{date_column} DESC"]
    if ann_column:
        order_cols.append(f"{ann_column} DESC")

    query = f"SELECT {', '.join(columns)} FROM {table} ORDER BY ts_code, {', '.join(order_cols)}"
    with ENGINE.connect() as conn:
        for row in conn.execute(text(query)):
            ts_code = row[0]
            if ts_code in result:
                continue
            item = {"latest_date": row[1] or _NO_DATA_DATE}
            if ann_column:
                item["latest_ann_date"] = row[2] or _NO_DATA_DATE
            result[ts_code] = item
    return result


def fetch_latest_disclosed_reports(periods: list[str | None] = None) -> dict[str, dict[str, str]]:
    periods = periods or _recent_report_periods()
    result: dict[str, dict[str, str]] = {}
    for period in periods:
        try:
            df = _pro_query("disclosure_date", end_date=period)
        except Exception as exc:
            logger.warning(f"读取财报披露计划失败 {period}: {exc}")
            continue
        if df is None or df.empty:
            continue
        # 用 ann_date（公告日期）判断是否已披露，actual_date 常为空
        date_col = "actual_date" if df["actual_date"].notna().any() else "ann_date"
        released_df = df[(df[date_col].notna()) & (df[date_col] != "")]
        if released_df.empty:
            continue
        released_df = released_df.sort_values(["ts_code", "end_date", date_col], ascending=[True, False, False])
        for _, row in released_df.iterrows():
            ts_code = row.get("ts_code")
            if not ts_code:
                continue
            end_date = row.get("end_date") or _NO_DATA_DATE
            release_date = row.get(date_col) or _NO_DATA_DATE
            current = result.get(ts_code)
            if current is None or (end_date, release_date) > (current["end_date"], current["actual_date"]):
                result[ts_code] = {
                    "end_date": end_date,
                    "actual_date": release_date,
                    "ann_date": row.get("ann_date") or _NO_DATA_DATE,
                    "pre_date": row.get("pre_date") or _NO_DATA_DATE,
                }
    logger.info(f"披露计划扫描完成，识别到 {len(result)} 只股票已有更新财报")
    return result


def _chunked_rows(df: pd.DataFrame, mapping: dict[str, str], extra: dict[str, str | None] = None) -> list[Dict]:
    if df is None or df.empty:
        return []

    extra = extra or {}
    records = []
    for _, row in df.iterrows():
        item = {dest: row.get(src) for dest, src in mapping.items()}
        item.update(extra)
        records.append(item)
    return records


def _save_rows(sql: str, rows: list[Dict]) -> None:
    if not rows:
        return
    with ENGINE.begin() as conn:
        conn.execute(text(sql), rows)


def save_income(ts_code: str, start_date: str, end_date: str, force: bool = False) -> None:
    if not force and not should_update_data(ts_code, "income_statement"):
        return
    df = _pro_query("income", ts_code=ts_code, start_date=start_date, end_date=end_date)
    rows = _chunked_rows(
        df,
        {
            "end_date": "end_date",
            "ann_date": "ann_date",
            "report_type": "report_type",
            "revenue": "revenue",
            "profit_dedt": "profit_dedt",
            "net_profit": "n_income",
            "eps": "basic_eps",
            "roe": "roe_wa",
        },
        {"ts_code": ts_code},
    )
    _save_rows(
        """
        INSERT OR REPLACE INTO income_statement
        (ts_code, end_date, ann_date, report_type, revenue, profit_dedt, net_profit, eps, roe)
        VALUES (:ts_code, :end_date, :ann_date, :report_type, :revenue, :profit_dedt, :net_profit, :eps, :roe)
        """,
        rows,
    )


def save_balance(ts_code: str, start_date: str, end_date: str, force: bool = False) -> None:
    if not force and not should_update_data(ts_code, "balance_sheet"):
        return
    df = _pro_query("balancesheet", ts_code=ts_code, start_date=start_date, end_date=end_date)
    rows = _chunked_rows(
        df,
        {
            "end_date": "end_date",
            "ann_date": "ann_date",
            "report_type": "report_type",
            "total_assets": "total_assets",
            "total_liab": "total_liab",
            "total_equity": "total_hldr_eqy_exc_min_int",
        },
        {"ts_code": ts_code},
    )
    _save_rows(
        """
        INSERT OR REPLACE INTO balance_sheet
        (ts_code, end_date, ann_date, report_type, total_assets, total_liab, total_equity)
        VALUES (:ts_code, :end_date, :ann_date, :report_type, :total_assets, :total_liab, :total_equity)
        """,
        rows,
    )


def save_cashflow(ts_code: str, start_date: str, end_date: str, force: bool = False) -> None:
    if not force and not should_update_data(ts_code, "cash_flow"):
        return
    df = _pro_query("cashflow", ts_code=ts_code, start_date=start_date, end_date=end_date)
    rows = _chunked_rows(
        df,
        {
            "end_date": "end_date",
            "ann_date": "ann_date",
            "report_type": "report_type",
            "n_cashflow_act": "n_cashflow_act",
            "n_cashflow_inv": "n_cashflow_inv",
            "n_cashflow_fin": "n_cashflow_fin",
            "c_pay_acq_const_fiolta": "c_pay_acq_const_fiolta",
            "free_cashflow": "free_cashflow",
        },
        {"ts_code": ts_code},
    )
    _save_rows(
        """
        INSERT OR REPLACE INTO cash_flow
        (ts_code, end_date, ann_date, report_type, n_cashflow_act, n_cashflow_inv,
         n_cashflow_fin, c_pay_acq_const_fiolta, free_cashflow)
        VALUES (:ts_code, :end_date, :ann_date, :report_type, :n_cashflow_act, :n_cashflow_inv,
                :n_cashflow_fin, :c_pay_acq_const_fiolta, :free_cashflow)
        """,
        rows,
    )


def save_daily_basic(ts_code: str, start_date: str, end_date: str, force: bool = False) -> None:
    if not force and not should_update_data(ts_code, "daily_basic"):
        return
    df = _pro_query("daily_basic", ts_code=ts_code, start_date=start_date, end_date=end_date)
    rows = _chunked_rows(
        df,
        {
            "trade_date": "trade_date",
            "close": "close",
            "pct_chg": "pct_chg",
            "total_share": "total_share",
            "float_share": "float_share",
            "free_share": "free_share",
            "total_mv": "total_mv",
            "circ_mv": "circ_mv",
            "vol": "vol",
            "amount": "amount",
            "turnover_rate": "turnover_rate",
            "pe": "pe",
            "pb": "pb",
        },
        {"ts_code": ts_code},
    )
    _save_rows(
        """
        INSERT OR REPLACE INTO daily_basic
        (ts_code, trade_date, close, pct_chg, total_share, float_share, free_share,
         total_mv, circ_mv, vol, amount, turnover_rate, pe, pb)
        VALUES (:ts_code, :trade_date, :close, :pct_chg, :total_share, :float_share, :free_share,
                :total_mv, :circ_mv, :vol, :amount, :turnover_rate, :pe, :pb)
        """,
        rows,
    )


def _get_all_stock_codes() -> list[str]:
    df = pd.read_sql("SELECT ts_code FROM stock_basic ORDER BY ts_code", ENGINE)
    return df["ts_code"].dropna().tolist()


def save_fina_indicator(ts_code: str, start_date: str, end_date: str, force: bool = False) -> None:
    if not force and not should_update_data(ts_code, "fina_indicator"):
        return
    df = _pro_query("fina_indicator", ts_code=ts_code, start_date=start_date, end_date=end_date)
    rows = _chunked_rows(
        df,
        {
            "end_date": "end_date",
            "ann_date": "ann_date",
            "report_type": "report_type",
            "roe": "roe",
            "roe_yearly": "roe_yearly",
            "grossprofit_margin": "grossprofit_margin",
            "netprofit_margin": "netprofit_margin",
            "debt_to_assets": "debt_to_assets",
            "npl_ratio": "npl_ratio",
            "provision_coverage": "provision_coverage",
            "net_interest_margin": "net_interest_margin",
            "core_capital_adequacy": "core_capital_adequacy",
            "dividend_yield": "dividend_yield",
            "ev": "ev",
            "vnb": "vnb",
            "persistency_rate": "persistency_rate",
            "solvency_ratio": "solvency_ratio",
            "investment_return": "investment_return",
        },
        {"ts_code": ts_code},
    )
    _save_rows(
        """
        INSERT OR REPLACE INTO fina_indicator
        (ts_code, end_date, ann_date, report_type, roe, roe_yearly,
         grossprofit_margin, netprofit_margin, debt_to_assets,
         npl_ratio, provision_coverage, net_interest_margin,
         core_capital_adequacy, dividend_yield,
         ev, vnb, persistency_rate, solvency_ratio, investment_return)
        VALUES (:ts_code, :end_date, :ann_date, :report_type, :roe, :roe_yearly,
                :grossprofit_margin, :netprofit_margin, :debt_to_assets,
                :npl_ratio, :provision_coverage, :net_interest_margin,
                :core_capital_adequacy, :dividend_yield,
                :ev, :vnb, :persistency_rate, :solvency_ratio, :investment_return)
        """,
        rows,
    )


def _update_stock(ts_code: str, start_date: str, end_date: str, include_finance: bool) -> None:
    if include_finance:
        save_income(ts_code, start_date, end_date, force=True)
        save_balance(ts_code, start_date, end_date, force=True)
        save_cashflow(ts_code, start_date, end_date, force=True)
        save_fina_indicator(ts_code, start_date, end_date, force=True)
    save_daily_basic(ts_code, start_date, end_date, force=True)


def fetch_all_history(start: str = _DEFAULT_START_DATE, end: str | None = None) -> bool:
    end = end or dt.datetime.now().strftime("%Y%m%d")
    try:
        codes = _get_all_stock_codes()
        failed = 0
        for index, code in enumerate(codes, start=1):
            logger.info(f"[{index}/{len(codes)}] 全量抓取 {code}")
            try:
                _update_stock(code, start, end, include_finance=True)
            except Exception as exc:
                failed += 1
                logger.warning(f"{code} 全量抓取异常：{exc}")
        if failed:
            logger.info(f"全量抓取完成，共 {len(codes)} 只，失败 {failed} 只")
        logger.info("全量历史数据抓取完成")
        return True
    except Exception as exc:
        logger.error(f"全量抓取失败：{exc}")
        return False


def load_all_stock_dates() -> dict[str, dict[str, str]]:
    result: dict[str, dict[str, str]] = {}
    with ENGINE.connect() as conn:
        for row in conn.execute(text("SELECT ts_code FROM stock_basic ORDER BY ts_code")):
            result[row[0]] = {
                "income_date": _NO_DATA_DATE,
                "income_ann_date": _NO_DATA_DATE,
                "balance_date": _NO_DATA_DATE,
                "balance_ann_date": _NO_DATA_DATE,
                "cashflow_date": _NO_DATA_DATE,
                "cashflow_ann_date": _NO_DATA_DATE,
                "fina_indicator_date": _NO_DATA_DATE,
                "fina_indicator_ann_date": _NO_DATA_DATE,
                "trade_date": _NO_DATA_DATE,
            }

    income_map = _latest_record_map("income_statement", "end_date", "ann_date")
    balance_map = _latest_record_map("balance_sheet", "end_date", "ann_date")
    cashflow_map = _latest_record_map("cash_flow", "end_date", "ann_date")
    fina_map = _latest_record_map("fina_indicator", "end_date", "ann_date")
    daily_map = _latest_record_map("daily_basic", "trade_date")

    defaults = {
        "income_date": _NO_DATA_DATE,
        "income_ann_date": _NO_DATA_DATE,
        "balance_date": _NO_DATA_DATE,
        "balance_ann_date": _NO_DATA_DATE,
        "cashflow_date": _NO_DATA_DATE,
        "cashflow_ann_date": _NO_DATA_DATE,
        "fina_indicator_date": _NO_DATA_DATE,
        "fina_indicator_ann_date": _NO_DATA_DATE,
        "trade_date": _NO_DATA_DATE,
    }

    for ts_code, item in income_map.items():
        result.setdefault(ts_code, defaults.copy())
        result[ts_code]["income_date"] = item["latest_date"]
        result[ts_code]["income_ann_date"] = item.get("latest_ann_date", _NO_DATA_DATE)

    for ts_code, item in balance_map.items():
        result.setdefault(ts_code, defaults.copy())
        result[ts_code]["balance_date"] = item["latest_date"]
        result[ts_code]["balance_ann_date"] = item.get("latest_ann_date", _NO_DATA_DATE)

    for ts_code, item in cashflow_map.items():
        result.setdefault(ts_code, defaults.copy())
        result[ts_code]["cashflow_date"] = item["latest_date"]
        result[ts_code]["cashflow_ann_date"] = item.get("latest_ann_date", _NO_DATA_DATE)

    for ts_code, item in fina_map.items():
        result.setdefault(ts_code, defaults.copy())
        result[ts_code]["fina_indicator_date"] = item["latest_date"]
        result[ts_code]["fina_indicator_ann_date"] = item.get("latest_ann_date", _NO_DATA_DATE)

    for ts_code, item in daily_map.items():
        result.setdefault(ts_code, defaults.copy())
        result[ts_code]["trade_date"] = item["latest_date"]

    return result


def get_global_latest_date(table: str, date_col: str) -> str | None:
    try:
        with ENGINE.connect() as conn:
            row = conn.execute(text(f"SELECT MAX({date_col}) FROM {table}")).fetchone()
        return row[0] if row and row[0] else None
    except Exception as exc:
        logger.warning(f"读取表 {table} 最新日期失败：{exc}")
        return None


def _date_range(start_date: str, end_date: str) -> list[str]:
    try:
        start = dt.datetime.strptime(start_date, "%Y%m%d").date()
        end = dt.datetime.strptime(end_date, "%Y%m%d").date()
    except Exception:
        return []

    days: list[str] = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y%m%d"))
        current += dt.timedelta(days=1)
    return days


def save_daily_basic_by_trade_date(trade_date: str, allowed_codes: set | None = None) -> int:
    df = _pro_query("daily_basic", trade_date=trade_date)
    if df is None or df.empty:
        return 0

    if allowed_codes:
        df = df[df["ts_code"].isin(allowed_codes)]
        if df.empty:
            return 0

    rows = _chunked_rows(
        df,
        {
            "ts_code": "ts_code",
            "trade_date": "trade_date",
            "close": "close",
            "pct_chg": "pct_chg",
            "total_share": "total_share",
            "float_share": "float_share",
            "free_share": "free_share",
            "total_mv": "total_mv",
            "circ_mv": "circ_mv",
            "vol": "vol",
            "amount": "amount",
            "turnover_rate": "turnover_rate",
            "pe": "pe",
            "pb": "pb",
        },
    )
    _save_rows(
        """
        INSERT OR REPLACE INTO daily_basic
        (ts_code, trade_date, close, pct_chg, total_share, float_share, free_share,
         total_mv, circ_mv, vol, amount, turnover_rate, pe, pb)
        VALUES (:ts_code, :trade_date, :close, :pct_chg, :total_share, :float_share, :free_share,
                :total_mv, :circ_mv, :vol, :amount, :turnover_rate, :pe, :pb)
        """,
        rows,
    )
    return len(rows)


def update_incremental() -> bool:
    today = dt.datetime.now().strftime("%Y%m%d")
    try:
        if not download_stock_basic():
            return False

        codes = _get_all_stock_codes()
        latest_map = load_all_stock_dates()
        allowed_codes = set(codes)
        disclosed_reports = fetch_latest_disclosed_reports()
        tasks: list[dict[str, object]] = []

        # ── 质量过滤：只采集有市场价值的高质量股票 ──
        STOCK_MIN_TOTAL_MV = 500000  # 50亿市值阈值（total_mv 单位=万元）
        try:
            with ENGINE.connect() as conn:
                latest_trade = conn.execute(text("SELECT MAX(trade_date) FROM daily_basic")).fetchone()[0]
                if latest_trade:
                    mv_rows = conn.execute(text(
                        f"SELECT ts_code, total_mv FROM daily_basic "
                        f"WHERE trade_date = '{latest_trade}' AND total_mv IS NOT NULL"
                    )).fetchall()
                    quality_codes = {row[0] for row in mv_rows if row[1] >= STOCK_MIN_TOTAL_MV}
                    skipped = len(codes) - len(quality_codes.intersection(codes))
                    if skipped > 0:
                        codes = [c for c in codes if c in quality_codes]
                        allowed_codes = set(codes)
                        logger.info(f"质量过滤: 跳过 {skipped} 只市值<50亿的股票，保留 {len(codes)} 只")
        except Exception as exc:
            logger.warning(f"市值过滤失败（继续全量更新）: {exc}")

        if not codes:
            logger.warning("质量过滤后无股票需要更新")
            return True

        for code in codes:
            latest_dates = latest_map.get(
                code,
                {
                    "income_date": _NO_DATA_DATE,
                    "income_ann_date": _NO_DATA_DATE,
                    "balance_date": _NO_DATA_DATE,
                    "balance_ann_date": _NO_DATA_DATE,
                    "cashflow_date": _NO_DATA_DATE,
                    "cashflow_ann_date": _NO_DATA_DATE,
                    "fina_indicator_date": _NO_DATA_DATE,
                    "fina_indicator_ann_date": _NO_DATA_DATE,
                    "trade_date": _NO_DATA_DATE,
                },
            )
            latest_income = latest_dates.get("income_date")
            latest_balance = latest_dates.get("balance_date")
            latest_cashflow = latest_dates.get("cashflow_date")
            latest_fina = latest_dates.get("fina_indicator_date")
            latest_trade = latest_dates.get("trade_date")
            latest_income_ann = latest_dates.get("income_ann_date")
            latest_balance_ann = latest_dates.get("balance_ann_date")
            latest_cashflow_ann = latest_dates.get("cashflow_ann_date")
            latest_fina_ann = latest_dates.get("fina_indicator_ann_date")

            disclosed = disclosed_reports.get(code)
            finance_released = disclosed is not None
            released_end_date = disclosed["end_date"] if disclosed else _NO_DATA_DATE
            released_actual_date = disclosed["actual_date"] if disclosed else _NO_DATA_DATE

            if finance_released:
                need_income = (
                    (released_end_date > (latest_income or _NO_DATA_DATE)) or (
                        released_end_date == (latest_income or _NO_DATA_DATE)
                        and released_actual_date > (latest_income_ann or _NO_DATA_DATE)
                    )
                )
                need_balance = (
                    (released_end_date > (latest_balance or _NO_DATA_DATE)) or (
                        released_end_date == (latest_balance or _NO_DATA_DATE)
                        and released_actual_date > (latest_balance_ann or _NO_DATA_DATE)
                    )
                )
                need_cashflow = (
                    (released_end_date > (latest_cashflow or _NO_DATA_DATE)) or (
                        released_end_date == (latest_cashflow or _NO_DATA_DATE)
                        and released_actual_date > (latest_cashflow_ann or _NO_DATA_DATE)
                    )
                )
                need_fina = (
                    (released_end_date > (latest_fina or _NO_DATA_DATE)) or (
                        released_end_date == (latest_fina or _NO_DATA_DATE)
                        and released_actual_date > (latest_fina_ann or _NO_DATA_DATE)
                    )
                )
            else:
                # quick 模式：没有披露计划时，只在完全无数据时尝试补齐，避免新股死循环
                need_income = not latest_income or latest_income == _NO_DATA_DATE
                need_balance = not latest_balance or latest_balance == _NO_DATA_DATE
                need_cashflow = not latest_cashflow or latest_cashflow == _NO_DATA_DATE
                need_fina = not latest_fina or latest_fina == _NO_DATA_DATE
            need_daily = should_update_by_local_date("daily_basic", latest_trade)

            tasks.append(
                {
                    "ts_code": code,
                    "need_income": int(need_income),
                    "need_balance": int(need_balance),
                    "need_cashflow": int(need_cashflow),
                    "need_fina": int(need_fina),
                    "need_daily": int(need_daily),
                    "income_start": get_incremental_start_date(latest_income, default_start=released_end_date),
                    "balance_start": get_incremental_start_date(latest_balance, default_start=released_end_date),
                    "cashflow_start": get_incremental_start_date(latest_cashflow, default_start=released_end_date),
                    "fina_start": get_incremental_start_date(latest_fina, default_start=released_end_date),
                    "daily_start": get_incremental_start_date(latest_trade),
                    "released_report_date": released_end_date,
                    "released_actual_date": released_actual_date,
                    "finance_end_date": released_end_date if finance_released else today,
                }
            )

        task_df = pd.DataFrame(tasks)
        task_df.to_csv("incremental_update_checklist.csv", index=False, encoding="utf-8-sig")
        logger.info(
            f"更新清单已生成 incremental_update_checklist.csv | "
            f"income:{int(task_df['need_income'].sum())} balance:{int(task_df['need_balance'].sum())} "
            f"cash:{int(task_df['need_cashflow'].sum())} fina:{int(task_df['need_fina'].sum())} "
            f"daily:{int(task_df['need_daily'].sum())} / total:{len(task_df)}"
        )

        need_daily_df = task_df[task_df["need_daily"] == 1]
        if not need_daily_df.empty:
            global_latest_trade = get_global_latest_date("daily_basic", "trade_date") or _NO_DATA_DATE
            daily_batch_start = get_incremental_start_date(global_latest_trade)
            if daily_batch_start <= today:
                total_daily_rows = 0
                for trade_date in _date_range(daily_batch_start, today):
                    inserted = save_daily_basic_by_trade_date(trade_date, allowed_codes=allowed_codes)
                    if inserted > 0:
                        total_daily_rows += inserted
                        logger.info(f"daily_basic 批量更新 {trade_date} 写入 {inserted} 行")
                logger.info(f"daily_basic 批量更新完成，共写入 {total_daily_rows} 行")

        finance_tasks = [row for row in tasks if row["need_income"] or row["need_balance"] or row["need_cashflow"] or row["need_fina"]]
        if finance_tasks:
            workers = min(MAX_WORKERS, 10)
            logger.info(f"财报并发更新开始，线程数 {workers}，基于披露计划触发更新")

            def _run_finance_task(task: dict[str, object]) -> str:
                code = str(task["ts_code"])
                try:
                    if task["need_income"]:
                        save_income(code, str(task["income_start"]), str(task["finance_end_date"]), force=True)
                    if task["need_balance"]:
                        save_balance(code, str(task["balance_start"]), str(task["finance_end_date"]), force=True)
                    if task["need_cashflow"]:
                        save_cashflow(code, str(task["cashflow_start"]), str(task["finance_end_date"]), force=True)
                    if task["need_fina"]:
                        save_fina_indicator(code, str(task["fina_start"]), str(task["finance_end_date"]), force=True)
                except Exception as exc:
                    logger.warning(f"{code} 财报采集异常：{exc}")
                return code

            completed = 0
            failed = 0
            with ThreadPoolExecutor(max_workers=workers) as executor:
                futures = [executor.submit(_run_finance_task, task) for task in finance_tasks]
                for future in as_completed(futures):
                    try:
                        _ = future.result()
                        completed += 1
                    except Exception as exc:
                        failed += 1
                        logger.warning(f"财报任务异常：{exc}")
                    if (completed + failed) % 50 == 0 or (completed + failed) == len(finance_tasks):
                        logger.info(f"财报并发进度 {completed}/{len(finance_tasks)}，失败 {failed}")

        logger.info("增量更新完成")
        return True
    except Exception as exc:
        logger.error(f"增量更新失败：{exc}")
        return False


def main() -> None:
    parser = argparse.ArgumentParser(description="Tushare 数据采集器")
    parser.add_argument(
        "--mode",
        choices=["quick", "full"],
        default=os.getenv("COLLECTOR_MODE", "quick"),
        help="quick: 按披露截止日快速增量更新；full: 全量历史更新",
    )
    parser.add_argument(
        "--start",
        default=_DEFAULT_START_DATE,
        help="full 模式起始日期，格式 YYYYMMDD，默认 20100101",
    )
    parser.add_argument(
        "--end",
        default=None,
        help="更新结束日期，格式 YYYYMMDD，默认当天",
    )
    args = parser.parse_args()

    logger.info(f"开始执行数据采集，模式: {args.mode}")
    if not create_finance_tables():
        return

    if not check_db_exists() or not check_table_has_data("stock_basic"):
        logger.info("首次初始化：下载股票列表并全量抓取历史数据")
        if download_stock_basic():
            fetch_all_history()
        return

    if args.mode == "full":
        logger.info(f"检测到已有数据库，执行全量更新：{args.start} -> {args.end or dt.datetime.now().strftime('%Y%m%d')}")
        if download_stock_basic():
            fetch_all_history(start=args.start, end=args.end)
        return

    logger.info("检测到已有数据库，执行快速更新（按披露截止日规则）")
    update_incremental()


if __name__ == "__main__":
    main()