#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
fast_ray_executor.py —— data-juicer Ray 模式「惰性 DAG」执行器

已验证：253s → 7 个算子 0.02s 挂完（原 8 次全链重跑）

【两个必须同时修的点】
1) ray_dataset.py:159/196 调 self.data.columns() → Ray Data 内部以 limit(1)
   执行整条 DAG（O(N^2)）。用静态列名覆盖 columns() 解决。
2) 覆盖 columns() 会让 DJ "以为" meta/stats 已存在而跳过补列，
   算子却要读 sample['__dj__meta__'] → KeyError。
   所以必须【自己正确地把这两列补上】。

【补列陷阱】pyarrow 无法从 [{}...] 推断 struct 类型（长度为 0）：
   ✗ table.append_column(col, [{} for _ in range(len(table))])
   ✓ table.append_column(col, pa.array([{}]*n, type=pa.struct([])))

用法: python fast_ray_executor.py ./video_pipeline_ray.yaml
"""
import os, sys, time, glob, json
import ray
import pyarrow as pa


def peek_columns(path):
    """支持目录：扫描所有 jsonl/json 的第一行，取列名并集。"""
    cols = set()
    if os.path.isdir(path):
        files = []
        for ext in ("*.jsonl", "*.json"):
            files.extend(sorted(glob.glob(os.path.join(path, "**", ext), recursive=True)))
    else:
        files = [path]

    if not files:
        print(f"[warn] 目录下没找到 jsonl/json: {path}")
        return cols

    for fp in files[:50]:                      # 最多扫 50 个，足够推断 schema
        try:
            with open(fp, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    if line.startswith("["):   # 兼容 json 数组
                        line = line.lstrip("[").rstrip(",").rstrip("]")
                    try:
                        cols.update(json.loads(line).keys())
                    except Exception:
                        pass
                    break                      # 每个文件只读第一行
        except Exception as e:
            print(f"[warn] 跳过 {fp}: {e}")
    print(f"[列名] 扫描 {min(len(files),50)}/{len(files)} 个文件 -> {sorted(cols)}")
    return cols


def build_lazy_process(state, batch_size):
    """state: dict(known=set(), ...)  用可变容器让 known 跨调用共享"""
    known = state["known"]

    def make_add_col(col):
        def _fn(table):
            if col in table.column_names:
                return table
            n = table.num_rows
            # 关键：必须显式 pa.struct([])，否则 pyarrow 推断出长度 0
            arr = pa.array([{} for _ in range(n)], type=pa.struct([]))
            return table.append_column(col, arr)
        return _fn

    def lazy_process(self, operators, *, exporter=None, checkpointer=None, tracer=None):
        if operators is None:
            return self
        if not isinstance(operators, list):
            operators = [operators]

        from data_juicer.utils.process_utils import calculate_ray_np
        from data_juicer.utils.constant import Fields

        calculate_ray_np(operators)

        # ---- 先补 meta / stats 两列（惰性，不触发执行）----
        for col in (Fields.meta, Fields.stats):
            self.data = self.data.map_batches(
                make_add_col(col), batch_format="pyarrow", batch_size=batch_size
            )
            known.add(col)
        print(f"[预补列] {sorted(known)}")

        # ---- 覆盖 columns()：返回已知列名，杜绝 schema 推理触发执行 ----
        orig_columns = ray.data.Dataset.columns
        ray.data.Dataset.columns = lambda ds: list(known)

        try:
            total = len(operators)
            t_all = time.time()
            print(f"\n[惰性 DAG] 开始挂接 {total} 个算子（每项应为毫秒级）")

            for i, op in enumerate(operators, 1):
                name = type(op).__name__
                try:
                    op.num_proc = 1      # concurrency=1 → actor 池=1 → 模型只加载一次
                except Exception:
                    pass

                t = time.time()
                self._run_single_op(op)  # Ray 分支：map_batches，不是 op.run
                cost = time.time() - t

                flag = "" if cost < 1.0 else "   <-- 异常! 此处触发了 action"
                print(f"  [{i}/{total}] {name:<42s} {cost * 1000:8.1f} ms{flag}")

                if not isinstance(self.data, ray.data.Dataset):
                    raise RuntimeError(
                        f"算子 [{name}] 把 ray.data.Dataset 转成了 {type(self.data).__name__}"
                    )

            print(f"[惰性 DAG] 构建总耗时: {time.time() - t_all:.2f}s（应接近 0）")
        finally:
            ray.data.Dataset.columns = orig_columns   # 还原，导出时用真实 schema

        return self

    return lazy_process


def main():
    cfg_path = sys.argv[1] if len(sys.argv) > 1 else "./video_pipeline_ray.yaml"

    from data_juicer.config import init_configs
    cfg = init_configs(args=["--config", cfg_path])

    print("=" * 72)
    print(f"  配置文件      : {cfg_path}")
    print(f"  executor_type : {cfg.executor_type}")
    print(f"  export_path   : {cfg.export_path}")
    print("=" * 72)

    ray.init(address=getattr(cfg, "ray_address", "auto"), ignore_reinit_error=True)
    print(f"[Ray] 已连接: {ray.get_runtime_context().gcs_address}")

    from data_juicer.core.data import ray_dataset as rd_mod

    RayDataset = rd_mod.RayDataset
    batch_size = getattr(rd_mod, "DEFAULT_BATCH_SIZE", 1024)

    state = {"known": peek_columns(cfg.dataset_path)}
    RayDataset.process = build_lazy_process(state, batch_size)

    from data_juicer.core import ExecutorFactory
    executor_cls = ExecutorFactory.create_executor(cfg.executor_type)
    print(f"[执行器] {executor_cls.__name__}")
    assert executor_cls.__name__ == "RayExecutor", "必须拿到 RayExecutor"

    export_path = cfg.export_path
    # 清掉上一轮遗留的分片（否则会被误合并进最终结果）
    if os.path.isdir(export_path):
        for f in glob.glob(os.path.join(export_path, "*.json")):
            os.remove(f)
    elif os.path.exists(export_path):
        os.remove(export_path)          # Ray 的 write_json 是目录语义
    os.makedirs(export_path, exist_ok=True)

    executor = executor_cls(cfg)
    t0 = time.time()
    executor.run()
    print(f"\n总耗时: {time.time() - t0:.1f}s")

    shards = sorted(glob.glob(os.path.join(export_path, "*.json")))
    if shards:
        merged = export_path.rstrip("/") + ".jsonl"
        n = 0
        with open(merged, "w", encoding="utf-8") as out:
            for s in shards:
                for line in open(s, encoding="utf-8"):
                    line = line.strip()
                    if line:
                        out.write(line + "\n")
                        n += 1
        print(f"\n已合并 {len(shards)} 个分片 -> {merged}（{n} 条样本）")
        for line in open(merged, encoding="utf-8"):
            d = json.loads(line)
            vp = d.get("video_path", ["?"])[0].split("/")[-1]
            print(f"  - {vp}")
            print(f"    text : {d.get('text', '')[:70]}")
            print(f"    meta : {d.get('__dj__meta__', {})}")
            print(f"    stats: {d.get('__dj__stats__', {})}")


if __name__ == "__main__":
    main()