"""
量化数据导入核心模块
连接 Docker PostgreSQL (192.168.27.11:5438)
从 Tushare 拉取数据并批量导入
"""
import time
import logging
from datetime import datetime, timedelta
from typing import Optional, List, Dict

import tushare as ts
import pandas as pd
import psycopg2
from psycopg2 import sql
from psycopg2.extras import execute_values
from sqlalchemy import create_engine

from config import DB_CONFIG, TUSHARE_TOKEN, BATCH_SIZE, START_DATE, END_DATE, PASSWORD_ENCODED

# ============================================================
# 日志配置
# ============================================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("import_data.log", encoding="utf-8"),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger(__name__)

# ============================================================
# 初始化连接
# ============================================================

# psycopg2 原生连接 (用于执行 DDL / 精细控制)
def get_pg_connection():
    """获取 psycopg2 原生连接"""
    return psycopg2.connect(
        host=DB_CONFIG["host"],
        port=DB_CONFIG["port"],
        database=DB_CONFIG["database"],
        user=DB_CONFIG["user"],
        password=DB_CONFIG["password"],
    )


# SQLAlchemy 引擎 (用于 DataFrame.to_sql)
_sqlalchemy_engine = None


def get_sqlalchemy_engine():
    """获取 SQLAlchemy 引擎 (单例)"""
    global _sqlalchemy_engine
    if _sqlalchemy_engine is None:
        db_url = (
            f"postgresql://{DB_CONFIG['user']}:{PASSWORD_ENCODED}"
            f"@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
        )
        _sqlalchemy_engine = create_engine(db_url, pool_size=5, max_overflow=10)
    return _sqlalchemy_engine


# Tushare Pro API (单例)
_ts_pro = None


def get_ts_pro():
    """获取 Tushare Pro API 实例 (单例)"""
    global _ts_pro
    if _ts_pro is None:
        ts.set_token(TUSHARE_TOKEN)
        _ts_pro = ts.pro_api()
        logger.info("Tushare Pro API 初始化完成")
    return _ts_pro


# ============================================================
# 通用导入工具函数
# ============================================================

def normalize_columns(df: pd.DataFrame) -> pd.DataFrame:
    """
    标准化列名：Tushare 返回的列名可能有大小写差异，统一转小写
    """
    df.columns = [c.lower() for c in df.columns]
    return df


def safe_float(val):
    """安全转换为浮点数，NaN -> None"""
    try:
        if pd.isna(val):
            return None
        return float(val)
    except (ValueError, TypeError):
        return None


def _get_table_columns(table_name: str, conn) -> set:
    """
    查询数据库表的实际列名集合
    """
    cursor = conn.cursor()
    try:
        cursor.execute(
            """
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name = %s
            """,
            (table_name,),
        )
        cols = {row[0] for row in cursor.fetchall()}
    finally:
        cursor.close()
    return cols


