#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""csindex 公告解析 - 重建中证1000/2000 历史成份股曾经入选集 (治幸存者偏差)

数据源:
  1. 列表 API: POST https://www.csindex.com.cn/csindex-home/announcement/queryAnnouncementByVo
     payload: {lang, classlist, indexlist, page:{desc,key,page,rows}, related_topics, typelist}
  2. 详情 API: GET https://www.csindex.com.cn/csindex-home/announcement/queryAnnouncementById?id={nid}
     返 {content (HTML), enclosureList:[{fileUrl, fileName}]}
  3. 附件: PDF (定期调整 多指数合并) 或 xlsx/xls (临时调整 单/多指数)

实证覆盖 (2026-07-22 实地验证):
  CSI 1000 (000852) 2014-09 发布:
    - 2018-07-11 起 csindex 有调整公告 (28 份: 7 定期 + 21 临时)
    - 2014-2018 期间的 June/Dec 定期调整不在 csindex (gap: ~7 round)
  CSI 2000 (932000) 2023-08-10 发布:
    - csindex 上无任何 CSI 2000 样本调整公告 (定期调整 PDF 中均无 CSI 2000 section)
    - 仅 launch xlsx 有初始 2000 只样本 (2023-08-10)
    - 历史/最新调整需查指数详情页 "拟生效样本" 或 wind/choice (不在本脚本范围)

输出:
  data/index_const_hist/000852_announce_union.parquet  - CSI 1000 曾经入选集
  data/index_const_hist/932000_announce_union.parquet  - CSI 2000 (initial only)
  schema: updateDate / index_code / code / code_name / adjust_type / notice_id / source

约束:
  - 直连不走代理 (unset)
  - 串行 + sleep(1.0~1.5s)
  - UA header
