from typing import Optional, Union
from chuangliangTool.db_base import Material1_db, Material2_db, Material3_db, Material4_db


def get_material(material_param: Optional[Union[int, str]] = None) -> Optional[str]:
    """
    获取素材
    参数:
        material_param: 可以是 material_id(int), material_name(str) 或 file_md5(str)
    返回:
        素材
    """
    if material_param is None:
        raise ValueError("必须提供 material_id, material_name 或 file_md5 参数")

    # 根据参数类型构建查询条件
    if isinstance(material_param, int):
        # 如果是整数，认为是material_id
        sql_template = f"SELECT * FROM chuangliang_ad_material.material WHERE material_id = '{material_param}'"
    else:
        # 如果是字符串，需要判断是material_name还是file_md5
        # 这里假设file_md5是32位的MD5值，material_name不是32位
        if len(material_param) == 32 and all(c in '0123456789abcdefABCDEF' for c in material_param):
            # 如果是32位且只包含十六进制字符，认为是file_md5
            sql_template = f"SELECT * FROM chuangliang_ad_material.material WHERE file_md5 = '{material_param}'"
        else:
            # 否则认为是material_name
            sql_template = f"SELECT * FROM chuangliang_ad_material.material WHERE material_name = '{material_param}'"

    material_dbs = [Material1_db, Material2_db, Material3_db, Material4_db]
    db_names = ["Material1_db", "Material2_db", "Material3_db", "Material4_db"]

    for db, db_name in zip(material_dbs, db_names):
        try:
            result = db.query_database_select_to_dataframe(sql_template)
            if result is not None and len(result) > 0:
                # 返回整个DataFrame，让调用者处理
                return result
        except Exception as e:
            print(f"查询{db_name}出错: {e}")

    return None


if __name__ == '__main__':
    # 测试不同的参数类型
    material_data = get_material("豪品素材-城市生活情感街拍短视频 (86)")
    print("查询结果:", material_data)