Agent/examples/auto_put_ad_mini/setup_region_dropdown.py

"""给飞书「自动化账户」表地域列设置下拉多选（选项来自 regions_all.json）。

用法:
    cd examples/auto_put_ad_mini
    python3 setup_region_dropdown.py             # 设置下拉（幂等，可重复跑）
    python3 setup_region_dropdown.py --dry-run   # 只打印选项，不调飞书

选项集(顺序):
    1. 不限(全国) / 排地域(回退投放模板) / 中国未知(156)
    2. 其余省级/直辖市/特别行政区按拼音排序(依赖 pypinyin,缺失回退 Unicode)
    排除: 市级/区级、国外(9000)、中国(1156)

说明:
    本脚本刻意保持自包含(仅依赖 httpx + python-dotenv)，不引入 agent 框架，
    方便在任意环境独立执行。token 获取逻辑与 tools/feishu_doc 保持一致。
"""

from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path

import httpx
from dotenv import load_dotenv

# 拼音排序依赖（可选）：未安装时回退到 Unicode 排序
try:
    from pypinyin import lazy_pinyin as _lazy_pinyin
except ImportError:
    _lazy_pinyin = None

_HERE = Path(__file__).parent
load_dotenv(_HERE / ".env")

# 飞书「自动化账户」表
DEFAULT_SPREADSHEET_TOKEN = "D8f4sKEb2hfBnNttX1ucZTklnHf"
DEFAULT_SHEET_ID = "y3uCcz"
DEFAULT_RANGE = "A1:N300"

FEISHU_BASE_URL = "https://open.feishu.cn/open-apis"
_REGION_DICTIONARY_PATH = (
    _HERE / "data" / "tencent_constants" / "regions_all.json"
)

# 顶层节点中排除的 id: 国外(9000)、中国(1156,整包节点,与"不限"语义重叠)
_EXCLUDED_TOP_IDS = {9000, 1156}
# 固定前导项，保持 parse_geo_location_config 语义（不限→全国 / 排地域→模板），
# 中国未知(156)紧随其后，其余省级按名称排序。
_SPECIAL_OPTIONS = ["不限", "排地域", "中国未知"]
_EXCLUDED_NAMES = {"中国未知"}
# 下拉选项胶囊颜色（与 conditionValues 一一对应；highlightValidData=true 时飞书强制要求）
_COLOR_PALETTE = [
    "#1FB6C1", "#F006C2", "#FB16C3", "#FFB6C1", "#A6D6FF",
    "#FFB84D", "#4CD3A0", "#8A7DFF", "#FF8FA3", "#5B8DEF",
    "#F7B500", "#00C2A8", "#B08CF9", "#FF7A45", "#36CFC9",
    "#597EF7", "#F759AB", "#7CB305", "#FAAD14", "#13C2C2",
]
# 表头在第 2 行，数据从第 3 行开始
HEADER_ROW_INDEX = 1
DATA_START_INDEX = 2
# 地域列下拉覆盖到第 300 行（与 sync 的 DEFAULT_RANGE A1:N300 对齐）
MAX_ROW = 300

# tenant token 缓存
_token_cache: dict = {}
_TOKEN_TTL = 1500  # 25 分钟


def _get_tenant_token() -> str:
    """获取飞书 tenant_access_token（带缓存，逻辑同 tools.feishu_doc）。"""
    cached = _token_cache.get("tenant")
    if cached and time.time() - cached["ts"] < _TOKEN_TTL:
        return cached["token"]
    app_id = os.getenv("FEISHU_APP_ID", "cli_a955e97067f85cb3")
    app_secret = os.getenv("FEISHU_APP_SECRET", "NQaG4ci1plXRDTgwCqrLJgMLLoA2tdF8")
    resp = httpx.post(
        f"{FEISHU_BASE_URL}/auth/v3/tenant_access_token/internal",
        json={"app_id": app_id, "app_secret": app_secret},
        timeout=30,
    )
    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"获取飞书 token 失败: {data}")
    token = data["tenant_access_token"]
    _token_cache["tenant"] = {"token": token, "ts": time.time()}
    return token


def _auth_headers(token: str) -> dict:
    return {"Authorization": f"Bearer {token}"}


def _pinyin_sort_key(name: str) -> str:
    """按拼音排序键；pypinyin 缺失时回退 Unicode 码位排序。"""
    if _lazy_pinyin is not None:
        return "".join(_lazy_pinyin(name))
    return name


def build_region_options() -> list[str]:
    """构建下拉选项：不限、排地域、中国未知固定在前，其余省级按拼音排序。"""
    regions = json.loads(_REGION_DICTIONARY_PATH.read_text(encoding="utf-8"))
    names: list[str] = []
    for r in regions:
        if r.get("parent_id") != 0:
            continue
        rid = int(r["id"])
        if rid in _EXCLUDED_TOP_IDS:
            continue
        name = str(r.get("name") or "").strip()
        if not name or name in _EXCLUDED_NAMES:
            continue
        names.append(name)
    names.sort(key=_pinyin_sort_key)
    return _SPECIAL_OPTIONS + names


def _column_letter(idx: int) -> str:
    """0-based 列索引转列字母（A→0，Z→25，AA→26）。"""
    letters = ""
    idx += 1
    while idx:
        idx, rem = divmod(idx - 1, 26)
        letters = chr(65 + rem) + letters
    return letters