def batch_insert(table_name: str, df: pd.DataFrame, conn, conflict_columns: List[str]):
    """
    使用 execute_values 批量 UPSERT (INSERT ... ON CONFLICT)
    - table_name: 目标表名
    - df: 待导入 DataFrame
    - conn: psycopg2 连接
    - conflict_columns: 冲突列 (唯一约束列)，冲突时更新其他列

    自动过滤 DataFrame 中数据库表不存在的列，避免
    'column "xxx" of relation "yyy" does not exist' 错误。
    """
    if df.empty:
        logger.warning(f"  {table_name}: 空数据，跳过")
        return 0

    # ---- 动态过滤：只保留数据库表中存在的列 ----
    db_columns = _get_table_columns(table_name, conn)
    if not db_columns:
        # 表可能不存在，回退到不过滤（后续会抛出真实错误）
        logger.warning(f"  {table_name}: 未查询到表结构，使用原始列")
    else:
        extra_cols = [c for c in df.columns if c not in db_columns]
        if extra_cols:
            logger.warning(
                f"  {table_name}: 过滤掉表中不存在的列 {extra_cols}"
            )
        df = df[[c for c in df.columns if c in db_columns]]

    if df.empty:
        logger.warning(f"  {table_name}: 过滤后无可用列，跳过")
        return 0

    columns = list(df.columns)

    # 过滤后验证冲突列仍存在
    missing_conflict = [c for c in conflict_columns if c not in columns]
    if missing_conflict:
        logger.error(
            f"  {table_name}: 冲突列 {missing_conflict} 不在可用列中，跳过"
        )
        return 0

    # ---- 按冲突列去重 ----
    # PostgreSQL 的 ON CONFLICT DO UPDATE 不允许同一命令中出现重复冲突键：
    #   "ON CONFLICT DO UPDATE command cannot affect row a second time"
    # Tushare 财务接口 (income/balancesheet/cashflow/fina_indicator) 对同一股票
    # 同一报告期可能返回多行数据（如多次公告修正，ann_date 不同但冲突键相同），
    # 必须先在批内去重。
    #
    # 为保证保留的是"最新公告"的数据而非仅依赖 Tushare 返回顺序：
    # 若存在公告日期列 (ann_date / f_ann_date)，先按公告日期升序排序，
    # 再 drop_duplicates(keep="last") 即可稳定保留最新一条 (公告日期最大)，
    # 且缺失公告日期的行 (NaT) 会排在最后，仅当无公告日期时才被保留。
    #
    # 注意：去重必须在此处 (NaN->None 替换之前) 执行，
    # 此时日期列仍为 datetime64 类型，sort_values(na_position="last")
    # 能正确处理 NaT；若在替换之后排序，object 类型混合日期/None 排序不可靠。
    #
    # 关键修复：PostgreSQL 对 VARCHAR/数值列会做隐式类型转换，因此数据库眼里
    #   report_type = '1' (str) 与 report_type = 1 (int) 是"同一个约束值"；
    # 但 pandas 的 drop_duplicates 认为 '1' != 1，导致去重不彻底、仍报错。
    # 因此去重前先将非日期冲突列统一转成 str 再判断重复 (不修改原始 DataFrame)。
    dedup_cols = [c for c in conflict_columns if c in columns]
    before_dedup = len(df)
    date_cols = [c for c in ["f_ann_date", "ann_date"] if c in columns]
    if date_cols:
        df = df.sort_values(date_cols, na_position="last")
    if dedup_cols:
        # 规范化去重键：非 datetime64 列统一 astype(str)，消除类型不一致的隐患
        key_df = df[dedup_cols].copy()
        for c in key_df.columns:
            if not pd.api.types.is_datetime64_any_dtype(key_df[c]):
                key_df[c] = key_df[c].astype(str)
        dup_mask = key_df.duplicated(keep="last")
        df = df[~dup_mask]
    after_dedup = len(df)
    if after_dedup < before_dedup:
        logger.warning(
            f"  {table_name}: 检测到 {before_dedup - after_dedup} 行重复冲突键，"
            f"已保留最新公告记录 (去重后 {after_dedup} 行)"
        )

    # pd.NaT / numpy NaN 等无法被 psycopg2 识别，统一替换为 Python None
    df = df.where(pd.notna(df), None)

    rows = [tuple(row) for row in df.itertuples(index=False)]

    # 构建 ON CONFLICT 子句（使用 sql.Identifier 防止注入/语法错误）
    conflict_identifiers = sql.SQL(", ").join(map(sql.Identifier, conflict_columns))
    # 构建 UPDATE SET 子句 (排除冲突列)
    update_cols = [c for c in columns if c not in conflict_columns]
    if not update_cols:
        # 没有需要更新的列，使用 DO NOTHING
        upsert_sql = sql.SQL(
            "INSERT INTO {table} ({cols}) VALUES %s "
            "ON CONFLICT ({conflict}) DO NOTHING"
        ).format(
            table=sql.Identifier(table_name),
            cols=sql.SQL(", ").join(map(sql.Identifier, columns)),
            conflict=conflict_identifiers,
        )
    else:
        update_set = sql.SQL(", ").join(
            sql.SQL("{col} = EXCLUDED.{col}").format(col=sql.Identifier(c))
            for c in update_cols
        )
        upsert_sql = sql.SQL(
            "INSERT INTO {table} ({cols}) VALUES %s "
            "ON CONFLICT ({conflict}) DO UPDATE SET {update_set}"
        ).format(
            table=sql.Identifier(table_name),
            cols=sql.SQL(", ").join(map(sql.Identifier, columns)),
            conflict=conflict_identifiers,
            update_set=update_set,
        )

    cursor = conn.cursor()
    try:
        execute_values(cursor, upsert_sql.as_string(cursor), rows, page_size=BATCH_SIZE)
        conn.commit()
        logger.info(f"  {table_name}: 成功导入 {len(rows)} 条记录")
        return len(rows)
    except Exception as e:
        conn.rollback()
        # 兜底：若批内仍存在重复冲突键 (如极端类型差异未被识别)，
        # 回退为逐行 INSERT。单行命令永远不会触发
        # "ON CONFLICT DO UPDATE command cannot affect row a second time"。
        if "cannot affect row a second time" in str(e):
            logger.warning(
                f"  {table_name}: 批内仍存在重复冲突键，回退为逐行导入 ({len(rows)} 行)"
            )
            try:
                inserted = 0
                for r in rows:
                    execute_values(cursor, upsert_sql.as_string(cursor), [r], page_size=1)
                    inserted += 1
                conn.commit()
                logger.info(f"  {table_name}: 逐行成功导入 {inserted} 条记录")
                return inserted
            except Exception as e2:
                conn.rollback()
                logger.error(f"  {table_name}: 逐行导入失败 - {e2}")
                raise
        logger.error(f"  {table_name}: 批量导入失败 - {e}")
        raise
    finally:
        cursor.close()


def fetch_with_retry(fetch_func, max_retries: int = 3, delay: float = 2.0):
    """
    带重试的数据获取装饰器
    - fetch_func: 数据获取函数 (返回 DataFrame)
    - max_retries: 最大重试次数
    - delay: 重试间隔(秒)
    """
    for attempt in range(max_retries):
        try:
            result = fetch_func()
            if result is not None and not result.empty:
                return result
            logger.warning(f"  第 {attempt+1} 次获取返回空数据，重试...")
        except Exception as e:
            logger.warning(f"  第 {attempt+1} 次获取失败: {e}")
        if attempt < max_retries - 1:
            time.sleep(delay * (attempt + 1))  # 递增延迟
    return pd.DataFrame()


# ============================================================
# 1. 导入股票基本信息
# ============================================================

