from bson import ObjectId
from tools.database.mongo import mainDB
from langchain.tools import tool

@tool
def QueryOriginalScript(session_id: str, only_exist: bool = False):
    """
    查询原始剧本内容或是否存在
    Args:
        session_id: 会话id
        only_exist: 是否只查询存在的剧本
        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            original_script (str): 原始剧本内容。仅当 only_exist 为 False 时返回该字段。
            exist (bool): 原始剧本内容是否存在。
    
    """
    # c = mainDB.agent_writer_session.count_documents({})
    # print(f"查询到的原始剧本session_id: {session_id}, only_exist: {only_exist} count:{c}")
    if only_exist:
        script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "original_script": {"$exists": True, "$ne": ""}})
        # print(f"exist: {script}")
        return {
            "original_script": "",
            "exist": script is not None,
        }
    else:
        script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id)}, {"original_script": 1})
        original_script = ""
        if script:
            original_script = script["original_script"] or ''
        print(f"查询到的原始剧本字符长度: {len(original_script)}")
        return {
                "original_script": original_script,
                "exist": original_script != '',
            }