import json
import logging
from mcp.types import Tool, TextContent
from .adbpg import DatabaseManager

logger = logging.getLogger(__name__)

def get_memory_tools() -> list[Tool]:
    return [
        Tool(
            name = "adbpg_llm_memory_add",
            description = "Execute llm_memory add operation",
            # 参数：messages json, user_id text, run_id text, agent_id text, metadata json
            # 增加新的记忆
            inputSchema={
                "type": "object",
                "properties": {
                    "messages": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "role": {"type": "string"},
                                "content": {"type": "string"}
                            },
                            "required": ["role", "content"]
                        },
                        "description": "List of messages objects (e.g., conversation history)"
                    },
                    "user_id": {
                        "type": "string",
                        "description": "the user_id"
                    },
                    "run_id": {
                        "type": "string",
                        "description": "the run_id"
                    },
                    "agent_id": {
                        "type": "string",
                        "description": "the agent_id"
                    },
                    "metadata": {
                        "type": "object",
                        "description": "the metatdata json"
                    },
                    "memory_type": {
                        "type": "string",
                        "description": "the memory_type text"
                    },
                    "prompt": {
                        "type": "string",
                        "description": "the prompt"
                    },
                    "infer": {
                        "type": "boolean",
                        "description": "use LLM to extract the memory"
                    }
                },
                "required": ["messages"]
            }
        ),
        Tool(
            name = "adbpg_llm_memory_get_all",
            description = "Execute llm_memory get_all operation",
            # 参数：user_id text, run_id text, agent_id text
            # 获取某个用户或者某个agent的所有记忆
            inputSchema={
                "type": "object",
                "properties": {
                    "user_id": {
                        "type": "string",
                        "description": "The user_id"
                    },
                    "run_id": {
                        "type": "string",
                        "description": "The run_id"
                    },
                    "agent_id": {
                        "type": "string",
                        "description": "The agent_id"
                    }
                },
                "required": []
            }
        ),
        Tool(
            name = "adbpg_llm_memory_search",
            description = "Execute llm_memory search operation",
            # 参数：query text, user_id text, run_id text, agent_id text, filter json, treshold number, rerank boolean
            # 获取与给定 query 相关的记忆
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "llm_memory relevant query"
                    },
                    "user_id": {
                        "type": "string",
                        "description": "The search of user_id"
                    },
                    "run_id": {
                        "type": "string",
                        "description": "The search of run_id"
                    },
                    "agent_id": {
                        "type": "string",
                        "description": "The search of agent_id"
                    },
                    "filter": {
                        "type": "object",
                        "description": "The search of filter"
                    },
                    "limits": {
                        "type": "integer",
                        "description": "The number of memory to return"
                    },
                    "threshold": {
                        "type": "number",
                        "minimum": 0.0,
                        "maximum": 1.0,
                        "description": "The threshold of similarity or memory results to return"
                    },
                     "rerank": {
                        "type": "boolean",
                        "description": "Whether to rerank the memory results"
                    },
                },
                "required": ["query"]
            }
        ),
        Tool(
            name = "adbpg_llm_memory_update",
            description = "Update a single memory item's content.",
            inputSchema={
                "type": "object",
                "properties": {
                    "memory_id": {
                        "type": "string",
                        "description": "The memory item id"
                    },
                    "updated_memory": {
                        "type": "string",
                        "description": "The updated memory item content"
                    }
                },
                "required": ["memory_id", "updated_memory"]
            }
        ),
        Tool(
            name = "adbpg_llm_memory_delete",
            description = "Delete a single memory item.",
            # 参数：memory_id text
            # Delete a specific memory item
            inputSchema={
                "type": "object",
                "properties": {
                    "memory_id": {
                        "type": "string",
                        "description": "The memory item id"
                    }
                },
                "required": ["memory_id"]
            }
        ),
        Tool(
            name = "adbpg_llm_memory_delete_all",
            description = "Execute llm_memory delete_all operation",
            # 参数：user_id text, run_id text, agent_id text
            # 删除某个用户或者agent的所有记忆
            inputSchema={
                "type": "object",
                "properties": {
                    "user_id": {
                        "type": "string",
                        "description": "The user_id"
                    },
                    "run_id": {
                        "type": "string",
                        "description": "The run_id"
                    },
                    "agent_id": {
                        "type": "string",
                        "description": "The agent_id"
                    }
                },
                "required": []
            }
        ),
        Tool(
            name = "adbpg_llm_memory_get_history",
            description = "Get the memory operation history of a memory item",
            inputSchema={
                "type": "object",
                "properties": {
                    "memory_id": {
                        "type": "string",
                        "description": "The memory id of a memory item"
                    }
                },
                "required": ["memory_id"]
            }
        ),
        Tool(
            name = "adbpg_llm_memory_delete_history",
            description = "Delete the memory operation history of a memory item",
            inputSchema={
                "type": "object",
                "properties": {
                    "memory_id": {
                        "type": "string",
                        "description": "The memory id of a memory item"
                    }
                },
                "required": ["memory_id"]
            }
        ),
        Tool(
            name="adbpg_llm_memory_delete_all_history",
            description="Delete all memory operation history",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        ),
        Tool(
            name="adbpg_llm_memory_set_custom_category",
            description="Set the custom category for the memory.",
            inputSchema={
                "type": "object",
                "properties": {
                    "custom_category": {
                        "type": "object",
                        "description": "The custom category for the memory."
                    }
                }
            }
        ),
        Tool(
            name="adbpg_llm_memory_get_custom_category",
            description="Get the custom category for the memory.",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        )
    ]