def import_stock_basic():
    """
    导入股票基本信息 (stock_basic)
    Tushare: stock_basic
    """
    # 确保数据库表结构已初始化
    init_database()

    logger.info("=" * 60)
    logger.info("[1/7] 导入股票基本信息 (stock_basic) ...")

    pro = get_ts_pro()
    conn = get_pg_connection()

    try:
        # 获取全量股票基本信息
        df = pro.stock_basic(
            exchange="",
            list_status="L",
            fields="ts_code,symbol,name,area,industry,market,list_status,list_date,is_hs,act_name,act_ent_type",
        )
        if df is None or df.empty:
            logger.warning("未获取到股票基本信息")
            return

        df = normalize_columns(df)

        # 转换日期格式
        if "list_date" in df.columns:
            df["list_date"] = pd.to_datetime(df["list_date"], format="%Y%m%d", errors="coerce")

        logger.info(f"  获取到 {len(df)} 条股票基本信息")

        # 批量导入
        conflict_cols = ["ts_code"]
        batch_insert("stock_basic", df, conn, conflict_cols)

        # 也尝试获取退市的股票
        try:
            df_d = pro.stock_basic(
                exchange="",
                list_status="D",
                fields="ts_code,symbol,name,area,industry,market,list_status,list_date,is_hs",
            )
            if df_d is not None and not df_d.empty:
                df_d = normalize_columns(df_d)
                if "list_date" in df_d.columns:
                    df_d["list_date"] = pd.to_datetime(
                        df_d["list_date"], format="%Y%m%d", errors="coerce"
                    )
                batch_insert("stock_basic", df_d, conn, conflict_cols)
                logger.info(f"  额外导入 {len(df_d)} 条退市股票信息")
        except Exception as e:
            logger.warning(f"  获取退市股票信息失败: {e}")

    except Exception as e:
        logger.error(f"  导入股票基本信息失败: {e}")
        raise
    finally:
        conn.close()


# ============================================================
# 2. 导入交易日历
# ============================================================

def import_trade_cal(start_date: Optional[str] = None, end_date: Optional[str] = None):
    """
    导入交易日历 (trade_cal)
    Tushare: trade_cal
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    logger.info("=" * 60)
    logger.info(f"[2/7] 导入交易日历 (trade_cal): {start_date} ~ {end_date}")

    pro = get_ts_pro()
    conn = get_pg_connection()

    try:
        # 上交所
        for exchange, ex_name in [("SSE", "上交所"), ("SZSE", "深交所")]:
            df = pro.trade_cal(
                exchange=exchange,
                start_date=start_date.replace("-", ""),
                end_date=end_date.replace("-", ""),
            )
            if df is None or df.empty:
                logger.warning(f"  {ex_name} 交易日历为空")
                continue

            df = normalize_columns(df)

            # 转换日期
            for col in ["cal_date", "pretrade_date"]:
                if col in df.columns:
                    df[col] = pd.to_datetime(df[col], format="%Y%m%d", errors="coerce")

            # 重命名 is_open (Tushare 返回 0/1 整数)
            if "is_open" in df.columns:
                df["is_open"] = df["is_open"].astype(int)

            conflict_cols = ["exchange", "cal_date"]
            batch_insert("trade_cal", df, conn, conflict_cols)
            logger.info(f"  {ex_name}: {len(df)} 条交易日历")

    except Exception as e:
        logger.error(f"  导入交易日历失败: {e}")
        raise
    finally:
        conn.close()

    time.sleep(0.5)  # API 频率限制


# ============================================================
# 3. 导入日线行情 (核心表)
# ============================================================

def import_daily_for_stock(ts_code: str, start_date: str, end_date: str, conn) -> int:
    """
    导入单只股票的日线行情（保留用于单只股票补充/重试）
    返回导入的记录数
    """
    pro = get_ts_pro()

    def fetch():
        return pro.daily_vip(
            ts_code=ts_code,
            start_date=start_date.replace("-", ""),
            end_date=end_date.replace("-", ""),
        )

    df = fetch_with_retry(fetch, max_retries=2)
    if df is None or df.empty:
        return 0

    df = normalize_columns(df)

    if "trade_date" in df.columns:
        df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")

    numeric_cols = [
        "open", "high", "low", "close", "pre_close", "change", "pct_chg",
        "vol", "amount",
    ]
    for col in numeric_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    for col in ["turnover_rate", "volume_ratio", "ma5", "ma10", "ma20", "ma_v_5", "ma_v_10", "ma_v_20"]:
        if col not in df.columns:
            df[col] = None

    conflict_cols = ["ts_code", "trade_date"]
    return batch_insert("daily", df, conn, conflict_cols)


def _normalize_daily_df(df: pd.DataFrame) -> pd.DataFrame:
    """标准化 daily DataFrame 的列和类型 (供 import_daily_by_date 复用)"""
    df = normalize_columns(df)
    if "trade_date" in df.columns:
        df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")
    numeric_cols = [
        "open", "high", "low", "close", "pre_close", "change", "pct_chg",
        "vol", "amount",
    ]
    for col in numeric_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
    for col in ["turnover_rate", "volume_ratio", "ma5", "ma10", "ma20", "ma_v_5", "ma_v_10", "ma_v_20"]:
        if col not in df.columns:
            df[col] = None
    return df


def import_daily_by_date(
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    conn=None,
    sleep_interval: float = 0.3,
    skip_existing: bool = True,
):
    """
    按交易日批量导入日线行情 (高效模式)
    使用 pro.daily_vip(trade_date='YYYYMMDD') 一次性拉取全市场当日数据
    大幅减少 API 调用次数: 约250交易日/年 × 16年 ≈ 4000次 (原来需要 5000股票 × 16年 = 80000次)

    参数:
    - skip_existing: 是否跳过数据库中已有数据的交易日 (默认 True，避免重复导入)

    返回: 失败的交易日列表
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    own_conn = conn is None
    if own_conn:
        conn = get_pg_connection()

    # 从 trade_cal 获取交易日列表
    try:
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT DISTINCT cal_date FROM trade_cal
            WHERE is_open = 1
              AND cal_date >= %s AND cal_date <= %s
            ORDER BY cal_date
            """,
            (start_date, end_date),
        )
        trade_dates = [row[0] for row in cursor.fetchall()]
        cursor.close()
    finally:
        if own_conn:
            conn.close()
            conn = get_pg_connection()

    total_cal = len(trade_dates)
    if total_cal == 0:
        logger.warning(f"  日期范围 {start_date} ~ {end_date} 内无交易日")
        if own_conn:
            conn.close()
        return []

    # --- 跳过已有数据的交易日 ---
    skipped_count = 0
    if skip_existing:
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT DISTINCT trade_date FROM daily
                WHERE trade_date >= %s AND trade_date <= %s
                """,
                (start_date, end_date),
            )
            existing_dates = set(row[0] for row in cursor.fetchall())
            cursor.close()
            if existing_dates:
                # 统一格式为 date 类型再比较
                trade_dates_filtered = [td for td in trade_dates if td not in existing_dates]
                skipped_count = len(trade_dates) - len(trade_dates_filtered)
                trade_dates = trade_dates_filtered
        except Exception as e:
            logger.warning(f"  查询已有交易日失败，将导入全部: {e}")
            # 查询失败时回退到全量导入（安全策略）

    total = len(trade_dates)
    if total == 0:
        logger.info(f"  日期范围 {start_date} ~ {end_date}: 所有 {total_cal} 个交易日均已有数据，无需导入")
        if own_conn:
            conn.close()
        return []

    if skipped_count > 0:
        logger.info(
            f"  日期范围 {start_date} ~ {end_date}: "
            f"共 {total_cal} 个交易日，跳过 {skipped_count} 个已有数据，"
            f"待导入 {total} 个"
        )
    else:
        logger.info(f"  日期范围 {start_date} ~ {end_date}: 共 {total} 个交易日")

    pro = get_ts_pro()
    success_count = 0
    fail_list = []

    for i, td in enumerate(trade_dates, 1):
        try:
            td_str = td.strftime("%Y%m%d") if hasattr(td, "strftime") else str(td).replace("-", "")

            def fetch():
                return pro.daily_vip(trade_date=td_str)

            df = fetch_with_retry(fetch, max_retries=3)
            if df is None or df.empty:
                logger.warning(f"  [{td_str}] 返回空数据 (可能非交易日或API限制)")
                continue

            df = _normalize_daily_df(df)
            conflict_cols = ["ts_code", "trade_date"]
            batch_insert("daily", df, conn, conflict_cols)
            success_count += 1

        except Exception as e:
            logger.error(f"  [{td}] 导入失败: {e}")
            fail_list.append(str(td))
            try:
                conn.rollback()
            except Exception:
                pass

        if i % 50 == 0 or i == total:
            logger.info(f"  进度: {i}/{total} 成功={success_count} 失败={len(fail_list)}")

        time.sleep(sleep_interval)

    if own_conn:
        conn.close()

    logger.info(
        f"  日线行情按日期导入完成: 成功 {success_count}/{total} 个交易日"
    )
    if fail_list:
        logger.warning(f"  失败日期({len(fail_list)}): {fail_list[:20]}...")
    return fail_list


