import inspect
import json
from typing import Dict, List, Any
from data_juicer.ops.base_op import OPERATORS

# ========= 配置区 =========
# 1. 基类公共参数：所有Mapper/Filter通用，独立公共面板
BASE_COMMON_PARAMS = {
    "video_key": {"type": "str", "default": "video_path", "desc": "样本中视频列表字段名"},
    "audio_key": {"type": "str", "default": "audios", "desc": "样本中音频列表字段名"},
    "text_key": {"type": "str", "default": "text", "desc": "样本文本字段名"},
}

# 2. 高级运行时参数：全局自动注入，默认折叠，仅高级用户修改
ADVANCED_RUNTIME_PARAMS = {
    "accelerator": {"type": "str", "default": None, "desc": "加速设备 cpu/gpu"},
    "num_proc": {"type": "int", "default": 1, "desc": "单算子并行进程数"},
    "cpu_required": {"type": "float", "default": None, "desc": "Ray每个worker需要CPU资源"},
    "gpu_required": {"type": "float", "default": None, "desc": "Ray每个worker需要GPU资源"},
    "mem_required": {"type": "str", "default": None, "desc": "内存资源限制"},
    "work_dir": {"type": "str", "default": None, "desc": "算子工作目录"},
    "skip_op_error": {"type": "bool", "default": False, "desc": "是否跳过算子执行异常"},
    "turbo": {"type": "bool", "default": False, "desc": "turbo加速开关"},
    "history_key": {"type": "str", "default": "history", "desc": "对话历史字段"},
    "prompt_key": {"type": "str", "default": "prompt", "desc": "prompt字段"},
    "query_key": {"type": "str", "default": "query", "desc": "query字段"},
    "response_key": {"type": "str", "default": "response", "desc": "response字段"},
    "instruction_key": {"type": "str", "default": "instruction", "desc": "指令字段"},
    "index_key": {"type": "str", "default": None, "desc": "索引字段"},
}

# 3. 已知枚举字典（DJ1.4.4 视频算子，可持续扩充）
OP_ENUM_MAP: Dict[str, Dict[str, List[Any]]] = {
    "video_extract_frames_mapper": {
        "output_format": ["path", "bytes"],
        "frame_sampling_method": ["all_keyframes", "uniform"],
    },
    "video_duration_filter": {
        "any_or_all": ["any", "all"],
    },
    "video_aesthetics_filter": {
        "reduce_mode": ["avg", "max", "min"],
        "any_or_all": ["any", "all"],
    },
    "video_frames_text_similarity_filter": {
        "reduce_mode": ["avg", "max", "min"],
        "any_or_all": ["any", "all"],
    }
}

# ======================
def get_op_param_schema(op_name: str, op_cls) -> Dict[str, Any]:
    """获取单个算子的完整schema"""
    sig = inspect.signature(op_cls.__init__)
    native_params = []
    enum_map = OP_ENUM_MAP.get(op_name, {})

    for param_name, param in sig.parameters.items():
        # 过滤 self, *args, **kwargs
        if param_name in ("self", "args", "kwargs"):
            continue

        default_val = param.default
        if default_val is inspect.Parameter.empty:
            default_str = "required"
        else:
            default_str = repr(default_val)

        anno = param.annotation
        enum_list = enum_map.get(param_name, None)

        native_params.append({
            "name": param_name,
            "annotation": str(anno),
            "default": default_str,
            "enum": enum_list,
            "ui_type": "select" if enum_list else "input",
            "description": ""
        })

    op_info = {
        "op_name": op_name,
        "class_name": op_cls.__name__,
        "op_category": op_cls.__bases__[0].__name__,
        "accelerator": getattr(op_cls, "_accelerator", "cpu"),
        "is_batched_op": getattr(op_cls, "_batched_op", False),
        "docstring": op_cls.__doc__.strip().splitlines()[0] if op_cls.__doc__ else "",
        "native_params": native_params,  # 主表单：算子原生参数
        "base_common_params": BASE_COMMON_PARAMS,  # 公共面板
        "advanced_runtime_params": ADVANCED_RUNTIME_PARAMS  # 高级折叠面板
    }
    return op_info


def dump_all_ops_schema(save_json_path: str = "dj_ops_schema.json"):
    """遍历全部算子，生成schema并保存json"""
    # 加载全部算子，保证懒加载算子注册
    from data_juicer.ops import load
    all_ops = []
    for op_name, op_cls in sorted(OPERATORS.modules.items()):
        schema = get_op_param_schema(op_name, op_cls)
        all_ops.append(schema)

    final_payload = {
        "version": "data-juicer 1.4.4",
        "total_ops_count": len(all_ops),
        "operators": all_ops
    }
    # 写入json
    with open(save_json_path, "w", encoding="utf-8") as f:
        json.dump(final_payload, f, ensure_ascii=False, indent=2)
    print(f"✅ Schema saved to {save_json_path}, total operators: {len(all_ops)}")
    return final_payload


def filter_op_by_keyword(schema_data, keyword: str):
    """快速筛选，例如 keyword=video 只返回视频算子"""
    res = [op for op in schema_data["operators"] if keyword in op["op_name"]]
    print(f"\n===== filter: contains [{keyword}], count={len(res)} =====")
    for item in res:
        print(f"OpName: {item['op_name']}, native params: {[p['name'] for p in item['native_params']]}")
    return res


if __name__ == "__main__":
    schema_data = dump_all_ops_schema()
    # 筛选所有video算子，打印预览
    video_ops = filter_op_by_keyword(schema_data, "video")