import pymongo

def get_mongo_col(col="config_beta"):
    """
    获取操作mongo集合
    """
    mongo_3012 = "mongodb://dashboard:9qRV71M%25Trx1CI@m3012.hw.mongos.m.com:3012/dashboard?authSource=dashboard"
    client = pymongo.MongoClient(mongo_3012, maxPoolSize=500)
    db = client["dashboard"]
    col = db[col]  # data_package(7234543961263854561)  goods_feature("7234800834201014260")
    return col

def set_txt2txt_price_config():
    # 拿到 collection 对象
    collection = get_mongo_col()  # 默认 col="config_beta"

    doc = collection.find_one(
        {"name": "txt2txt_price_config"},
        {"_id": 0, "data": 1}
    )
    if not doc:
        raise RuntimeError("找不到配置 document")

    data = doc["data"]

    data["doubao-seed-1-6-251015"] = {
        "completion_price": 0.0168,
        "prompt_price": 0.00168,
        "currency": "￥"
    }

    data["doubao-seed-1-8-251228"] = {
        "completion_price": 0.0216,
        "prompt_price": 0.00216,
        "currency": "￥"
    }

    # 3. 一次性把整个 data 写回去
    result = collection.update_one(
        {"name": "txt2txt_price_config"},
        {"$set": {"data": data}}
    )
    print(f"matched={result.matched_count}, modified={result.modified_count}")


    query = {"name": "txt2txt_price_config"}
    projection = {"_id": 0, "data": 1}

    doc = collection.find_one(query, projection)


    if doc:
        data_value = doc.get("data")
        print(data_value)
    else:
        print("未找到 name 等于 txt2txt_price_config 的记录")               

def set_txt2txt_config():
    # 拿到 collection 对象
    collection = get_mongo_col()  # 默认 col="config_beta"

    doc = collection.find_one(
        {"name": "txt2txt_config"},
        {"_id": 0, "data": 1}
    )
    if not doc:
        raise RuntimeError("找不到配置 document")

    data = doc["data"]

    data.append({
        "stream_selected": 0,
        "label": "doubao-seed-1-6-251015",
        "value": "doubao-seed-1-6-251015",
        "wechat_selected": 0,
        "wechat_label": 0,
        "selected": 0,
        "max_tokens": 4000,
        "is_think": 0,
        "web_label": 0,
        "stream_label": 0,
        "manufacturer": "字节"
    })

    data.append({
        "stream_selected": 0,
        "label": "doubao-seed-1-8-251228",
        "value": "doubao-seed-1-8-251228",
        "wechat_selected": 0,
        "wechat_label": 0,
        "selected": 0,
        "max_tokens": 4000,
        "is_think": 0,
        "web_label": 0,
        "stream_label": 0,
        "manufacturer": "字节"
    })

    # 3. 一次性把整个 data 写回去
    result = collection.update_one(
        {"name": "txt2txt_config"},
        {"$set": {"data": data}}
    )
    print(f"matched={result.matched_count}, modified={result.modified_count}")

    query = {"name": "txt2txt_config"}
    projection = {"_id": 0, "data": 1}

    doc = collection.find_one(query, projection)

    if doc:
        data_value = doc.get("data")
        print(data_value)
    else:
        print("未找到 name 等于 txt2txt_config 的记录")       

set_txt2txt_price_config()
set_txt2txt_config()