def import_daily_by_year(
    start_year: int = 2010,
    end_year: int = 2025,
    sleep_interval: float = 0.3,
):
    """
    按年份逐批导入日线行情 (按交易日循环拉取全市场数据)
    不再需要 stock_list 参数 — 每次 API 调用拉取当日全市场数据
    """
    logger.info("=" * 60)
    logger.info(
        f"[3/7] 按年导入日线行情 (按交易日): {start_year} ~ {end_year}"
    )

    for year in range(start_year, end_year + 1):
        year_start = f"{year}-01-01"
        year_end = f"{year}-12-31"
        logger.info(f"--- 导入 {year} 年日线行情 ---")
        import_daily_by_date(
            start_date=year_start,
            end_date=year_end,
            sleep_interval=sleep_interval,
        )
        logger.info(f"  {year} 年完成\n")

    logger.info("  所有年份日线行情导入完成!")


# ============================================================
# 4. 导入每日指标 (daily_basic) - 估值/基本面
# ============================================================

def import_daily_basic(
    ts_code: Optional[str] = None,
    trade_date: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    conn=None,
):
    """
    导入每日指标 (daily_basic)
    Tushare: daily_basic
    - 支持按单只股票导入
    - 支持按日期范围全市场导入
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    pro = get_ts_pro()
    own_conn = conn is None
    if own_conn:
        conn = get_pg_connection()

    try:
        # 构建参数
        kwargs = {
            "start_date": start_date.replace("-", ""),
            "end_date": end_date.replace("-", ""),
        }
        if ts_code:
            kwargs["ts_code"] = ts_code
        if trade_date:
            kwargs["trade_date"] = trade_date.replace("-", "")

        def fetch():
            return pro.daily_basic_vip(**kwargs)

        df = fetch_with_retry(fetch, max_retries=2)
        if df is None or df.empty:
            return 0

        df = normalize_columns(df)

        if "trade_date" in df.columns:
            df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")

        numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
        for col in numeric_cols:
            df[col] = pd.to_numeric(df[col], errors="coerce")

        conflict_cols = ["ts_code", "trade_date"]
        return batch_insert("daily_basic", df, conn, conflict_cols)

    finally:
        if own_conn:
            conn.close()


def import_daily_basic_by_date(
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
):
    """
    按日期批量导入每日指标 (全市场)
    Tushare daily_basic 接口可按交易日获取全市场数据，比较高效
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    logger.info("=" * 60)
    logger.info(f"[4/7] 导入每日指标 (daily_basic): {start_date} ~ {end_date}")

    # 获取交易日列表
    conn = get_pg_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT DISTINCT cal_date FROM trade_cal
            WHERE is_open = 1
              AND cal_date >= %s AND cal_date <= %s
            ORDER BY cal_date
            """,
            (start_date, end_date),
        )
        trade_dates = [row[0].strftime("%Y%m%d") for row in cursor.fetchall()]
        cursor.close()
    finally:
        conn.close()

    total = len(trade_dates)
    logger.info(f"  共 {total} 个交易日")

    conn = get_pg_connection()
    for i, td in enumerate(trade_dates, 1):
        try:
            pro = get_ts_pro()

            def fetch_daily_basic():
                return pro.daily_basic_vip(trade_date=td)

            df = fetch_with_retry(fetch_daily_basic, max_retries=3)
            if df is not None and not df.empty:
                df = normalize_columns(df)
                if "trade_date" in df.columns:
                    df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")
                conflict_cols = ["ts_code", "trade_date"]
                batch_insert("daily_basic", df, conn, conflict_cols)
        except Exception as e:
            logger.warning(f"  [{td}] 导入失败: {e}")
            conn.rollback()

        if i % 20 == 0 or i == total:
            logger.info(f"  进度: {i}/{total}")
        time.sleep(0.3)

    conn.close()
    logger.info("  每日指标导入完成")


# ============================================================
# 5. 导入复权因子
# ============================================================

def import_adj_factor(
    ts_code: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    conn=None,
):
    """
    导入复权因子 (adj_factor)
    Tushare: adj_factor
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    pro = get_ts_pro()
    own_conn = conn is None
    if own_conn:
        conn = get_pg_connection()

    try:
        kwargs = {
            "start_date": start_date.replace("-", ""),
            "end_date": end_date.replace("-", ""),
        }
        if ts_code:
            kwargs["ts_code"] = ts_code

        def fetch():
            return pro.adj_factor_vip(**kwargs)

        df = fetch_with_retry(fetch, max_retries=2)
        if df is None or df.empty:
            return 0

        df = normalize_columns(df)

        if "trade_date" in df.columns:
            df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")

        conflict_cols = ["ts_code", "trade_date"]
        return batch_insert("adj_factor", df, conn, conflict_cols)

    finally:
        if own_conn:
            conn.close()