def _execute_memory_tool(sql: str, params: list, db: DatabaseManager) -> list[TextContent]:
    try:
        conn = db.get_llm_memory_connection()
        with conn.cursor() as cursor:
            cursor.execute(sql, params)
            if cursor.description:
                result = cursor.fetchone()[0]
                if isinstance(result, (dict, list)):
                    result = json.dumps(result, ensure_ascii=False)
                return [TextContent(type="text", text=result)]
            else:
                return [TextContent(type="text", text="LLM memory tool executed successfully.")]
    except Exception as e:
        return [TextContent(type="text", text=f"Error executing LLM memory tool: {str(e)}")]
    

async def call_memory_tool(name: str, arguments: dict, db: DatabaseManager) -> list[TextContent]:
    """调用 LLM Memory 工具"""
    user_id = arguments.get("user_id")
    run_id = arguments.get("run_id")
    agent_id = arguments.get("agent_id")

    if name == "adbpg_llm_memory_add":
        messages = arguments.get("messages")
        if not messages:
            raise ValueError("messages is required")
        if not any([user_id, run_id, agent_id]):
            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
        sql = "SELECT adbpg_llm_memory.add(%s::json, %s::text, %s::text, %s::text, %s::json, %s::text, %s::text, %s::boolean)"
        params = [
            json.dumps(messages, ensure_ascii=False),
            user_id, run_id, agent_id,
            json.dumps(arguments.get("metadata"),ensure_ascii=False) if arguments.get("metadata") else None,
            arguments.get("memory_type"),
            arguments.get("prompt"),
            arguments.get("infer")
        ]
    
    elif name == "adbpg_llm_memory_search":
        query = arguments.get("query")
        if not query:
            raise ValueError("query is required")
        if not any([user_id, run_id, agent_id]):
            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
       
        sql = "SELECT adbpg_llm_memory.search(%s::text, %s::text, %s::text, %s::text, %s::json, %s::int, %s::float, %s::boolean)"
        params = [
            query, user_id, run_id, agent_id,
            json.dumps(arguments.get("filter"), ensure_ascii=False) if arguments.get("filter") else None,
            arguments.get("limits"), arguments.get("threshold"), arguments.get("rerank")
        ]

    elif name == "adbpg_llm_memory_get_all":
        if not any([user_id, run_id, agent_id]):
            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
        sql = "SELECT adbpg_llm_memory.get_all(%s::text, %s::text, %s::text)"
        params = [user_id, run_id, agent_id]
    
    elif name == "adbpg_llm_memory_delete":
        memory_id = arguments.get("memory_id")
        sql = "SELECT adbpg_llm_memory.delete(%s::text)"
        params = [memory_id]
    
    elif name == "adbpg_llm_memory_delete_all":
        if not any([user_id, run_id, agent_id]):
            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
        sql = "SELECT adbpg_llm_memory.delete_all(%s::text, %s::text, %s::text)"
        params = [user_id, run_id, agent_id]
    
    elif name == "adbpg_llm_memory_get_history":
        sql = "SELECT adbpg_llm_memory.get_history(%s::text)"
        params = [arguments.get("memory_id")]
    
    elif name == "adbpg_llm_memory_delete_history":
        sql = "SELECT adbpg_llm_memory.delete_history(%s::text)"
        params = [arguments.get("memory_id")]

    elif name == "adbpg_llm_memory_delete_all_history":
        sql = "SELECT adbpg_llm_memory.delete_all_history()"
        params = []
    
    elif name == "adbpg_llm_memory_set_custom_category":
        sql = "SELECT adbpg_llm_memory.set_custom_category(%s::json)"
        params = [json.dumps(arguments.get("custom_category"),ensure_ascii=False) if arguments.get("custom_category") else None]

    elif name == "adbpg_llm_memory_get_custom_category":
        sql = "SELECT adbpg_llm_memory.get_custom_category()"
        params = []

    else:
        raise ValueError(f"Unknown memory tool: {name}")
    
    return _execute_memory_tool(sql, params, db)