# -*- coding: utf-8 -*-
"""
gRPC 模型快速配置：mode 列表、标签文件映射、精度配置。
新增/修改 mode 时只需改本文件或通过 get_grpc_config(custom_*) 覆盖。
"""
from typing import Dict, List, Optional

# 模型标签文件映射（与 model_inference/labels 下文件名一致）
MODEL_LABEL_MAP = {
    "blood": "blood.txt",
    "chinamap": "chinamap.txt",
    "discomfort": "discomfort.txt",
    "discomfort_excrement": "discomfort_excrement.txt",
    "emblem": "emblem.txt",
    "fire": "fire.txt",
    "flag_classification": "flag_classification.txt",
    "flag_detect": "flag_detect.txt",
    "gamble": "gamble.txt",
    "media1": "media1.txt",
    "media2": "media2.txt",
    "military": "military.txt",
    "nsfw": "nsfw.txt",
    "politics": "politics.txt",
    "porn_classification": "porn_classification.txt",
    "prohibitlogo": "prohibitlogo.txt",
    "qrcode": "qrcode.txt",
    "religion_activity": "religion_activity.txt",
    "religion_clothes": "religion_clothes.txt",
    "religion_symbol": "religion_symbol.txt",
    "tattoo": "tattoo.txt",
    "tiananmen": "tiananmen.txt",
    "trypophobia": "trypophobia.txt",
}

# 模型精度：FP16 / FP32，未列出的默认 FP32
MODEL_DATA_TYPES = {
    "blood": "FP16",
    "chinamap": "FP16",
    "discomfort": "FP32",
    "discomfort_excrement": "FP32",
    "flag_classification": "FP16",
    "flag_detect": "FP16",
    "gamble": "FP32",
    "military": "FP32",
    "porn_classification": "FP16",
    "fire": "FP16",
    "prohibitlogo": "FP32",
    "qrcode": "FP16",
    "tattoo": "FP16",
    "tiananmen": "FP16",
    "trypophobia": "FP16",
    "politics": "FP32",
    "media1": "FP16",
    "media2": "FP16",
    "nsfw": "FP16",
    "emblem": "FP16",
    "religion_symbol": "FP16",
    "religion_activity": "FP16",
    "religion_clothes": "FP16",
}

# 所有支持的 mode（infer(mode, image) 的 mode 取值）
MODES_LIST = list(MODEL_LABEL_MAP.keys()) + ["porn", "flag"]

# mode=="all" 时实际执行的列表（与 client_test_grpc 一致）
MODES_ALL = [
    "fire", "qrcode", "prohibitlogo", "blood", "military",
    "chinamap", "discomfort", "discomfort_excrement", "gamble", "tattoo",
    "tiananmen", "trypophobia", "porn", "flag", "politics", "media1", "media2",
    "nsfw", "emblem", "religion_symbol", "religion_clothes",
]


def get_grpc_config(
    custom_label_map: Optional[Dict[str, str]] = None,
    custom_data_types: Optional[Dict[str, str]] = None,
    custom_modes_list: Optional[List[str]] = None,
    custom_modes_all: Optional[List[str]] = None,
) -> Dict:
    """
    返回合并后的 gRPC 配置，便于用户覆盖或扩展 mode。
    :return: dict 含 label_map, data_types, modes_list, modes_all
    """
    label_map = dict(MODEL_LABEL_MAP)
    if custom_label_map:
        label_map.update(custom_label_map)
    data_types = dict(MODEL_DATA_TYPES)
    if custom_data_types:
        data_types.update(custom_data_types)
    modes_list = custom_modes_list if custom_modes_list is not None else list(MODES_LIST)
    modes_all = custom_modes_all if custom_modes_all is not None else list(MODES_ALL)
    return {
        "label_map": label_map,
        "data_types": data_types,
        "modes_list": modes_list,
        "modes_all": modes_all,
    }