"""
import argparse
import json
import logging
import os
import re
import socket
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# ======================== 硬约束 ========================
for _k in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY", "all_proxy", "ALL_PROXY"):
    os.environ.pop(_k, None)
socket.setdefaulttimeout(30)

try:
    sys.stdout.reconfigure(line_buffering=True)
except (AttributeError, ValueError):
    pass

import pandas as pd  # noqa: E402
import openpyxl  # noqa: E402
import pdfplumber  # noqa: E402

# ======================== 配置 ========================
BASE = "https://www.csindex.com.cn/csindex-home"
LIST_URL = f"{BASE}/announcement/queryAnnouncementByVo"
DETAIL_URL = f"{BASE}/announcement/queryAnnouncementById"
UA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0 Safari/537.36"

DEFAULT_OUT_DIR = "/Users/chufeng/.openclaw/sanguo_projects/sanguo_vnpy_v2/data/index_const_hist"
OUT_DIR = Path(os.environ.get("INDEX_HIST_OUT_DIR", DEFAULT_OUT_DIR))
CACHE_DIR = Path(os.environ.get("CSINDEX_CACHE", "/tmp/csindex_raw"))
DETAIL_CACHE_DIR = CACHE_DIR / "detail"

# CSI 1000 相关调整公告 id (实证已发现, 脚本也支持自动发现)
CSI1000_REGULAR_IDS = [14796, 15044, 15267, 15471, 15690, 3006000, 3006137]
CSI1000_TEMP_IDS = [
    12446, 12965, 13070, 13212, 13281, 13334, 13765, 14092,
    14842, 15019, 15232, 15342, 15357, 15390, 15575, 15613, 15648,
    1006042, 3006027, 3006041, 3006120,
]
# CSI 2000 launch xlsx (initial 2000 stocks)
CSI2000_LAUNCH_ID = 14883

# ======================== logging ========================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("csindex")


# ======================== HTTP ========================
def _post_json(url: str, payload: dict, timeout: int = 30) -> dict:
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url, data=data, method="POST",
        headers={"Content-Type": "application/json", "User-Agent": UA},
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def _get_json(url: str, timeout: int = 30) -> dict:
    req = urllib.request.Request(url, headers={"User-Agent": UA})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def _download(url: str, target: Path, timeout: int = 60) -> int:
    """下载含中文名 URL (自动 quote, 防双重编码)"""
    p = urllib.parse.urlsplit(url)
    path_q = p.path if "%" in p.path else urllib.request.quote(p.path)
    quoted = urllib.parse.urlunsplit((p.scheme, p.netloc, path_q, p.query, p.fragment))
    req = urllib.request.Request(quoted, headers={"User-Agent": UA})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        data = r.read()
    target.parent.mkdir(parents=True, exist_ok=True)
    with open(target, "wb") as f:
        f.write(data)
    return len(data)


# ======================== 1. 列表 ========================
def _fetch_paginated(payload_base: dict, label: str) -> List[dict]:
    """通用分页拉取(单线程 + sleep)"""
    all_items = []
    page = 1
    while True:
        payload = dict(payload_base)
        payload["page"] = {"desc": "", "key": "", "page": page, "rows": 100}
        try:
            d = _post_json(LIST_URL, payload)
        except Exception as e:
            log.error(f"[{label}] page {page} err: {e}; retry once after 5s")
            time.sleep(5)
            d = _post_json(LIST_URL, payload)
        if d.get("code") != "200":
            log.error(f"[{label}] page {page} API err: {d}")
            break
        items = d.get("data") or []
        all_items.extend(items)
        total = d.get("total") or 0
        log.info(f"  [{label}] page {page}: +{len(items)} (cum={len(all_items)}/{total})")
        if not items or len(all_items) >= total:
            break
        page += 1
        time.sleep(1.0)
    return all_items


def fetch_all_notices(cache_path: Path, force: bool = False) -> List[dict]:
    """拉全量公告列表 (cached)"""
    if cache_path.exists() and not force:
        log.info(f"使用缓存列表: {cache_path}")
        with open(cache_path, encoding="utf-8") as f:
            return json.load(f)

    log.info(f"分页拉取全量公告: {LIST_URL}")
    all_items = _fetch_paginated(
        {"lang": "cn", "classlist": [], "indexlist": [],
         "related_topics": [], "typelist": []},
        label="all",
    )

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    with open(cache_path, "w", encoding="utf-8") as f:
        json.dump(all_items, f, ensure_ascii=False)
    log.info(f"缓存列表: {cache_path} (total={len(all_items)})")
    return all_items


def fetch_notices_by_index(index_code: str, cache_path: Path, force: bool = False) -> List[dict]:
    """按 indexCode 精准分页拉公告列表 (cached)

    实证: indexCode='000852' 返 96 条(44 调样, 回溯到 2007), 远优于 title 过滤的 28 条。
    """
    if cache_path.exists() and not force:
        log.info(f"使用缓存 indexCode 列表: {cache_path}")
        with open(cache_path, encoding="utf-8") as f:
            return json.load(f)

    log.info(f"分页拉取 indexCode={index_code} 公告: {LIST_URL}")
    items = _fetch_paginated(
        {"lang": "cn", "classlist": [], "indexlist": [], "indexCode": index_code,
         "related_topics": [], "typelist": []},
        label=f"idx={index_code}",
    )

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    with open(cache_path, "w", encoding="utf-8") as f:
        json.dump(items, f, ensure_ascii=False)
    log.info(f"缓存 indexCode 列表: {cache_path} (total={len(items)})")
    return items


def filter_csi1000_notices(notices: List[dict]) -> List[dict]:
    """筛 CSI 1000 调整公告 (theme=指数调样 + title 含 中证1000)"""
    out = []
    seen_ids = set()
    for x in notices:
        title = x.get("title") or ""
        if x.get("theme") != "指数调样":
            continue
        if "中证1000" not in title:
            continue
        if "调整" not in title:
            continue
        # 排除 "不实施" 通知
        if "不实施" in title:
            continue
        if x["id"] in seen_ids:
            continue
        seen_ids.add(x["id"])
        out.append(x)
    out.sort(key=lambda x: x.get("publishDate", ""))
    return out


def filter_adjustment_notices(notices: List[dict], keyword: str = "") -> List[dict]:
    """通用调整公告筛选: theme=指数调样 + (可选) title 含 keyword

    用于 indexCode 拉取的列表(已按指数过滤,无需 title 匹配)。
    """
    out = []
    seen_ids = set()
    for x in notices:
        if x.get("theme") != "指数调样":
            continue
        title = x.get("title") or ""
        if keyword and keyword not in title:
            continue
        if "不实施" in title:
            continue
        if x["id"] in seen_ids:
            continue
        seen_ids.add(x["id"])
        out.append(x)
    out.sort(key=lambda x: x.get("publishDate", ""))
    return out


# ======================== 2. 详情 + 下载 ========================
def fetch_detail(nid: int, cache_dir: Path) -> dict:
    """详情 (cached by id)"""
    cache_file = cache_dir / f"detail_{nid}.json"
    if cache_file.exists():
        with open(cache_file, encoding="utf-8") as f:
            return json.load(f)
    try:
        d = _get_json(f"{DETAIL_URL}?id={nid}")
    except Exception as e:
        log.error(f"detail {nid} err: {e}")
        return {}
    cache_file.parent.mkdir(parents=True, exist_ok=True)
    with open(cache_file, "w", encoding="utf-8") as f:
        json.dump(d, f, ensure_ascii=False)
    time.sleep(1.0)
    return d


def extract_file_urls(detail_data: dict) -> List[Tuple[str, str]]:
    """从 enclosureList + content 内嵌 href 抽 (fileName, fileUrl)"""
    urls = []
    for e in (detail_data.get("enclosureList") or []):
        fu = e.get("fileUrl")
        fn = e.get("fileName") or "file"
        if fu:
            urls.append((fn, fu))
    content = detail_data.get("content") or ""
    for m in re.finditer(r'href="([^"]+)"', content):
        fu = m.group(1)
        if re.search(r"\.(xlsx|pdf|xls|csv)(\?|$)", fu, re.I):
            urls.append(("inline", fu))
    # 去重保序
    seen, out = set(), []
    for fn, fu in urls:
        if fu not in seen:
            seen.add(fu)
            out.append((fn, fu))
    return out


def download_notice_files(nid: int, publish_date: str, urls: List[Tuple[str, str]], cache_dir: Path) -> List[Path]:
    """下载该公告的所有附件,返本地路径列表"""
    out = []
    for fn, fu in urls:
        ext = fu.rsplit(".", 1)[-1].lower().split("?")[0]
        local = cache_dir / f"{nid}_{publish_date}.{ext}"
        if not local.exists():
            try:
                sz = _download(fu, local)
                log.info(f"    dl [{sz}B] {fn} -> {local.name}")
            except Exception as ex:
                log.error(f"    dl ERR {fu}: {ex}")
                continue
        else:
            log.info(f"    cache {local.name}")
        out.append(local)
    return out


# ======================== 3. 解析 ========================
def _norm_code(c) -> Optional[str]:
    """规整为 6 位字符串代码 (前补 0)"""
    if c is None:
        return None
    s = str(c).strip()
    # 排除空/NA
    if not s or s.lower() in ("nan", "none", ""):
        return None
    # 纯数字 (可能 int 转 str 失去前导 0)
    if s.isdigit():
        s = s.zfill(6)
        if len(s) > 6:
            return None  # 异常长
        return s
    # H开头 (港股代码, CSI 1000 不含, 跳过)
    if s.startswith("H") or s.startswith("688") and not s.isdigit():
        return None
    # 6 位字母数字 (如 SH/SZ 前缀)
    m = re.search(r"(\d{6})", s)
    return m.group(1) if m else None


def parse_xlsx_adjustments(path: Path, target_index_code: str = "000852") -> Tuple[List[dict], List[dict]]:
    """解析 xlsx/xls -> (add_rows, remove_rows)
    每个 row = {code, code_name}
    Sheet 名: 调入/调出 或 换入/换出 (反向)
    """
    add_rows, remove_rows = [], []

    if path.suffix.lower() == ".xls":
        xl = pd.ExcelFile(path)
        df_map = {sn: xl.parse(sn, dtype=str) for sn in xl.sheet_names}
    else:
        wb = openpyxl.load_workbook(path, data_only=True)
        df_map = {}
        for sn in wb.sheetnames:
            ws = wb[sn]
            rows = list(ws.iter_rows(values_only=True))
            if not rows:
                continue
            df_map[sn] = pd.DataFrame(rows[1:], columns=rows[0])

    for sn, df in df_map.items():
        if df.empty or len(df.columns) < 4:
            continue
        # 判定方向
        sn_norm = sn.strip()
        if sn_norm in ("调入", "换入", "新增"):
            direction = "add"
        elif sn_norm in ("调出", "换出", "删除"):
            direction = "remove"
        else:
            continue
        # 筛 target index
        col0 = df.iloc[:, 0].astype(str).str.strip()
        mask = col0 == target_index_code
        sub = df[mask]
        for _, row in sub.iterrows():
            code = _norm_code(row.iloc[2])
            name = str(row.iloc[3]).strip() if row.iloc[3] is not None else ""
            if not code:
                continue
            r = {"code": code, "code_name": name}
            if direction == "add":
                add_rows.append(r)
            else:
                remove_rows.append(r)
    return add_rows, remove_rows


def _extract_section_records(full_text: str, target_section: str) -> Tuple[List[dict], List[dict]]:
    """从纯文本中定位 target_section 调整名单 + 抽 4 列/2 列 records。
    PDF (pdfplumber extract_text) 和 HTML content (strip 标签后) 共用此逻辑。

    header_re 已泛化: 支持任意"中证XXX"/"沪深XXX"/"上证XXX"/"深证XXX"/"中证N位数字" 等指数简称。
    target_section 精确匹配 group(1)。
    """
    # 定位所有 section header (泛化: 数字代号 + 中文简称均可)
    header_re = re.compile(
        r"(沪深300|沪深[一-龥]{1,10}|"
        r"中证\d{1,4}|中证A\d+|中证[一-龥]{1,10}|"
        r"上证\d+|上证[一-龥]{1,10}|"
        r"深证\d+|深证[一-龥]{1,10}|"
        r"科创50|北证\d+|国证\d+|国证[一-龥]{1,10})"
        r"\s*指数样本调整名单[:：]?"
    )
    headers = list(header_re.finditer(full_text))
    if not headers:
        return [], []

    # 找 target section
    target_text = None
    for i, m in enumerate(headers):
        if m.group(1) == target_section:
            start = m.end()
            end = headers[i + 1].start() if i + 1 < len(headers) else len(full_text)
            target_text = full_text[start:end].strip()
            break
    if not target_text:
        return [], []

    # 解析 4 列 (调出代码 调出名称 调入代码 调入名称)
    add_rows, remove_rows = [], []
    in_data = False
    last_side = None  # 处理 2 列 (只有一侧) 的情况
    # HTML table 横向表头 ("调出|调入" 同行): 左侧 side 通常='remove', 右侧='add'
    # 用于:csindex content HTML 表格, 某行 4 列中右侧 2 列为 nbsp 空时, fallback 把左 2 列归到 left_side
    header_orientation = None  # 'horizontal' or None
    left_side = None
    for ln in target_text.split("\n"):
        ln = ln.strip()
        if not ln:
            continue
        # 表头行跳过 (含横向 "调出 调入" 同行表头, csindex HTML table 典型)
        if re.search(r"证券代码|股票代码|调出名单|调入名单|指数代码|换出|换入|"
                     r"调出\s+调入|调入\s+调出", ln):
            in_data = True
            if "调出名单" in ln and "调入名单" not in ln:
                last_side = "remove"
                header_orientation = None
            elif "调入名单" in ln and "调出名单" not in ln:
                last_side = "add"
                header_orientation = None
            elif re.search(r"调出.*调入", ln):
                # 横向表头 "调出 调入" (HTML table): 左 remove 右 add
                header_orientation = "horizontal"
                left_side = "remove"
                last_side = None
            elif re.search(r"调入.*调出", ln):
                header_orientation = "horizontal"
                left_side = "add"
                last_side = None
            else:
                last_side = None
            continue
        if not in_data:
            continue
        # 4 列 (调出代码 调出名称 调入代码 调入名称), 名称用 .+? 支持含空格 (如 "ST 锦化")
        m4 = re.match(r"^(\d{6})\s+(.+?)\s+(\d{6})\s+(.+)$", ln)
        if m4:
            rc, rn, ac, an = m4.group(1), m4.group(2), m4.group(3), m4.group(4)
            remove_rows.append({"code": rc, "code_name": rn})
            add_rows.append({"code": ac, "code_name": an})
            last_side = None
            continue
        # 2 列 (单侧)
        m2 = re.match(r"^(\d{6})\s+(\S+)", ln)
        if m2:
            r = {"code": m2.group(1), "code_name": m2.group(2)}
            if last_side == "add":
                add_rows.append(r)
            elif last_side == "remove":
                remove_rows.append(r)
            elif header_orientation == "horizontal" and left_side in ("add", "remove"):
                # HTML 横向表头下, 单独 2 列通常左侧 (右侧 nbsp 空)
                (add_rows if left_side == "add" else remove_rows).append(r)
    return add_rows, remove_rows


def parse_pdf_adjustments(path: Path, target_section: str = "中证1000") -> Tuple[List[dict], List[dict]]:
    """解析 PDF 的指定指数 section
    target_section: '中证1000' / '中证2000' / '中证全指' / '中证能源' 等
    返 (add_rows, remove_rows)
    """
    with pdfplumber.open(path) as pdf:
        full_text = "\n".join((p.extract_text() or "") for p in pdf.pages)
    return _extract_section_records(full_text, target_section)


def parse_content_adjustments(content_html: str, target_section: str) -> Tuple[List[dict], List[dict]]:
    """解析 csindex 公告 detail.content (HTML) 嵌入的调整名单表格。

    用途: csindex 早期公告 (如 2009-12-14 nid=1208 "中证行业指数调整名单") 无 PDF/xlsx 附件,
    调整名单直接以 HTML <table> 嵌入 content 字段。此函数把 HTML 转 text 后复用
    _extract_section_records 抽 add/remove。

    HTML → text 策略: </tr> → \n (每个 tr 一行) → strip 所有标签 → &nbsp; → space → 压缩 [ \t]+
    保留 \n 以让 _extract_section_records 按行扫描。
    """
    if not content_html:
        return [], []
    try:
        # csindex 老公告 HTML 无 </tr> 闭标签 (HTML5 隐式闭合), 用 <tr> 开标签作行边界
        # \x00 marker 防止 raw HTML 里 \n 缩进被压空白时丢行边界
        text = re.sub(r"(?i)<\s*tr[^>]*>", "\x00", content_html)
        # 去所有标签
        text = re.sub(r"<[^>]+>", " ", text)
        # HTML entity → space
        text = re.sub(r"(?i)&nbsp;", " ", text)
        # 压所有空白 (含 raw HTML \n 缩进) 为单个空格
        text = re.sub(r"\s+", " ", text)
        # \x00 marker 恢复为 \n (每个 tr 一行)
        text = text.replace("\x00", "\n")
        # 行内首尾空白
        text = "\n".join(ln.strip() for ln in text.split("\n"))
        return _extract_section_records(text, target_section)
    except Exception as e:
        log.error(f"    parse_content_adjustments ERR: {e}")
        return [], []


# ======================== 4. 主流程 ========================
def process_notice(nid: int, detail_cache: Path, file_cache: Path,
                   target_index_code: str, target_section: str) -> List[dict]:
    """处理单个公告 -> 返回 records 列表
    record: {updateDate, index_code, code, code_name, adjust_type, notice_id, source}

    流程:
      1. 取 detail
      2. 优先解析 PDF/xlsx 附件 (extract_file_urls + download_notice_files)
      3. 附件 records 空时 fallback: 解析 detail.content HTML 嵌入的调整名单表格
         (解锁 2009-2014 早期公告无附件但 content 内嵌 table 的情况)
    """
    detail = fetch_detail(nid, detail_cache)
    if not detail or not detail.get("data"):
        log.warning(f"  {nid}: detail empty")
        return []
    data = detail["data"]
    publish_date = (data.get("publishDate") or "").split("T")[0]
    title = data.get("title") or ""

    urls = extract_file_urls(data)
    files = download_notice_files(nid, publish_date, urls, file_cache)
    if not files:
        log.info(f"  {nid} {publish_date}: no attachments; will try content HTML fallback; title={title}")

    records = []
    for f in files:
        ext = f.suffix.lower()
        try:
            if ext == ".pdf":
                add, rem = parse_pdf_adjustments(f, target_section=target_section)
            elif ext in (".xlsx", ".xls"):
                add, rem = parse_xlsx_adjustments(f, target_index_code=target_index_code)
            else:
                continue
        except Exception as e:
            log.error(f"    parse {f.name} ERR: {e}")
            continue
        for r in add:
            records.append({
                "updateDate": publish_date, "index_code": target_index_code,
                "code": r["code"], "code_name": r["code_name"],
                "adjust_type": "add", "notice_id": nid, "source": f.name,
            })
        for r in rem:
            records.append({
                "updateDate": publish_date, "index_code": target_index_code,
                "code": r["code"], "code_name": r["code_name"],
                "adjust_type": "remove", "notice_id": nid, "source": f.name,
            })

    # Fallback: 附件空 或 附件解析 0 records 时, 解析 detail.content HTML 嵌入的表格
    if not records:
        content_html = data.get("content") or ""
        try:
            add, rem = parse_content_adjustments(content_html, target_section)
        except Exception as e:
            log.error(f"    content fallback {nid} ERR: {e}")
            add, rem = [], []
        if add or rem:
            log.info(f"    content HTML fallback hit: add={len(add)}, remove={len(rem)}")
            for r in add:
                records.append({
                    "updateDate": publish_date, "index_code": target_index_code,
                    "code": r["code"], "code_name": r["code_name"],
                    "adjust_type": "add", "notice_id": nid, "source": f"content_html_{nid}",
                })
            for r in rem:
                records.append({
                    "updateDate": publish_date, "index_code": target_index_code,
                    "code": r["code"], "code_name": r["code_name"],
                    "adjust_type": "remove", "notice_id": nid, "source": f"content_html_{nid}",
                })

    log.info(f"  {nid} {publish_date}: parsed add={sum(1 for r in records if r['adjust_type']=='add')}, "
             f"remove={sum(1 for r in records if r['adjust_type']=='remove')}, title={title}")
    return records


def fetch_akshare_current(index_code: str) -> List[dict]:
    """拉 akshare 当前快照 -> records(adjust_type='current')

    akshare 返列: 日期/指数代码/指数名称/指数英文名称/成分券代码/成分券名称/...
    成分券代码 = 倒数第 2 不是 iloc[0](=日期), 用 header 名定位稳健。
    """
    import akshare as ak
    log.info(f"拉 akshare 当前快照 index={index_code}...")
    df = ak.index_stock_cons_csindex(symbol=index_code)
    # header 定位
    cols = list(df.columns)
    code_col = next((c for c in cols if "成分券代码" in str(c) or "股票代码" in str(c) or "证券代码" in str(c)), None)
    name_col = next((c for c in cols if "成分券名称" in str(c) or "股票名称" in str(c) or "证券名称" in str(c)), None)
    if code_col is None:
        # 兜底: 第 5 列(成分券代码 实证位置)
        code_col = cols[4] if len(cols) > 4 else cols[0]
    if name_col is None and len(cols) > 5:
        name_col = cols[5]
    out = []
    for _, row in df.iterrows():
        raw = row[code_col]
        s = str(raw).strip()
        if not s or not s.isdigit():
            continue
        code = s.zfill(6)
        if len(code) != 6:
            continue
        name = str(row[name_col]).strip() if name_col else ""
        out.append({
            "updateDate": "current", "index_code": index_code,
            "code": code, "code_name": name,
            "adjust_type": "current", "notice_id": 0, "source": "akshare.index_stock_cons_csindex",
        })
    log.info(f"  akshare current {index_code}: {len(out)} stocks")
    return out


def parse_launch_xlsx(path: Path, index_code: str, publish_date: str, notice_id: int) -> List[dict]:
    """解析 launch xlsx (单 sheet, header 6 列: 指数代码/指数简称/指数英文简称/证券代码/证券中文简称/证券英文名称)

    用 header 定位"证券代码"和"证券中文简称"列, 不依赖固定 col index。
    """
    wb = openpyxl.load_workbook(path, data_only=True)
    out = []
    for sn in wb.sheetnames:
        ws = wb[sn]
        rows = list(ws.iter_rows(values_only=True))
        if not rows:
            continue
        header = rows[0]
        code_idx = next((i for i, h in enumerate(header) if h and "证券代码" in str(h)), 3)
        name_idx = next((i for i, h in enumerate(header) if h and "证券中文简称" in str(h)), 4)
        for row in rows[1:]:
            if not row:
                continue
            code = _norm_code(row[code_idx] if len(row) > code_idx else None)
            if not code:
                continue
            name = str(row[name_idx]).strip() if len(row) > name_idx and row[name_idx] else ""
            out.append({
                "updateDate": publish_date, "index_code": index_code,
                "code": code, "code_name": name,
                "adjust_type": "initial", "notice_id": notice_id, "source": path.name,
            })
        break  # 只用第一个 sheet
    return out


# ======================== 通用指数处理 (G1+G2 扩展) ========================
# 内置指数 -> PDF section 名映射 (用于 parse_pdf_adjustments 精确定位 section)
# 实证 csindex 公告系统 PDF section 用指数简称 (无"指数"后缀), 如 "中证1000 指数样本调整名单"
# 行业指数系列 section 名实证 (2026-07-28):
#   来源: csindex 公告 id=1208 (2009-12-14 "中证行业指数调整名单") HTML content 嵌入文本
#   method: regex `(中证[一-龥]{1,10}指数)样本调整名单` 对 content 去标签后文本匹配
#   实证 10 个 section 名: 中证能源/中证原材料/中证工业/中证可选消费/中证主要消费/
#                          中证医药卫生/中证金融地产/中证信息技术/中证电信/中证公用事业
#   注意: akshare 当前 000928 简称="800能源" (后期改名中证800行业系列),
#         但 csindex 公告 content PDF section 文本仍用"中证能源"等早期名称。
#   parse_pdf_adjustments header_re = `中证[一-龥]{1,10}\s*指数样本调整名单` 贪婪回溯后
#         group(1) = "中证能源" 等 (剔"指数"后缀), 故 MAP 值不带"指数"二字。
#   ⚠️ 已知限制: csindex 对 000928-000937 几乎无 PDF 调整公告 (fetch_notices_by_index
#         只 6 条, theme=指数调样 3 条均无附件), 此 MAP 主要供未来潜在 PDF / HTML content
#         解析扩展使用。当前 was_removed 治偏差增益来自 xlsx 单指数临时调整 (不经此 MAP)。
INDEX_SECTION_MAP: Dict[str, str] = {
    # 已有 (CSI 1000/2000 走 main --only 路径, 不走 --indices)
    "000852": "中证1000",
    "932000": "中证2000",
    # G1: 中证全指 (全市场池, 解锁策略 02)
    "000985": "中证全指",
    # G2: 中证一级行业指数 (行业轮动, 解锁策略 03) - 中证行业指数系列 (000928-000937)
    # section 名 1208 公告 HTML content 实证 (见上方注释)
    "000928": "中证能源",
    "000929": "中证原材料",
    "000930": "中证工业",
    "000931": "中证可选消费",
    "000932": "中证主要消费",
    "000933": "中证医药卫生",
    "000934": "中证金融地产",
    "000935": "中证信息技术",
    "000936": "中证电信",
    "000937": "中证公用事业",
    # 000938 不是行业指数 (中证民企ESG 50 等), 留作可选
}

# 多指数合并公告 -> 共享给哪些 index_code 的映射
# 实证 (2026-07-28): csindex API 的 indexCode 字段对部分 code 绑定不全
# 例: 公告 id=1208 "中证行业指数调整名单" (2009-12-14) 涵盖 000928-000937 全部 10 个行业指数,
#     但 fetch_notices_by_index("000936") 返回的列表里没 1208 (csindex 后台绑定漏),
#     需在此显式补入, 否则 000936 治偏差完全失效。
SHARED_NOTICES_BY_INDEX: Dict[int, List[str]] = {
    1208: ["000928", "000929", "000930", "000931", "000932",
           "000933", "000934", "000935", "000936", "000937"],
}


def process_generic_index(
    idx_code: str,
    target_section: str,
    out_dir: Path,
    cache_dir: Path,
    refresh_index: bool = False,
    max_notices: int = 0,
) -> pd.DataFrame:
    """通用指数处理流程 (复用 CSI 1000 的 4 步逻辑, 幂等可重跑):

    1. fetch_notices_by_index(idx_code) 拉公告列表
    2. filter_adjustment_notices 筛 theme=指数调样
    3. process_notice 解析 PDF/xlsx (target_section 定位 PDF section)
       - max_notices > 0 时, 只处理前 N 条 (按 publishDate 升序, 最近的 N 条),
         避免 000985 等公告多的指数跑数小时 (PDF section 多不匹配, 治偏差增益微小)
    4. fetch_akshare_current 加 current 行 (兜底治偏差: 即使无公告也有当前快照)

    返回 announce_union DataFrame (也写 {idx_code}_announce_union.parquet)
    """
    detail_cache = cache_dir / "detail"
    file_cache = cache_dir / "files"

    log.info(f"\n========== generic index: {idx_code} (section={target_section}) ==========")

    # 1. indexCode 公告列表
    notices_idx = fetch_notices_by_index(
        idx_code, cache_dir / f"notices_{idx_code}.json", force=refresh_index,
    )
    adj = filter_adjustment_notices(notices_idx)
    # 补共享的多指数合并公告 (csindex indexCode 字段对部分 code 绑定不全, 见 SHARED_NOTICES_BY_INDEX)
    existing_ids = {x["id"] for x in adj}
    for shared_nid, codes in SHARED_NOTICES_BY_INDEX.items():
        if idx_code in codes and shared_nid not in existing_ids:
            adj.insert(0, {"id": shared_nid})
            log.info(f"  补共享公告 nid={shared_nid} (csindex indexCode 未绑定)")
    log.info(f"  indexCode={idx_code} theme=指数调样 notices: {len(adj)}")

    # 2. 逐公告解析 (PDF/xlsx), max_notices 截断
    if max_notices > 0 and len(adj) > max_notices:
        # 保留最近 max_notices 条 (adj 已按 publishDate 升序, 取末尾即最新)
        adj_trimmed = adj[-max_notices:]
        log.info(f"  max_notices={max_notices}: 截断 {len(adj)} -> {len(adj_trimmed)} (保留最新)")
        adj = adj_trimmed

    records: List[dict] = []
    for x in adj:
        nid = x["id"]
        try:
            recs = process_notice(nid, detail_cache, file_cache, idx_code, target_section)
            records.extend(recs)
        except Exception as e:
            log.error(f"  {idx_code} nid={nid} FAILED: {e}")
    log.info(f"  after notices: {len(records)} records "
             f"(add={sum(1 for r in records if r['adjust_type']=='add')}, "
             f"remove={sum(1 for r in records if r['adjust_type']=='remove')})")

    # 3. akshare 当前快照 (兜底, 保证至少有 current 行)
    try:
        records.extend(fetch_akshare_current(idx_code))
    except Exception as e:
        log.warning(f"  akshare {idx_code} current 失败 (非致命): {e}")

    # 4. 写 parquet
    df = pd.DataFrame(records, columns=[
        "updateDate", "index_code", "code", "code_name", "adjust_type", "notice_id", "source"])
    out_path = out_dir / f"{idx_code}_announce_union.parquet"
    df.to_parquet(out_path, index=False)
    distinct = df["code"].nunique() if len(df) else 0
    n_add = (df["adjust_type"] == "add").sum() if len(df) else 0
    n_rem = (df["adjust_type"] == "remove").sum() if len(df) else 0
    n_cur = (df["adjust_type"] == "current").sum() if len(df) else 0
    log.info(f"  {idx_code} result: {out_path.name} total={len(df)} "
             f"(add={n_add}, remove={n_rem}, current={n_cur}, distinct={distinct})")
    return df


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--out-dir", default=str(OUT_DIR))
    parser.add_argument("--cache-dir", default=str(CACHE_DIR))
    parser.add_argument("--refresh-list", action="store_true", help="强制重新拉全量公告列表")
    parser.add_argument("--refresh-index", action="store_true", help="强制重新拉 indexCode 公告列表")
    parser.add_argument("--only", choices=["1000", "2000", "both"], default="both")
    parser.add_argument(
        "--indices",
        help="通用模式: 逗号分隔的指数代码 (如 000985,000928~000937), "
             "用 INDEX_SECTION_MAP 处理 (代替默认 1000/2000 硬编码路径)",
    )
    parser.add_argument(
        "--max-notices-per-index", type=int, default=0,
        help="通用模式: 每指数最多处理的公告数 (0=不限, >0 取最新 N 条). "
             "避免 000985 公告多的指数跑数小时",
    )
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    cache_dir = Path(args.cache_dir)
    detail_cache = cache_dir / "detail"
    file_cache = cache_dir / "files"

    # ===== 通用模式: --indices =====
    if args.indices:
        codes = [c.strip() for c in args.indices.split(",") if c.strip()]
        log.info(f"通用模式: indices={codes}, max_notices_per_index={args.max_notices_per_index}")
        for code in codes:
            section = INDEX_SECTION_MAP.get(code)
            if not section:
                log.warning(f"  {code} 不在 INDEX_SECTION_MAP, 跳过 (请补充映射)")
                continue
            try:
                process_generic_index(
                    code, section, out_dir, cache_dir,
                    refresh_index=args.refresh_index,
                    max_notices=args.max_notices_per_index,
                )
            except Exception as e:
                log.error(f"  {code} 通用处理 FAILED: {e}")
        # 总结
        log.info("\n========== DONE (generic) ==========")
        log.info(f"OUT_DIR: {out_dir.resolve()}")
        for f in sorted(out_dir.glob("*_announce_union.parquet")):
            log.info(f"  {f.name}: {f.stat().st_size} bytes")
        return

    # ===== 默认模式: CSI 1000/2000 (硬编码, 不动) =====
    # 1. 拉/缓存全量列表(兜底) + indexCode 精准列表(主源)
    notices = fetch_all_notices(cache_dir / "all_notices.json", force=args.refresh_list)
    notices_000852 = fetch_notices_by_index(
        "000852", cache_dir / "notices_000852.json", force=args.refresh_index,
    )

    # ========== CSI 1000 ==========
    if args.only in ("1000", "both"):
        log.info("\n========== CSI 1000 (000852) ==========")
        # 主源: indexCode=000852 精准拉, theme=指数调样
        adj_idx = filter_adjustment_notices(notices_000852)
        log.info(f"indexCode=000852 调整公告: {len(adj_idx)}")
        # 兜底: 旧 title 过滤(防 indexCode 接口变动)
        filtered_old = filter_csi1000_notices(notices)
        log.info(f"title 过滤兜底: {len(filtered_old)}")

        all_ids = (
            {x["id"] for x in adj_idx}
            | {x["id"] for x in filtered_old}
            | set(CSI1000_REGULAR_IDS)
            | set(CSI1000_TEMP_IDS)
        )
        all_ids = sorted(all_ids)
        log.info(f"total CSI 1000 notice ids to process: {len(all_ids)}")

        records_1000 = []
        for nid in all_ids:
            try:
                recs = process_notice(nid, detail_cache, file_cache, "000852", "中证1000")
                records_1000.extend(recs)
            except Exception as e:
                log.error(f"  {nid} FAILED: {e}")

        # 加 current (akshare 现快照, header 定位列, 修 iloc[0] bug)
        try:
            records_1000.extend(fetch_akshare_current("000852"))
        except Exception as e:
            log.warning(f"akshare 000852 当前快照失败 (非致命): {e}")

        df_1000 = pd.DataFrame(records_1000, columns=[
            "updateDate", "index_code", "code", "code_name", "adjust_type", "notice_id", "source"])
        out_1000 = out_dir / "000852_announce_union.parquet"
        df_1000.to_parquet(out_1000, index=False)
        distinct_1000 = df_1000["code"].nunique()
        n_add = (df_1000["adjust_type"] == "add").sum()
        n_rem = (df_1000["adjust_type"] == "remove").sum()
        n_cur = (df_1000["adjust_type"] == "current").sum()
        log.info(f"\nCSI 1000 result: {out_1000}")
        log.info(f"  total records: {len(df_1000)} (add={n_add}, remove={n_rem}, current={n_cur})")
        log.info(f"  distinct codes (ever-included): {distinct_1000}")
        log.info(f"  date range: {df_1000['updateDate'].min()} ~ {df_1000['updateDate'].max()}")
        log.info(f"  distinct notice_ids: {df_1000['notice_id'].nunique()}")

    # ========== CSI 2000 ==========
    if args.only in ("2000", "both"):
        log.info("\n========== CSI 2000 (932000) ==========")
        log.warning("GAP: csindex 公告系统未发布任何 CSI 2000 样本调整公告")
        log.warning("     仅可从 2023-08-10 launch xlsx 获取初始 2000 只样本")
        records_2000 = []

        # launch xlsx (initial) — header 定位列, 修 row[0] bug
        try:
            detail = fetch_detail(CSI2000_LAUNCH_ID, detail_cache)
            if detail.get("data"):
                data = detail["data"]
                publish_date = (data.get("publishDate") or "2023-08-10").split("T")[0]
                urls = extract_file_urls(data)
                files = download_notice_files(CSI2000_LAUNCH_ID, publish_date, urls, file_cache)
                for f in files:
                    if f.suffix.lower() != ".xlsx":
                        continue
                    recs = parse_launch_xlsx(f, "932000", publish_date, CSI2000_LAUNCH_ID)
                    records_2000.extend(recs)
                    log.info(f"  launch xlsx {f.name}: {len(recs)} initial stocks")
                    break
        except Exception as e:
            log.error(f"CSI 2000 launch xlsx ERR: {e}")

        # current snapshot via akshare (header 定位列)
        try:
            records_2000.extend(fetch_akshare_current("932000"))
        except Exception as e:
            log.warning(f"akshare CSI 2000 当前快照失败: {e}")

        df_2000 = pd.DataFrame(records_2000, columns=[
            "updateDate", "index_code", "code", "code_name", "adjust_type", "notice_id", "source"])
        out_2000 = out_dir / "932000_announce_union.parquet"
        df_2000.to_parquet(out_2000, index=False)
        distinct_2000 = df_2000["code"].nunique()
        n_init = (df_2000["adjust_type"] == "initial").sum()
        n_cur = (df_2000["adjust_type"] == "current").sum()
        log.info(f"\nCSI 2000 result: {out_2000}")
        log.info(f"  total records: {len(df_2000)} (initial={n_init}, current={n_cur})")
        log.info(f"  distinct codes: {distinct_2000}")
        log.info(f"  GAP: 无调整公告,初始集 ∪ 当前集 (中间调整无记录)")

    # 总结
    log.info("\n========== DONE ==========")
    log.info(f"OUT_DIR: {out_dir.resolve()}")
    for f in sorted(out_dir.glob("*_announce_union.parquet")):
        log.info(f"  {f.name}: {f.stat().st_size} bytes")


if __name__ == "__main__":
    sys.exit(main())