from app.cache.memory.interest_memory import InterestMemoryCache
from app.core.memory.agent.utils.llm_tools import WriteState
from app.core.memory.agent.utils.write_tools import write
from app.core.logging_config import get_agent_logger

logger = get_agent_logger(__name__)


async def write_node(state: WriteState) -> WriteState:
    """
        Write data to the database/file system.

        Args:
            state: WriteState containing messages, end_user_id, memory_config, and language

        Returns:
            dict: Contains 'write_result' with status and data fields
        """
    messages = state.get('messages', [])
    end_user_id = state.get('end_user_id', '')
    memory_config = state.get('memory_config', '')
    language = state.get('language', 'zh')  # 默认中文

    # Convert LangChain messages to structured format expected by write()
    structured_messages = []
    for msg in messages:
        if hasattr(msg, 'type') and hasattr(msg, 'content'):
            # Map LangChain message types to role names
            role = 'user' if msg.type == 'human' else 'assistant' if msg.type == 'ai' else msg.type
            structured_messages.append({
                "role": role,
                "content": msg.content  # content is now guaranteed to be a string
            })

    try:
        result = await write(
            messages=structured_messages,
            end_user_id=end_user_id,
            memory_config=memory_config,
            language=language,
        )
        logger.info(f"Write completed successfully! Config: {memory_config.config_name}")

        # 写入 neo4j 成功后，删除该用户的兴趣分布缓存，确保下次请求重新生成
        for lang in ["zh", "en"]:
            deleted = await InterestMemoryCache.delete_interest_distribution(
                end_user_id=end_user_id,
                language=lang,
            )
            if deleted:
                logger.info(f"Invalidated interest distribution cache: end_user_id={end_user_id}, language={lang}")

        write_result = {
            "status": "success",
            "data": structured_messages,
            "config_id": memory_config.config_id,
            "config_name": memory_config.config_name,
        }
        return {"write_result": write_result}

    except Exception as e:
        logger.error(f"Data_write failed: {e}", exc_info=True)
        write_result = {
            "status": "error",
            "message": str(e),
        }
        return {"write_result": write_result}