def _cell_text(value) -> str:
    """飞书单元格值（可能是富文本 list/dict）转纯文本。"""
    if value is None:
        return ""
    if isinstance(value, list):
        parts = []
        for item in value:
            if isinstance(item, dict):
                parts.append(str(item.get("text") or item.get("link") or ""))
            else:
                parts.append(str(item))
        return "".join(parts).strip()
    if isinstance(value, dict):
        return str(value.get("text") or value.get("link") or "").strip()
    return str(value).strip()


def _get_sheet_row_count(spreadsheet_token: str, sheet_id: str) -> int:
    """查询 sheet 的实际网格行数（数据验证 range 不能超出网格）。"""
    token = _get_tenant_token()
    resp = httpx.get(
        f"{FEISHU_BASE_URL}/sheets/v2/spreadsheets/{spreadsheet_token}/metainfo",
        headers=_auth_headers(token),
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"读取表格元数据失败: {json.dumps(data, ensure_ascii=False)[:500]}")
    for s in (data.get("data") or {}).get("sheets") or []:
        if s.get("sheetId") == sheet_id:
            return int(s.get("rowCount") or 0)
    raise RuntimeError(f"未找到 sheet: {sheet_id}")


def _find_geo_column_letter(
    spreadsheet_token: str, sheet_id: str, cell_range: str
) -> str:
    """读取表头，返回地域列字母（优先「地域」，回退「投放地域」）。"""
    token = _get_tenant_token()
    url = (
        f"{FEISHU_BASE_URL}/sheets/v2/spreadsheets/{spreadsheet_token}/values/"
        f"{sheet_id}!{cell_range}"
    )
    resp = httpx.get(url, headers=_auth_headers(token), timeout=30)
    resp.raise_for_status()
    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"读取飞书配置表失败: {json.dumps(data, ensure_ascii=False)[:500]}")
    values = (data.get("data") or {}).get("valueRange") or {}
    rows = values.get("values") or []
    if len(rows) <= HEADER_ROW_INDEX:
        raise RuntimeError("飞书配置表缺少表头行")
    headers = [_cell_text(v) for v in rows[HEADER_ROW_INDEX]]
    mapping = {name: idx for idx, name in enumerate(headers) if name}
    col = "地域" if "地域" in mapping else "投放地域"
    if col not in mapping:
        raise RuntimeError(f"飞书表缺少地域列; 当前表头={headers}")
    return _column_letter(mapping[col])


def _set_dropdown(
    spreadsheet_token: str,
    sheet_id: str,
    col_letter: str,
    options: list[str],
    row_count: int,
    dry_run: bool = False,
) -> dict:
    """幂等设置下拉：PUT 更新接口同时支持创建/更新，直接覆盖目标 range。"""
    end_row = min(max(row_count, DATA_START_INDEX + 1), MAX_ROW)
    data_range = f"{sheet_id}!{col_letter}{DATA_START_INDEX + 1}:{col_letter}{end_row}"

    if dry_run:
        return {"range": data_range, "options": len(options), "dry_run": True}

    token = _get_tenant_token()
    headers = _auth_headers(token)
    url = (
        f"{FEISHU_BASE_URL}/sheets/v2/spreadsheets/{spreadsheet_token}"
        f"/dataValidation/{sheet_id}"
    )

    # highlightValidData=true 时 colors 必须与选项一一对应
    colors = [_COLOR_PALETTE[i % len(_COLOR_PALETTE)] for i in range(len(options))]
    body = {
        "ranges": [data_range],
        "dataValidationType": "list",
        "dataValidation": {
            "conditionValues": options,
            "options": {
                "multipleValues": True,
                "highlightValidData": True,
                "colors": colors,
            },
        },
    }
    resp = httpx.put(
        url,
        headers={**headers, "Content-Type": "application/json"},
        json=body,
        timeout=30,
    )
    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"设置下拉失败: {json.dumps(data, ensure_ascii=False)[:500]}")
    return {"range": data_range, "options": len(options), "dry_run": False}


def main() -> int:
    parser = argparse.ArgumentParser(description="给飞书「自动化账户」表地域列设置下拉多选")
    parser.add_argument("--spreadsheet-token", default=DEFAULT_SPREADSHEET_TOKEN)
    parser.add_argument("--sheet-id", default=DEFAULT_SHEET_ID)
    parser.add_argument("--range", default=DEFAULT_RANGE)
    parser.add_argument("--dry-run", action="store_true", help="只打印选项，不调用飞书 API")
    args = parser.parse_args()

    options = build_region_options()
    print(f"下拉选项共 {len(options)} 个:")
    for i, name in enumerate(options, start=1):
        print(f"  {i:2d}. {name}")

    col_letter = _find_geo_column_letter(
        args.spreadsheet_token, args.sheet_id, args.range
    )
    print(f"地域列: {col_letter}")

    row_count = _get_sheet_row_count(args.spreadsheet_token, args.sheet_id)
    print(f"表格行数: {row_count}")

    result = _set_dropdown(
        args.spreadsheet_token,
        args.sheet_id,
        col_letter,
        options,
        row_count=row_count,
        dry_run=args.dry_run,
    )
    if result.get("dry_run"):
        print("(dry-run) 未调用飞书 API。正式设置请去掉 --dry-run 重新运行。")
    else:
        print(
            f"✅ 下拉已设置: range={result['range']} 选项数={result['options']}"
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())