#!/usr/bin/env python
# coding=utf-8
"""根据表名从 DataWorks 获取生产代码，保存到 production_code/ 目录。

用法：
    fetch-table-code loghubods.dwd_recsys_alg_exposure_base_20250108
    fetch-table-code loghubods.dwd_recsys_alg_exposure_base_20250108 --force
    fetch-table-code loghubods.dwd_recsys_alg_exposure_base_20250108 --recursive
    fetch-table-code loghubods.dwd_recsys_alg_exposure_base_20250108 --recursive --depth 5
"""

import os
import sys
import argparse


def main():
    parser = argparse.ArgumentParser(description="获取表的 DataWorks 生产代码")
    parser.add_argument("table_name", help="表名，格式: project.table 或 table")
    parser.add_argument("--force", action="store_true", help="跳过缓存，强制从 API 拉取")
    parser.add_argument("--recursive", "-r", action="store_true", help="递归获取所有上游表的代码")
    parser.add_argument("--depth", type=int, default=3, help="递归最大深度（默认 3）")
    parser.add_argument("--output-dir", type=str, default=None,
                        help="产出基础目录（默认 $ODPS_SKILLS_DATA_DIR 或 cwd）")
    args = parser.parse_args()

    # 设置 output-dir（需在 import dataworks 之前，因为模块级会读取 env）
    if args.output_dir:
        os.environ["ODPS_SKILLS_DATA_DIR"] = args.output_dir

    from odps_skills.dataworks import DataWorksClient

    dw = DataWorksClient()

    if args.recursive:
        dw.get_node_code_recursive(args.table_name, max_depth=args.depth, force=args.force)
    else:
        results = dw.get_node_code(args.table_name, force=args.force)
        if not results:
            print(f"未找到 '{args.table_name}' 的生产代码")
            sys.exit(1)
        for r in results:
            print(f"任务: {r['task_name']}  代码长度: {len(r['content'])} chars")


if __name__ == "__main__":
    main()