def import_adj_factor_batch(
    stock_list: List[str],
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
):
    """批量导入复权因子"""
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    total = len(stock_list)
    logger.info("=" * 60)
    logger.info(f"[5/7] 导入复权因子 (adj_factor): {total} 只股票")

    conn = get_pg_connection()
    for i, ts_code in enumerate(stock_list, 1):
        try:
            import_adj_factor(ts_code=ts_code, start_date=start_date, end_date=end_date, conn=conn)
        except Exception as e:
            logger.warning(f"  [{ts_code}] 复权因子导入失败: {e}")
            conn.rollback()

        if i % 100 == 0 or i == total:
            logger.info(f"  进度: {i}/{total}")
        time.sleep(0.25)

    conn.close()
    logger.info("  复权因子导入完成")


# ============================================================
# 6. 导入财务数据 (利润表、资产负债表、现金流量表、财务指标)
# ============================================================

def import_financial_statements(
    stock_list: List[str],
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
):
    """
    按股票批量导入三大报表 + 财务指标
    - income: 利润表
    - balancesheet: 资产负债表
    - cashflow: 现金流量表
    - fina_indicator: 财务指标
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    total = len(stock_list)
    logger.info("=" * 60)
    logger.info(
        f"[6/7] 导入财务数据: {start_date} ~ {end_date}, "
        f"共 {total} 只股票"
    )

    period_start = start_date.replace("-", "")
    period_end = end_date.replace("-", "")
    pro = get_ts_pro()
    conn = get_pg_connection()

    for i, ts_code in enumerate(stock_list, 1):
        for table_name, fetch_method in [
            ("income", pro.income_vip),
            ("balancesheet", pro.balancesheet_vip),
            ("cashflow", pro.cashflow_vip),
            ("fina_indicator", pro.fina_indicator_vip),
        ]:
            try:
                if table_name == "fina_indicator":
                    # fina_indicator 参数略有不同
                    df = fetch_method(
                        ts_code=ts_code,
                        start_date=period_start,
                        end_date=period_end,
                    )
                else:
                    df = fetch_method(
                        ts_code=ts_code,
                        start_date=period_start,
                        end_date=period_end,
                    )

                if df is None or df.empty:
                    continue

                df = normalize_columns(df)

                # 转换日期列
                for col in ["ann_date", "f_ann_date", "end_date"]:
                    if col in df.columns:
                        df[col] = pd.to_datetime(df[col], format="%Y%m%d", errors="coerce")

                if table_name == "fina_indicator":
                    conflict_cols = ["ts_code", "end_date"]
                else:
                    conflict_cols = ["ts_code", "end_date", "report_type"]

                batch_insert(table_name, df, conn, conflict_cols)

            except Exception as e:
                logger.warning(f"  [{ts_code}] {table_name}: {e}")
                conn.rollback()

        if i % 50 == 0 or i == total:
            logger.info(f"  财务数据进度: {i}/{total}")
        time.sleep(0.3)

    conn.close()
    logger.info("  财务数据导入完成")


# ============================================================
# 7. 导入指数日线行情
# ============================================================

def import_index_daily(
    index_codes: Optional[List[str]] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
):
    """
    导入指数日线行情 (index_daily)
    默认导入主要指数：上证指数、深证成指、沪深300、中证500、创业板指、科创50
    """
    if index_codes is None:
        index_codes = [
            "000001.SH",  # 上证指数
            "399001.SZ",  # 深证成指
            "000300.SH",  # 沪深300
            "000905.SH",  # 中证500
            "399006.SZ",  # 创业板指
            "000688.SH",  # 科创50
        ]
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    logger.info("=" * 60)
    logger.info(f"[7/7] 导入指数日线行情: {len(index_codes)} 个指数")

    pro = get_ts_pro()
    conn = get_pg_connection()

    for idx_code in index_codes:
        try:
            def fetch():
                return pro.index_daily_vip(
                    ts_code=idx_code,
                    start_date=start_date.replace("-", ""),
                    end_date=end_date.replace("-", ""),
                )

            df = fetch_with_retry(fetch, max_retries=2)
            if df is not None and not df.empty:
                df = normalize_columns(df)
                if "trade_date" in df.columns:
                    df["trade_date"] = pd.to_datetime(
                        df["trade_date"], format="%Y%m%d", errors="coerce"
                    )
                conflict_cols = ["ts_code", "trade_date"]
                n = batch_insert("index_daily", df, conn, conflict_cols)
                logger.info(f"  {idx_code}: {n} 条")
            else:
                logger.warning(f"  {idx_code}: 无数据")
        except Exception as e:
            logger.error(f"  {idx_code}: {e}")
            conn.rollback()
        time.sleep(0.3)

    conn.close()
    logger.info("  指数日线行情导入完成")


# ============================================================
# 8. 初始化数据库 Schema
# ============================================================

# ---- 表结构迁移映射: 为已存在的旧表补充新列 ----
# key: 表名, value: 需要确保存在的列 -> (列定义类型, 注释)
_SCHEMA_MIGRATIONS = {
    "income": {
        "fv_value_chg_gain": "NUMERIC(20,4)",
    },
    "balancesheet": {
        "total_share": "NUMERIC(20,4)",
    },
    "cashflow": {
        "finan_exp": "NUMERIC(20,4)",
    },
    "fina_indicator": {
        "ca_turn": "NUMERIC(16,4)",
    },
}


def _migrate_schema(conn):
    """
    迁移已存在的旧表: 使用 ALTER TABLE ... ADD COLUMN IF NOT EXISTS
    补齐 Tushare API 返回但旧 schema 缺失的列。

    - 对已存在的表生效 (CREATE TABLE IF NOT EXISTS 不会修改旧表)
    - ADD COLUMN IF NOT EXISTS 幂等，可安全重复执行
    - 与其他存储引擎不同，PostgreSQL 的 ADD COLUMN 是 O(1) 元数据操作
    """
    cursor = conn.cursor()
    try:
        for table_name, columns in _SCHEMA_MIGRATIONS.items():
            # 先检查表是否存在
            cursor.execute(
                "SELECT to_regclass(%s)",
                (table_name,),
            )
            if cursor.fetchone()[0] is None:
                continue

            for col_name, col_type in columns.items():
                cursor.execute(
                    sql.SQL("ALTER TABLE {} ADD COLUMN IF NOT EXISTS {} {}").format(
                        sql.Identifier(table_name),
                        sql.Identifier(col_name),
                        sql.SQL(col_type),
                    )
                )
                logger.info(f"  迁移: {table_name}.{col_name} 列已确认 ({col_type})")

        conn.commit()
    except Exception as e:
        conn.rollback()
        logger.warning(f"  表结构迁移失败 (可忽略，导入时自动过滤): {e}")
    finally:
        cursor.close()


def init_database():
    """
    执行 DDL，创建所有表结构
    """
    logger.info("=" * 60)
    logger.info("初始化数据库 Schema ...")

    # 尝试创建数据库（如果不存在）
    try:
        admin_conn = psycopg2.connect(
            host=DB_CONFIG["host"],
            port=DB_CONFIG["port"],
            database="postgres",
            user=DB_CONFIG["user"],
            password=DB_CONFIG["password"],
        )
        admin_conn.autocommit = True
        cursor = admin_conn.cursor()

        cursor.execute(
            "SELECT 1 FROM pg_database WHERE datname = %s",
            (DB_CONFIG["database"],),
        )
        if cursor.fetchone() is None:
            cursor.execute(
                sql.SQL("CREATE DATABASE {} ENCODING 'UTF8'").format(
                    sql.Identifier(DB_CONFIG["database"])
                )
            )
            logger.info(f"  数据库 {DB_CONFIG['database']} 创建成功")
        else:
            logger.info(f"  数据库 {DB_CONFIG['database']} 已存在")

        cursor.close()
        admin_conn.close()
    except Exception as e:
        logger.warning(f"  创建数据库步骤跳过 (可能无权限): {e}")

    # 执行 DDL
    import os as _os
    import re as _re

    schema_path = _os.path.join(
        _os.path.dirname(_os.path.abspath(__file__)), "schema.sql"
    )
    with open(schema_path, "r", encoding="utf-8") as f:
        ddl_sql = f.read()

    # 先去掉单行注释，再按分号分句
    ddl_sql = _re.sub(r"^\s*--.*$", "", ddl_sql, flags=_re.MULTILINE)
    statements = [s.strip() for s in ddl_sql.split(";") if s.strip()]

    conn = get_pg_connection()
    try:
        cursor = conn.cursor()
        for stmt in statements:
            cursor.execute(stmt)
        conn.commit()
        cursor.close()
        logger.info("  Schema 初始化完成")

        # 迁移已存在的旧表: 补齐 Tushare API 新增列
        _migrate_schema(conn)
    except Exception as e:
        conn.rollback()
        logger.error(f"  Schema 初始化失败: {e}")
        raise
    finally:
        conn.close()

# ============================================================
# 9. 获取所有股票列表 (辅助)
# ============================================================

def get_all_stock_codes(include_delisted: bool = False) -> List[str]:
    """
    从 Tushare 获取所有 A 股股票代码列表
    """
    pro = get_ts_pro()
    codes = []

    for status, label in [("L", "上市"), ("D", "退市"), ("P", "暂停")]:
        if status != "L" and not include_delisted:
            continue
        try:
            df = pro.stock_basic(
                exchange="",
                list_status=status,
                fields="ts_code",
            )
            if df is not None and not df.empty:
                codes.extend(df["ts_code"].tolist())
        except Exception as e:
            logger.warning(f"  获取 {label} 股票列表失败: {e}")

    logger.info(f"  获取到 {len(codes)} 只股票代码")
    return codes


def get_stock_codes_from_db(conn=None) -> List[str]:
    """
    从已导入的 stock_basic 表获取股票代码列表
    """
    own_conn = conn is None
    if own_conn:
        conn = get_pg_connection()

    try:
        cursor = conn.cursor()
        cursor.execute("SELECT ts_code FROM stock_basic WHERE list_status = 'L' ORDER BY ts_code")
        codes = [row[0] for row in cursor.fetchall()]
        cursor.close()
        return codes
    finally:
        if own_conn:
            conn.close()


# ============================================================
# 10. 一键全量导入
# ============================================================

def full_import(
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    import_financials: bool = True,
    stock_codes: Optional[List[str]] = None,
):
    """
    一键全量导入：
    1. 初始化 Schema
    2. 股票基本信息
    3. 交易日历
    4. 日线行情
    5. 每日指标(估值)
    6. 复权因子
    7. 财务数据 (可选)
    8. 指数日线行情

    参数:
    - start_date, end_date: 数据范围
    - import_financials: 是否导入财务数据 (耗时较长)
    - stock_codes: 指定股票列表，不传则全量导入
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    start_time = datetime.now()
    logger.info("=" * 70)
    logger.info(f"  开始全量数据导入: {start_date} ~ {end_date}")
    logger.info(f"  PostgreSQL: {DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}")
    logger.info("=" * 70)

    # Step 0: 初始化 Schema
    init_database()

    # Step 1: 股票基本信息
    import_stock_basic()

    # Step 2: 交易日历
    import_trade_cal(start_date, end_date)

    # 获取股票列表
    if stock_codes is None:
        stock_codes = get_stock_codes_from_db()
        if not stock_codes:
            logger.error("无法获取股票列表，请先导入 stock_basic")
            return

    # Step 3: 日线行情 (按年导入，按交易日循环拉取全市场数据)
    import_daily_by_year(
        start_year=int(start_date[:4]),
        end_year=int(end_date[:4]),
    )

    # Step 4: 每日指标 (按日期导入)
    import_daily_basic_by_date(start_date, end_date)

    # Step 5: 复权因子
    import_adj_factor_batch(stock_codes, start_date, end_date)

    # Step 6: 财务数据
    if import_financials:
        import_financial_statements(stock_codes, start_date, end_date)

    # Step 7: 指数日线行情
    import_index_daily(start_date=start_date, end_date=end_date)

    elapsed = datetime.now() - start_time
    logger.info("=" * 70)
    logger.info(f"  全量数据导入完成! 总耗时: {elapsed}")
    logger.info("=" * 70)


