"""共用：股票码规范化、覆盖池、幂等写因子表、注册 factor_metadata（自适应列）。"""
import json

import pandas as pd

import db


def to_prefix(ts_code: str) -> str:
    """600000.SH -> SH600000；已是前缀式(SZ002625)或纯代码则原样返回。"""
    s = str(ts_code).strip()
    if "." in s:
        num, exch = s.split(".", 1)
        return f"{exch.upper()}{num}"
    return s


def load_universe() -> set[str]:
    """KG 覆盖池 ts_code 集合（600000.SH 形态），来自基座视图 v_factor_universe。"""
    df = db.read_pg("SELECT ts_code FROM v_factor_universe")
    return set(df["ts_code"].astype(str).str.strip())


def trading_days(start: str, end: str) -> list:
    """目标区间交易日历——用热度表(153，日频、覆盖广)的 distinct trade_date 近似。"""
    df = db.read_mysql("heat",
                       "SELECT DISTINCT trade_date FROM stock_fund_heat_scores "
                       "WHERE trade_date BETWEEN %s AND %s ORDER BY trade_date",
                       (start, end))
    return list(pd.to_datetime(df["trade_date"]))


_CREATE_FACTOR_TABLE = """
CREATE TABLE IF NOT EXISTS {t} (
  trade_date DATE NOT NULL,
  stock_code VARCHAR(15) NOT NULL,
  factor_value DOUBLE,
  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
  PRIMARY KEY (trade_date, stock_code),
  KEY idx_stock (stock_code)
) ENGINE=InnoDB
"""


def ensure_table(table: str) -> None:
    with db.factor_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(_CREATE_FACTOR_TABLE.format(t=table))
        conn.commit()


def write_factor(table: str, df: pd.DataFrame, mode: str) -> None:
    """df 列 [trade_date, stock_code, factor_value] → 幂等写因子表。
    daily/history 都是「删涉及日期区间 → 批插」。stock_code 统一转前缀式。"""
    if df is None or df.empty:
        print(f"  {table}: 无数据（跳过）")
        return
    df = df.dropna(subset=["trade_date", "stock_code", "factor_value"]).copy()
    df["stock_code"] = df["stock_code"].map(to_prefix)
    df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date
    df = df.drop_duplicates(["trade_date", "stock_code"], keep="last")
    if df.empty:
        print(f"  {table}: 清洗后无数据")
        return
    dmin, dmax = df["trade_date"].min(), df["trade_date"].max()
    ensure_table(table)
    rows = list(df[["trade_date", "stock_code", "factor_value"]]
                .itertuples(index=False, name=None))
    with db.factor_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(f"DELETE FROM {table} WHERE trade_date BETWEEN %s AND %s", (dmin, dmax))
            cur.executemany(
                f"INSERT INTO {table} (trade_date, stock_code, factor_value) VALUES (%s,%s,%s)",
                rows)
        conn.commit()
    print(f"  {table}: 写入 {len(rows)} 行, 日期 {dmin}~{dmax}")


def register(factor_code, display_name, table, category, desc,
             author="akg-factor-bridge") -> None:
    """注册 factor_metadata。自适应实际存在的列（避免猜死 schema）——
    列名以线上 factor_metadata 为准，本函数只写它有的列。"""
    values = {
        "factor_code": factor_code, "display_name": display_name,
        "target_ds_name": "ds_a", "target_table_name": table,
        "author": author, "description": desc,
        "category": json.dumps(category, ensure_ascii=False),
        "frequency": "daily", "status": "active", "factor_type": "single",
    }
    with db.factor_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT column_name FROM information_schema.columns "
                        "WHERE table_schema=DATABASE() AND table_name='factor_metadata'")
            cols = {r[0] for r in cur.fetchall()}
            use = [k for k in values if k in cols]
            if "factor_code" not in use:
                raise RuntimeError("factor_metadata 无 factor_code 列？请核实线上 schema")
            ph = ",".join(["%s"] * len(use))
            upd = ",".join(f"{k}=VALUES({k})" for k in use if k != "factor_code")
            cur.execute(
                f"INSERT INTO factor_metadata ({','.join(use)}) VALUES ({ph}) "
                f"ON DUPLICATE KEY UPDATE {upd}", [values[k] for k in use])
        conn.commit()
    print(f"  注册 {factor_code} -> {table}（写入列: {sorted(use)}）")
    if "factor_type" not in cols:
        print("  ⚠️ factor_metadata 无 factor_type 列——平台 /mining/factors/all 按 "
              "factor_type IN('single','multiple') 过滤，缺列可能查不到本因子，请核实。")