import { AiToolExecutionErrorCode, } from "@superblocksteam/library-shared/types";
import { OperationQueue } from "../../../util/operation-queue.js";
import { ActionBasedPermissionKeyStrategy, } from "./entity-permissions.js";
import { PermissionLevel, ToolExecutionError, } from "./types.js";
/**
 * Middleware that captures compaction summaries for tool outputs.
 *
 * If a tool has a `getSummary` function, this middleware will call it
 * after execution and store the summary via the provided callback.
 *
 * Unlike the previous implementation, this does NOT mutate the tool output,
 * which fixes compatibility issues with type-safe providers like Bedrock.
 *
 * @param storeSummary - Callback to store summaries (typically context.storeSummary)
 * @returns ToolMiddleware that captures summaries
 */
export function summaryInjector(storeSummary) {
    return (tool) => {
        // Only wrap if the tool has a getSummary function
        if (!tool.getSummary) {
            return tool;
        }
        const getSummary = tool.getSummary;
        return {
            ...tool,
            execute: async (input, options) => {
                const output = await tool.execute(input, options);
                const summary = getSummary(input, output);
                storeSummary(options.toolCallId, summary);
                return output;
            },
        };
    };
}
/**
 * Central registry for tools with access control.
 *
 * Provides:
 * - Tool registration and instantiation
 * - Access control middleware
 * - User prompt handling
 * - Tool discovery and metadata
 */