# ============================================================
# 11. 进度查询 & 断点续传辅助函数
# ============================================================

def check_daily_progress(conn=None):
    """
    查看 daily 表的导入进度（按年份 + 股票数统计）
    返回各年份的记录数和独立股票数
    """
    own_conn = conn is None
    if own_conn:
        conn = get_pg_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("""
            SELECT 
                EXTRACT(YEAR FROM trade_date)::int AS year,
                COUNT(*) AS records,
                COUNT(DISTINCT ts_code) AS stocks
            FROM daily
            GROUP BY year
            ORDER BY year
        """)
        rows = cursor.fetchall()
        if not rows:
            logger.info("daily 表无数据")
            return {}

        logger.info(f"{'年份':<6} {'记录数':>12} {'股票数':>8}")
        logger.info("-" * 30)
        result = {}
        for year, records, stocks in rows:
            logger.info(f"{year:<6} {records:>12,} {stocks:>8,}")
            result[int(year)] = {"records": int(records), "stocks": int(stocks)}
        cursor.close()
        return result
    finally:
        if own_conn:
            conn.close()


def check_table_summary(conn=None):
    """
    查看所有表的导入概览
    """
    own_conn = conn is None
    if own_conn:
        conn = get_pg_connection()
    try:
        tables = [
            ("stock_basic", None),
            ("trade_cal", (("trade_cal", "cal_date"),)),
            ("daily", (("daily", "trade_date"),)),
            ("daily_basic", (("daily_basic", "trade_date"),)),
            ("adj_factor", (("adj_factor", "trade_date"),)),
            ("income", (("income", "end_date"),)),
            ("balancesheet", (("balancesheet", "end_date"),)),
            ("cashflow", (("cashflow", "end_date"),)),
            ("fina_indicator", (("fina_indicator", "end_date"),)),
            ("index_daily", (("index_daily", "trade_date"),)),
        ]
        cursor = conn.cursor()
        logger.info(f"{'表名':<20} {'记录数':>12} {'最早日期':>12} {'最晚日期':>12}")
        logger.info("-" * 62)
        for table_name, date_info in tables:
            try:
                cursor.execute(sql.SQL("SELECT COUNT(*) FROM {}").format(sql.Identifier(table_name)))
                count = cursor.fetchone()[0]
                if date_info:
                    tbl, col = date_info
                    cursor.execute(
                        sql.SQL("SELECT MIN({}), MAX({}) FROM {}").format(
                            sql.Identifier(col), sql.Identifier(col), sql.Identifier(tbl)
                        )
                    )
                    min_d, max_d = cursor.fetchone()
                    min_str = str(min_d)[:10] if min_d else "N/A"
                    max_str = str(max_d)[:10] if max_d else "N/A"
                    logger.info(f"{table_name:<20} {count:>12,} {min_str:>12} {max_str:>12}")
                else:
                    logger.info(f"{table_name:<20} {count:>12,}")
            except Exception as e:
                logger.warning(f"{table_name:<20} 查询失败: {e}")
        cursor.close()
    finally:
        if own_conn:
            conn.close()


