# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Toolbox Registry for the Agent service.
Manages thematic tool groupings (toolboxes) and the `requestToolbox` meta-tool.
"""

import logging
from typing import Dict, List, Any, Optional, Set
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


class ToolboxDefinition(BaseModel):
    """Definition of a thematic toolbox."""
    id: str = Field(description="Unique toolbox identifier (e.g. 'core', 'email', 'workflow')")
    label: str = Field(description="Human-readable label")
    description: str = Field(default="", description="What this toolbox provides")
    featureCode: Optional[str] = Field(default=None, description="Feature code if toolbox is feature-specific")
    tools: List[str] = Field(default_factory=list, description="Tool names belonging to this toolbox")
    isDefault: bool = Field(default=False, description="If true, toolbox is active by default")
    requiresConnection: Optional[str] = Field(
        default=None,
        description="Connection authority required (e.g. 'microsoft', 'google'). None = always available."
    )


class ToolboxRegistry:
    """Registry for toolbox definitions. Manages activation and tool lookup."""

    def __init__(self):
        self._toolboxes: Dict[str, ToolboxDefinition] = {}

    def registerToolbox(self, toolbox: ToolboxDefinition) -> None:
        """Register a toolbox definition."""
        if toolbox.id in self._toolboxes:
            logger.debug("Toolbox '%s' already registered, updating", toolbox.id)
        self._toolboxes[toolbox.id] = toolbox
        logger.debug("Registered toolbox: %s (%d tools, default=%s)", toolbox.id, len(toolbox.tools), toolbox.isDefault)

    def getToolbox(self, toolboxId: str) -> Optional[ToolboxDefinition]:
        """Get a toolbox by ID."""
        return self._toolboxes.get(toolboxId)

    def getAllToolboxes(self) -> List[ToolboxDefinition]:
        """Get all registered toolboxes."""
        return list(self._toolboxes.values())

    def getDefaultToolboxes(self) -> List[ToolboxDefinition]:
        """Get all default toolboxes (active at agent start)."""
        return [tb for tb in self._toolboxes.values() if tb.isDefault]

    def getActiveToolboxes(self, userConnections: List[str] = None) -> List[ToolboxDefinition]:
        """
        Get toolboxes available to the user based on their connections.
        Toolboxes without requiresConnection are always available.
        Toolboxes with requiresConnection are available only if the user has that connection.
        """
        available = []
        connectionAuthorities: Set[str] = set(userConnections or [])
        for tb in self._toolboxes.values():
            if tb.requiresConnection is None:
                available.append(tb)
            elif tb.requiresConnection in connectionAuthorities:
                available.append(tb)
        return available

    def getToolsForToolboxes(self, toolboxIds: List[str]) -> List[str]:
        """Get the union of all tool names for the given toolbox IDs."""
        tools: Set[str] = set()
        for tbId in toolboxIds:
            tb = self._toolboxes.get(tbId)
            if tb:
                tools.update(tb.tools)
        return sorted(tools)

    def getToolboxForTool(self, toolName: str) -> Optional[str]:
        """Find which toolbox a tool belongs to."""
        for tb in self._toolboxes.values():
            if toolName in tb.tools:
                return tb.id
        return None

    def toApiResponse(self, userConnections: List[str] = None) -> List[Dict[str, Any]]:
        """Serialize available toolboxes for API response."""
        available = self.getActiveToolboxes(userConnections)
        return [
            {
                "id": tb.id,
                "label": tb.label,
                "description": tb.description,
                "toolCount": len(tb.tools),
                "isDefault": tb.isDefault,
                "requiresConnection": tb.requiresConnection,
            }
            for tb in available
        ]


# Module-level singleton
_toolboxRegistry = ToolboxRegistry()


def getToolboxRegistry() -> ToolboxRegistry:
    """Get the global toolbox registry singleton."""
    return _toolboxRegistry


def _registerDefaultToolboxes() -> None:
    """Register the default set of toolboxes."""
    defaults = [
        ToolboxDefinition(
            id="core",
            label="Core Tools",
            description="Basic agent tools: search, read, write, web",
            isDefault=True,
            tools=[
                "readFile", "listFiles", "searchInFileContent", "listFolders",
                "webSearch", "readUrl", "writeFile", "deleteFile", "renameFile",
                "copyFile", "createFolder", "deleteFolder", "moveFile", "moveFolder",
                "renameFolder", "tagFile", "replaceInFile", "translateText",
                "detectLanguage", "queryFeatureInstance",
            ],
        ),
        ToolboxDefinition(
            id="ai",
            label="AI Tools",
            description="AI-powered analysis and generation",
            isDefault=True,
            tools=[
                "summarizeContent", "describeImage", "generateImage",
                "textToSpeech", "speechToText", "renderDocument",
                "createChart", "executeCode", "neutralizeData",
            ],
        ),
        ToolboxDefinition(
            id="datasources",
            label="Data Sources",
            description="Access external data sources and databases",
            isDefault=True,
            tools=[
                "listConnections", "browseDataSource", "searchDataSource",
                "downloadFromDataSource", "uploadToExternal",
                "browseContainer", "readContentObjects", "extractContainerItem",
            ],
        ),
        ToolboxDefinition(
            id="email",
            label="Email",
            description=(
                "Outlook mail management: send/draft new mails, reply or forward existing "
                "messages (preserves the conversation thread), move/copy/delete/archive "
                "mails, mark as read/unread, set follow-up flags, and list mail folders. "
                "Use replyToMail (NOT sendMail) when answering an existing message so the "
                "Outlook thread stays intact."
            ),
            requiresConnection="msft",
            isDefault=False,
            tools=[
                "sendMail",
                "replyToMail",
                "forwardMail",
                "moveMail",
                "deleteMail",
                "archiveMail",
                "setMailReadState",
                "flagMail",
                "listMailFolders",
            ],
        ),
        ToolboxDefinition(
            id="sharepoint",
            label="SharePoint",
            description="Access SharePoint sites, lists, and files",
            requiresConnection="msft",
            isDefault=False,
            tools=[
                "sharepoint_findDocuments", "sharepoint_readDocuments",
                "sharepoint_upload",
            ],
        ),
        ToolboxDefinition(
            id="clickup",
            label="ClickUp",
            description="Manage ClickUp tasks and projects",
            requiresConnection="clickup",
            isDefault=False,
            tools=[
                "clickup_listTasks",
                "clickup_listFields",
                "clickup_searchTasks",
                "clickup_getTask",
                "clickup_createTask",
                "clickup_updateTask",
                "clickup_uploadAttachment",
            ],
        ),
        ToolboxDefinition(
            id="jira",
            label="Jira",
            description="Manage Jira issues and projects",
            requiresConnection="jira",
            isDefault=False,
            tools=[
                "jira_connect", "jira_exportTickets", "jira_importTickets",
            ],
        ),
        ToolboxDefinition(
            id="workflow",
            label="Workflow",
            description="Graph manipulation tools for the visual editor",
            featureCode="graphicalEditor",
            isDefault=False,
            tools=[
                "readWorkflowGraph", "addNode", "removeNode", "connectNodes",
                "setNodeParameter", "listAvailableNodeTypes", "describeNodeType",
                "autoLayoutWorkflow", "validateGraph",
                "listWorkflowHistory", "readWorkflowMessages",
                "createWorkflow", "updateWorkflowMetadata", "createWorkflowFromFile",
                "exportWorkflowToFile", "deleteWorkflow",
            ],
        ),
        ToolboxDefinition(
            id="trustee",
            label="Trustee / Accounting",
            description="Trustee accounting tools: refresh data from external system (e.g. Abacus), query positions and journal entries",
            featureCode="trustee",
            isDefault=False,
            tools=[
                "trustee_refreshAccountingData",
            ],
        ),
    ]
    for tb in defaults:
        _toolboxRegistry.registerToolbox(tb)


_registerDefaultToolboxes()


REQUEST_TOOLBOX_TOOL_NAME = "requestToolbox"


def buildRequestToolboxDefinition(availableToolboxIds: List[str]) -> dict:
    """Build the tool definition dict for the requestToolbox meta-tool."""
    return {
        "name": REQUEST_TOOLBOX_TOOL_NAME,
        "description": (
            "Request additional specialized tools for the current task. "
            "Call this when you need tools from a specific toolbox that is not yet active. "
            "After calling, the requested tools will be available in the next round."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "toolboxId": {
                    "type": "string",
                    "enum": availableToolboxIds,
                    "description": "ID of the toolbox to activate",
                },
                "reason": {
                    "type": "string",
                    "description": "Brief reason why this toolbox is needed",
                },
            },
            "required": ["toolboxId"],
        },
    }