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

@tool
def QueryOriginalScript(session_id: str):
    """
    查询原始剧本内容是否存在
    Args:
        session_id: 会话id
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 原始剧本内容是否存在。
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "original_script": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }


def QueryDiagnosisAndAssessment(session_id: str):
    """
    查询诊断与资产评估报告是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 诊断与资产评估报告是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "diagnosis_and_assessment": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryAdaptationIdeas(session_id: str):
    """
    查询改编思路是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 改编思路是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "adaptation_ideas": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryScriptBible(session_id: str):
    """
    查询剧本圣经是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible": {"$exists": True}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryCoreOutline(session_id: str):
    """
    查询剧本圣经中的核心大纲是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的核心大纲是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.core_outline": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryCharacterProfile(session_id: str):
    """
    查询剧本圣经中的核心人物小传是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的核心人物小传是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.character_profile": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryCoreEventTimeline(session_id: str):
    """
    查询剧本圣经中的重大事件时间线是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的重大事件时间线是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.core_event_timeline": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryCharacterList(session_id: str):
    """
    查询剧本圣经中的总人物表是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的总人物表是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.character_list": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryEpisodeCount(session_id: str):
    """
    查询剧集创作情况
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            completed (int): 已完成的集数
            total (int): 总集数
    """
    total = mainDB.agent_writer_episodes.count_documents({"session_id": session_id})
    count = mainDB.agent_writer_episodes.count_documents({"session_id": session_id, "content": {"$exists": True, "$ne": ""}})
    return {
        "completed": count,
        "total": total,
    }

# def QuerySingleEpisodeContent(session_id: str):
#     """
#     查询单集完整内容
#     Args:
#         session_id: 会话id        
#     Returns:
#         Dict: 返回一个包含以下字段的字典：
#             exist (bool): 剧本圣经中的总人物表是否存在。    
#     """
#     pass