def get_missing_daily_dates(
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    conn=None,
) -> List[str]:
    """
    获取 daily 表中缺失的交易日列表
    对比 trade_cal 中 is_open=1 的日期和 daily 表已有的 trade_date，
    返回未导入的交易日列表。

    返回: 缺失交易日字符串列表 (YYYY-MM-DD 格式)
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    own_conn = conn is None
    if own_conn:
        conn = get_pg_connection()

    try:
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT tc.cal_date
            FROM trade_cal tc
            WHERE tc.is_open = 1
              AND tc.cal_date >= %s
              AND tc.cal_date <= %s
              AND NOT EXISTS (
                  SELECT 1 FROM daily d
                  WHERE d.trade_date = tc.cal_date
              )
            ORDER BY tc.cal_date
            """,
            (start_date, end_date),
        )
        missing_dates = [row[0].strftime("%Y-%m-%d") if hasattr(row[0], "strftime") else str(row[0])[:10]
                         for row in cursor.fetchall()]
        cursor.close()
        return missing_dates
    finally:
        if own_conn:
            conn.close()


def resume_daily_by_date(
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    sleep_interval: float = 0.3,
):
    """
    按缺失日期断点续传日线行情
    自动查询 daily 表已有的 trade_date 与 trade_cal 对比，
    只导入缺失日期的全市场数据。

    用法:
        resume_daily_by_date(start_date="2010-01-01", end_date="2025-12-31")
    """
    if start_date is None:
        start_date = START_DATE
    if end_date is None:
        end_date = END_DATE

    logger.info("=" * 60)
    logger.info(f"[断点续传] 检测缺失日期: {start_date} ~ {end_date}")

    missing_dates = get_missing_daily_dates(start_date, end_date)

    if not missing_dates:
        logger.info("  所有交易日数据已完整，无需续传!")
        return

    total = len(missing_dates)
    logger.info(f"  发现 {total} 个缺失交易日待导入")
    if total <= 20:
        logger.info(f"  缺失日期: {missing_dates}")
    else:
        logger.info(f"  缺失日期 (前20): {missing_dates[:20]}")

    # 按年份分组统计
    years_map: Dict[int, List[str]] = {}
    for d in missing_dates:
        y = int(d[:4])
        years_map.setdefault(y, []).append(d)
    for y in sorted(years_map):
        logger.info(f"    {y} 年: {len(years_map[y])} 个缺失交易日")

    # 逐日期导入
    conn = get_pg_connection()
    pro = get_ts_pro()
    success_count = 0
    fail_list = []

    for i, td_str in enumerate(missing_dates, 1):
        try:
            td_compact = td_str.replace("-", "")

            def fetch():
                return pro.daily_vip(trade_date=td_compact)

            df = fetch_with_retry(fetch, max_retries=3)
            if df is None or df.empty:
                logger.warning(f"  [{td_str}] 返回空数据，跳过")
                continue

            df = _normalize_daily_df(df)
            conflict_cols = ["ts_code", "trade_date"]
            batch_insert("daily", df, conn, conflict_cols)
            success_count += 1

        except Exception as e:
            logger.error(f"  [{td_str}] 导入失败: {e}")
            fail_list.append(td_str)
            try:
                conn.rollback()
            except Exception:
                pass

        if i % 50 == 0 or i == total:
            logger.info(f"  续传进度: {i}/{total} 成功={success_count} 失败={len(fail_list)}")

        time.sleep(sleep_interval)

    conn.close()
    logger.info(f"  断点续传完成: 成功 {success_count}/{total}")
    if fail_list:
        logger.warning(f"  失败日期({len(fail_list)}): {fail_list[:20]}...")
    return fail_list


if __name__ == "__main__":
    # 测试连接
    try:
        conn = get_pg_connection()
        logger.info(f"成功连接到 PostgreSQL: {DB_CONFIG['host']}:{DB_CONFIG['port']}")
        conn.close()
    except Exception as e:
        logger.error(f"无法连接到 PostgreSQL: {e}")
        logger.error("请确认 Docker 容器已启动，且 config.py 中的连接参数正确")