from pydantic import BaseModel
from typing import Dict
import requests, json, time, os
from datetime import datetime

class Property(BaseModel):
    type: str
    description: str

class InputSchema(BaseModel):
    type: str
    properties: Dict[str, Property]
    required: list

class AnthropicTool(BaseModel):
    name: str
    description: str
    input_schema: InputSchema

class OpenAITool(BaseModel):
    name: str
    description: str
    parameters: InputSchema

class AnthropicRetrieverTool:
    @staticmethod
    def create_tool() -> AnthropicTool:
        input_schema = InputSchema(
            type="object",
            properties={
                "query": Property(
                    type="string",
                    description="The user query to be used for retrieval. The query must be comprehensive and in Thai"
                )
            },
            required=["query"]
        )
        return AnthropicTool(
            name="retriever",
            description="Retrieve documents related to a given query from a vector database",
            input_schema=input_schema
        )

class OpenAIRetrieverTool:
    @staticmethod
    def create_tool() -> OpenAITool:
        parameters = InputSchema(
            type="object",
            properties={
                "query": Property(
                    type="string",
                    description="The user query to be used for retrieval. The query must be comprehensive and in Thai"
                )
            },
            required=["query"]
        )
        return OpenAITool(
            name="retriever",
            description="Retrieve documents related to a given query from a vector database",
            parameters=parameters
        )


class MCPTools:
    """MCP (Model Context Protocol) tool execution client.

    Connects to an MCP server to discover available tools and execute them.

    Args:
        mcp_url: Base URL of the MCP server
        user_chat_id: User/tenant chat ID for multi-tenant isolation
        user_context: User context dict passed to tool executions
        mcp_api_key: API key for MCP server. Defaults to MCP_API_KEY env var.
    """
    def __init__(self, mcp_url: str, user_chat_id: str, user_context: Dict = None, mcp_api_key: str = None, profiler=None):
        self.mcp_url = mcp_url
        self.user_chat_id = user_chat_id
        self.mcp_api_key = mcp_api_key or os.environ.get("MCP_API_KEY", "")
        self.profiler = profiler
        if self.profiler:
            self.profiler.mark("mcp_list_tools_start")
        self.tools = self._retrieve_tools_list()
        if self.profiler:
            self.profiler.mark("mcp_list_tools_end")
            self.profiler.record_duration("mcp_list_tools", self.profiler.between("mcp_list_tools_start", "mcp_list_tools_end"))
        self.user_context = user_context or {}

    def _retrieve_tools_list(self):
        res = requests.post(
            f"{self.mcp_url}/list_tools",
            headers={"x-api-key": self.mcp_api_key},
            json={"userChatID": self.user_chat_id}
        )
        if res.status_code != 200:
            raise ConnectionError(f"MCP list_tools failed (HTTP {res.status_code}): {res.text[:200]}")
        body = res.json()
        if "availableTools" not in body:
            raise KeyError(f"MCP response missing 'availableTools'. Got: {list(body.keys())}")
        tools = body["availableTools"]
        has_rag = any(t["name"] == "rag" for t in tools)
        if not has_rag:
            rag = {
                "type": "function",
                "name": "rag",
                "description": "RAG tool for semantic search",
                "endpoint": {"method": "POST", "path": "/rag"},
                "parameters": {
                    "type": "object",
                    "properties": {"query": {"type": "string"}},
                    "required": ["query"]
                }
            }
            tools.append(rag)
        return tools

    def execute_tool(self, tool_name: str, user_message=None, **data):
        if "userChatID" not in data:
            data["userChatID"] = self.user_chat_id
        if "dateTime" not in data:
            data["dateTime"] = datetime.now().strftime('%d/%m/%Y %H:%M:%S')
        if "userContext" not in data:
            data["userContext"] = self.user_context
        if tool_name == "send_email" and user_message:
            data["userMessage"] = user_message
        if tool_name == "rag":
            data["useBm25"] = True

        start = time.perf_counter()
        res = requests.post(
            f"{self.mcp_url}/{tool_name}",
            headers={"x-api-key": self.mcp_api_key},
            json=data
        )
        duration = time.perf_counter() - start

        if self.profiler:
            self.profiler.count("tool_calls")
            self.profiler.record_duration("mcp_tool_exec", duration)
            self.profiler.record_duration(f"mcp_{tool_name}", duration)

        return res.json()

    @staticmethod
    def parse_response(tool_name: str, response: dict, tool_input):
        metadata = None
        scores = None
        if tool_name == "rag":
            if response["status"] == "failed":
                return response["result"], None

            output, metadata = response["output"], response["metadata"]
            if isinstance(metadata, str):
                metadata = json.loads(metadata)
            scores = [item["score"] for item in metadata["documents"]]
        else:
            output = response["result"]
            metadata = None

        print(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": f"Tool {tool_name} results. input: {str(tool_input)} result: <result>{output[:500]}</result> scores: {str(scores)}",
            "exception": None
        }, ensure_ascii=False))

        return output, metadata


if __name__ == "__main__":
    print(AnthropicRetrieverTool.create_tool().model_dump())