#!/usr/bin/env python3
"""sector_enrich_cninfo.py — 用巨潮 cninfo 批量补全行业/主营业务映射
覆盖全部 active holding_strategies。限速 0.3s/只。
"""
import sqlite3, time, sys
import akshare as ak

DB = "/home/hmo/MoFin/data/mofin.db"

def main():
    conn = sqlite3.connect(DB)
    codes = [r[0] for r in conn.execute(
        "SELECT DISTINCT code FROM holding_strategies WHERE status='active'")]
    have = {r[0] for r in conn.execute("SELECT code FROM stock_sectors")}
    todo = [c for c in codes if c not in have]
    print(f"总 {len(codes)} 只，缺 {len(todo)} 只")

    ok = fail = 0
    for i, code in enumerate(todo, 1):
        # 港股 (5位 0/1 开头) cninfo 不支持
        if len(str(code)) == 5 and str(code)[0] in "01":
            continue
        try:
            df = ak.stock_profile_cninfo(symbol=str(code))
            if df is not None and len(df):
                row = df.iloc[0].to_dict()
                sector = str(row.get("所属行业") or "").strip()
                business = str(row.get("主营业务") or "").strip()[:200]
                if sector:
                    conn.execute(
                        "INSERT OR REPLACE INTO stock_sectors (code, sector_name, source, updated_at) "
                        "VALUES (?, ?, 'cninfo', datetime('now','localtime'))",
                        (code, sector))
                    if business:
                        # business 存到 sector_name 扩展？表只有 sector_name，把 business 塞进 source? 不，加列太多。
                        # sector_name 保留行业，business 记到 holding_strategies.sector_context（如果为空/被污染）
                        conn.execute(
                            "UPDATE holding_strategies SET sector_context=? WHERE code=? AND status='active' "
                            "AND (sector_context IS NULL OR sector_context='' OR sector_context LIKE '大盘上涨比%')",
                            (f"行业{sector} | {business[:60]}", code))
                    ok += 1
                else:
                    fail += 1
            else:
                fail += 1
        except Exception as e:
            fail += 1
            if fail <= 3:
                print(f"  {code} 失败: {str(e)[:80]}")
        if i % 20 == 0:
            conn.commit()
            print(f"  进度 {i}/{len(todo)} ok={ok} fail={fail}")
        time.sleep(0.3)
    conn.commit()
    total = conn.execute("SELECT COUNT(*) FROM stock_sectors").fetchone()[0]
    print(f"完成: 新增 {ok}, 失败 {fail}, stock_sectors 总数 {total}")
    conn.close()

if __name__ == "__main__":
    main()