export class ToolRegistry {
    factories = new Map();
    instances = new Map();
    accessRules = [];
    middleware = [];
    promptHandler;
    promptQueue = new OperationQueue();
    entityPermissionStore;
    scopeStrategy = new ActionBasedPermissionKeyStrategy();
    /**
     * Register a tool factory and optionally instantiate it immediately.
     *
     * @param factory - Tool factory to register
     * @param deps - Optional dependencies to instantiate the tool immediately
     * @throws Error if tool with same name already registered
     */
    register(factory, deps) {
        if (this.factories.has(factory.name) || this.instances.has(factory.name)) {
            throw new Error(`Tool "${factory.name}" already registered`);
        }
        this.factories.set(factory.name, factory);
        // If dependencies are provided, instantiate immediately
        if (deps !== undefined) {
            const instance = factory.create(deps);
            if (instance.name !== factory.name) {
                throw new Error(`Tool name mismatch: factory.name="${factory.name}" but tool.name="${instance.name}"`);
            }
            this.instances.set(factory.name, instance);
        }
    }
    /**
     * Register an already-instantiated tool (for tools not using factory pattern).
     *
     * @param tool - Tool instance to register (must have name and execute)
     * @throws Error if tool with same name already registered
     */
    registerTool(tool) {
        if (!tool.name) {
            throw new Error("Tool must have a name property");
        }
        if (this.instances.has(tool.name)) {
            throw new Error(`Tool "${tool.name}" already registered`);
        }
        // Coerce to Tool type - AI SDK tools may have different structure
        this.instances.set(tool.name, tool);
    }
    /**
     * Get list of all registered tool names.
     *
     * @returns Array of tool names
     */
    listTools() {
        return Array.from(this.instances.keys());
    }
    /**
     * Register multiple tool factories at once.
     *
     * @param factories - Array of tool factories to register
     */
    registerAll(factories) {
        for (const factory of factories) {
            this.register(factory);
        }
    }
    /**
     * Get a tool instance by name.
     *
     * Returns a wrapped version where execute() enforces access control.
     *
     * @param name - Tool name
     * @param context - Access control context (optional)
     * @returns Wrapped tool instance or undefined if not found
     */
    getTool(name, context) {
        const tool = this.getToolInternal(name);
        return tool ? this.wrapTool(tool, context ?? {}) : undefined;
    }
    /**
     * Get all tool instances.
     *
     * Returns wrapped versions where execute() enforces access control.
     *
     * @param context - Access control context (optional)
     * @returns Array of all wrapped tool instances
     */
    getAllTools(context) {
        return Array.from(this.instances.keys()).flatMap((name) => {
            const tool = this.getTool(name, context);
            return tool ? [tool] : [];
        });
    }
    /**
     * Get metadata for all tools.
     *
     * @returns Array of metadata for all instantiated tools
     */
    getToolMetadata(context) {
        return this.getAllTools(context).map((tool) => ({
            name: tool.name,
            description: tool.description,
            category: tool.category,
            defaultPermission: tool.defaultPermission,
        }));
    }
    /**
     * Get tools as a record for AI SDK.
     *
     * Returns wrapped versions where execute() enforces access control.
     * Filters out tools that are BLOCKED by access control rules.
     *
     * @param context - Access control context (optional, e.g., { mode: AiMode.PLAN })
     * @returns Record mapping tool names to wrapped tool instances
     */
    async getToolsRecord(context) {
        const ctx = context ?? {};
        // Filter out blocked tools
        const allowedTools = [];
        for (const tool of this.getAllTools(context)) {
            const permission = await this.checkPermission(tool.name, ctx);
            if (permission === PermissionLevel.BLOCK) {
                continue;
            }
            allowedTools.push([tool.name, tool]);
        }
        return Object.fromEntries(allowedTools);
    }
    /**
     * Add one or more access control rules.
     *
     * Rules are evaluated in order when checking permissions.
     * Most restrictive permission level wins.
     *
     * @param rules - Access control rule function(s)
     */
    addAccessRule(...rules) {
        this.accessRules.push(...rules);
    }
    /**
     * Add middleware to transform tools.
     *
     * Middleware is applied in the order it's added - first-added middleware
     * executes first (outermost wrapper).
     *
     * @param middleware - Middleware function(s)
     *
     * @example
     * registry.addMiddleware(loggingMiddleware);    // executes first
     * registry.addMiddleware(outputLimiter);         // executes second
     * registry.addMiddleware((tool) => profileTool(tool)); // executes last
     */
    addMiddleware(...middleware) {
        this.middleware.push(...middleware);
    }
    /**
     * Set handler for prompting user for confirmation.
     *
     * @param handler - Prompt handler function
     */
    setPromptHandler(handler) {
        this.promptHandler = (tool, input, _p, context, options) => this.promptQueue.enqueue(async () => {
            const shouldPromptForEntities = await this.shouldPromptForEntities(tool, input, context);
            if (!shouldPromptForEntities) {
                return;
            }
            const permissions = await this.getPendingPermissions(tool, input);
            return handler(tool, input, permissions, context, options);
        });
    }
    /**
     * Set entity permission store for tracking per-entity approvals.
     *
     * When set, tools with `getRequiredPermissions` defined will automatically
     * check if required permissions have been approved before execution.
     * If any permission is not approved, the user will be prompted.
     *
     * @param store - Entity permission store instance
     */
    setEntityPermissionStore(store) {
        this.entityPermissionStore = store;
    }
    /**
     * Get permission scope for a tool.
     *
     * Computes the scope string using the configured strategy.
     * This scope is used to isolate permissions between different tools/categories.
     *
     * @param tool - Tool to get scope for
     * @returns Permission scope string
     */
    getPermissionScope(permission, input) {
        const key = this.scopeStrategy.getPermissionKey(permission, input);
        return this.scopeStrategy.formatScope(key);
    }
    async getPendingPermissions(tool, input) {
        return (await Promise.all(((await tool.getRequiredPermissions?.(input)) ?? []).map(async (permission) => ({
            permission,
            isApproved: this.entityPermissionStore?.isApproved(permission.entityType, permission.entityId, this.getPermissionScope(permission, input)),
        }))))
            .filter(({ isApproved }) => !isApproved)
            .map(({ permission }) => permission);
    }
    /**
     * Check permission level for tool execution.
     *
     * Evaluates all access control rules and returns most restrictive permission.
     *
     * @param toolName - Name of tool to check
     * @param context - Access control context
     * @param input - Tool input
     * @returns Permission level for execution
     */
    async checkPermission(toolName, context, input) {
        // Start with tool's default permission
        const tool = this.getToolInternal(toolName);
        if (!tool) {
            return PermissionLevel.BLOCK;
        }
        let permission = tool?.defaultPermission ?? PermissionLevel.ALLOW;
        // Check all rules, most restrictive wins
        for (const rule of this.accessRules) {
            const result = await rule(tool, context, input);
            // BLOCK is most restrictive
            if (result === PermissionLevel.BLOCK) {
                return PermissionLevel.BLOCK;
            }
            // PROMPT is more restrictive than ALLOW
            if (result === PermissionLevel.PROMPT) {
                if (this.entityPermissionStore?.isToolCallApproved(toolName, input)) {
                    this.entityPermissionStore?.revokeToolCall(toolName, input);
                    return PermissionLevel.ALLOW;
                }
                permission = PermissionLevel.PROMPT;
            }
        }
        return permission;
    }
    /**
     * Execute a tool with access control.
     *
     * Checks permissions, prompts user if needed, validates input,
     * and executes the tool.
     *
     * @param toolName - Name of tool to execute
     * @param input - Input data for tool
     * @param context - Access control context
     * @returns Execution result with success status and data/error
     */
    async execute(toolName, input, context = {}, options) {
        const trimmedContextForErrorsReport = {
            ...context,
            ...(context.planContext &&
                context.planContext.attachments &&
                context.planContext.attachments.length > 0 && {
                planContext: {
                    ...context.planContext,
                    attachments: `Redacted attachments: ${context.planContext.attachments.length}`,
                },
            }),
        };
        // Get raw tool instance (not wrapped to avoid infinite recursion)
        const tool = this.getToolInternal(toolName);
        if (!tool) {
            return {
                type: "error",
                error: new ToolExecutionError(`Tool "${toolName}" not found`, {
                    toolName,
                    input,
                    context: trimmedContextForErrorsReport,
                    code: AiToolExecutionErrorCode.NOT_FOUND,
                }),
            };
        }
        // Validate input
        const validated = tool.inputSchema.safeParse(input);
        if (!validated.success) {
            return {
                type: "error",
                error: new ToolExecutionError(`Invalid input for tool "${toolName}": ${validated.error.message}, you should call it again with the correct input!`, 
                // context intentionally empty since the llm doesnt need this to re-call the tool on input schema failure
                {
                    toolName,
                    input,
                    context: {},
                    code: AiToolExecutionErrorCode.INVALID_INPUT,
                }),
            };
        }
        // Invoke onCall hook (fires before permission checks, can abort by throwing)
        if (tool.onCall) {
            try {
                await tool.onCall(input, options);
            }
            catch (error) {
                return {
                    type: "error",
                    error: new ToolExecutionError({
                        toolName,
                        input,
                        context: trimmedContextForErrorsReport,
                        cause: error,
                        code: AiToolExecutionErrorCode.INTERNAL_ERROR,
                    }),
                };
            }
        }
        // Check permission
        const permission = await this.checkPermission(toolName, context, input);
        // Handle BLOCK
        if (permission === PermissionLevel.BLOCK) {
            return {
                type: "error",
                error: new ToolExecutionError(`Tool "${toolName}" execution blocked by access control`, {
                    toolName,
                    input,
                    context: trimmedContextForErrorsReport,
                    code: AiToolExecutionErrorCode.NOT_ALLOWED,
                }),
            };
        }
        // Handle PROMPT
        if (permission === PermissionLevel.PROMPT) {
            // Check if we need to prompt (or if entities are already approved)
            const needsPrompt = await this.shouldPromptForEntities(tool, input, context);
            if (needsPrompt) {
                // Ensure prompt handler exists
                if (!this.promptHandler) {
                    return {
                        type: "error",
                        error: new ToolExecutionError(`Tool "${toolName}" requires user confirmation but no prompt handler registered`, {
                            toolName,
                            input,
                            context: trimmedContextForErrorsReport,
                            code: AiToolExecutionErrorCode.INTERNAL_ERROR,
                        }),
                    };
                }
                // Prompt user for approval
                await this.promptHandler(tool, input, [], context, options);
                return {
                    type: "error",
                    error: new ToolExecutionError(`Tool "${toolName}" execution blocked by access control`, {
                        toolName,
                        input,
                        context: trimmedContextForErrorsReport,
                        code: AiToolExecutionErrorCode.REQUIRES_USER_CONFIRMATION,
                    }),
                };
            }
            // If !needsPrompt, all entities already approved - proceed to execution
        }
        // Execute tool
        try {
            const output = await tool.execute(validated.data, options ?? {
                toolCallId: "",
                messages: [],
            });
            return {
                type: "success",
                data: output,
            };
        }
        catch (error) {
            // Preserve severity from cause if it's a ToolExecutionError
            const severity = error instanceof ToolExecutionError ? error.severity : undefined;
            return {
                type: "error",
                error: new ToolExecutionError({
                    toolName,
                    input,
                    context: trimmedContextForErrorsReport,
                    cause: error,
                    code: AiToolExecutionErrorCode.INTERNAL_ERROR,
                    severity,
                }),
            };
        }
    }
    /**
     * Determines whether a user prompt is required for executing a tool with PROMPT permission.
     *
     * If the tool supports entity extraction and all required entities have already been approved
     * for the relevant permission scope, no prompt is needed and this returns false.
     * If entity extraction is not supported or the entity permission store isn't configured, returns true.
     * If any required entity is not yet approved, returns true (prompt needed).
     *
     * @param tool - The tool to check.
     * @param input - Input from which to extract required entities.
     * @returns True if prompting the user is required; false if all approvals are in place.
     */
    async shouldPromptForEntities(tool, input, _context) {
        // No entity system configured, use standard prompting
        const entityPermissionStore = this.entityPermissionStore;
        if (!tool.getRequiredPermissions || !entityPermissionStore) {
            return true;
        }
        const permissions = await tool.getRequiredPermissions(input);
        if (permissions.length === 0) {
            return true;
        }
        // Check if any entities are not approved for this scope
        // Map each permission to a promise that resolves to true if not approved
        const approvalChecks = permissions.map((permission) => {
            const scope = this.getPermissionScope(permission, input);
            return !entityPermissionStore.isApproved(permission.entityType, permission.entityId, scope);
        });
        // Return true if any permission is not approved (needs prompt)
        return approvalChecks.some((needsPrompt) => needsPrompt);
    }
    /**
     * Store entity approvals after successful user prompt.
     *
     * Extracts entities from tool input and marks them as approved for this
     * tool's scope in the entity permission store. No-op if entity system
     * is not configured.
     *
     * @param tool - Tool that was approved
     * @param input - Tool input to extract entities from
     */
    async approveEntitiesAfterPrompt(tool, input, alwaysAllowPermissions) {
        // No entity system configured
        if (!tool.getRequiredPermissions || !this.entityPermissionStore) {
            return;
        }
        const permissions = await tool.getRequiredPermissions(input);
        const permissionsToApprove = alwaysAllowPermissions
            ? permissions.filter((permission) => alwaysAllowPermissions.some((alwaysAllow) => alwaysAllow.entityType === permission.entityType &&
                alwaysAllow.entityId === permission.entityId &&
                alwaysAllow.permissionType === permission.permissionType))
            : [];
        for (const permission of permissionsToApprove) {
            // Compute scope for each permission individually
            const scope = this.getPermissionScope(permission, input);
            this.entityPermissionStore.approve(permission.entityType, permission.entityId, scope);
        }
    }
    /**
     * Wrap a tool to enforce access control on execute().
     *
     * Returns a proxied version where execute() goes through the registry's
     * permission checks, ensuring access control is enforced regardless of
     * how the tool is accessed.
     *
     * @param tool - Raw tool instance
     * @param context - Access control context to use for permission checks
     * @returns Wrapped tool with access-controlled execute()
     */
    wrapTool(tool, context) {
        return {
            ...tool,
            execute: async (input, options) => {
                // Execute through registry to apply access control with captured context
                const result = await this.execute(tool.name, input, context, options);
                if (result.type === "error") {
                    throw result.error;
                }
                return result.data;
            },
        };
    }
    /**
     * Internal method to get raw tool instance without wrapping
     * or access control. Used internally to avoid infinite recursion.
     *
     * The returned tool will still have middleware applied.
     */
    getToolInternal(name) {
        const tool = this.instances.get(name);
        if (!tool) {
            return tool;
        }
        return this.applyMiddleware(tool);
    }
    applyMiddleware(tool) {
        let processedTool = tool;
        // Apply middleware in reverse order so first-added becomes outermost
        for (let i = this.middleware.length - 1; i >= 0; i--) {
            processedTool = this.middleware[i](processedTool);
        }
        return processedTool;
    }
    /**
     * Clear all registered tools and rules.
     *
     * Useful for testing or resetting state.
     */
    clear() {
        this.factories.clear();
        this.instances.clear();
        this.accessRules = [];
        this.middleware = [];
        this.promptHandler = undefined;
        this.entityPermissionStore = undefined;
    }
}
//# sourceMappingURL=registry.js.map