/*
 * Copyright 2008-2026 Async-IO.org
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package org.atmosphere.mcp.registry;

import org.atmosphere.mcp.annotation.McpParam;
import org.atmosphere.mcp.annotation.McpPrompt;
import org.atmosphere.mcp.annotation.McpResource;
import org.atmosphere.mcp.annotation.McpTool;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Registry for MCP tools, resources, and prompts. Supports both annotation-based
 * discovery (via {@link #scan(Object)}) and programmatic registration
 * (via {@link #registerTool(String, String, List, ToolHandler)}).
 */
public final class McpRegistry {

    private static final Logger logger = LoggerFactory.getLogger(McpRegistry.class);

    /**
     * Functional interface for dynamically registered tools.
     * Receives a map of argument name → value, returns the tool result.
     *
     * <p>Handlers that need the authenticated caller's identity — to gate
     * writes, scope tenant lookups, etc. — implement the sibling
     * {@link IdentityAwareToolHandler}. The registry invokes whichever
     * SAM the registered instance conforms to; existing single-arg
     * lambdas keep compiling unchanged.</p>
     */
    @FunctionalInterface
    public interface ToolHandler {
        Object execute(Map<String, Object> arguments) throws Exception;
    }

    /**
     * Identity-aware tool handler. {@code McpProtocolHandler} invokes this
     * variant when the registered handler implements it, passing the
     * servlet-resolved principal name ({@code null} for anonymous
     * callers). Used by {@code AdminMcpBridge} so write tools can forward
     * the caller identity to {@code ControlAuthorizer.authorize}.
     */
    @FunctionalInterface
    public interface IdentityAwareToolHandler extends ToolHandler {
        /** Execute with the caller's principal name ({@code null} when anonymous). */
        Object execute(Map<String, Object> arguments, String principal) throws Exception;

        /** Default delegates to the identity-aware form with principal=null. */
        @Override
        default Object execute(Map<String, Object> arguments) throws Exception {
            return execute(arguments, null);
        }
    }

    /**
     * Functional interface for dynamically registered resources.
     * Receives a map of argument name → value, returns the resource content.
     */
    @FunctionalInterface
    public interface ResourceHandler {
        Object read(Map<String, Object> arguments) throws Exception;
    }

    /**
     * Functional interface for dynamically registered prompts.
     * Receives a map of argument name → value, returns prompt messages.
     */
    @FunctionalInterface
    public interface PromptHandler {
        Object get(Map<String, Object> arguments) throws Exception;
    }

    /**
     * Metadata for a registered MCP tool.
     */
    public record ToolEntry(String name, String description, Method method, Object instance,
                            List<ParamEntry> params, ToolHandler handler) {
        /** Annotation-based constructor (method + instance). */
        public ToolEntry(String name, String description, Method method, Object instance,
                         List<ParamEntry> params) {
            this(name, description, method, instance, params, null);
        }

        /** Programmatic constructor (handler lambda). */
        public ToolEntry(String name, String description, List<ParamEntry> params,
                         ToolHandler handler) {
            this(name, description, null, null, params, handler);
        }

        /** Returns true if this tool uses a programmatic handler. */
        public boolean isDynamic() {
            return handler != null;
        }
    }

    /**
     * Metadata for a registered MCP resource.
     */
    public record ResourceEntry(String uri, String name, String description, String mimeType,
                                Method method, Object instance, List<ParamEntry> params,
                                ResourceHandler handler) {
        /** Annotation-based constructor. */
        public ResourceEntry(String uri, String name, String description, String mimeType,
                             Method method, Object instance, List<ParamEntry> params) {
            this(uri, name, description, mimeType, method, instance, params, null);
        }

        /** Programmatic constructor. */
        public ResourceEntry(String uri, String name, String description, String mimeType,
                             List<ParamEntry> params, ResourceHandler handler) {
            this(uri, name, description, mimeType, null, null, params, handler);
        }

        public boolean isDynamic() {
            return handler != null;
        }
    }

    /**
     * Metadata for a registered MCP prompt.
     */
    public record PromptEntry(String name, String description, Method method, Object instance,
                              List<ParamEntry> params, PromptHandler handler) {
        /** Annotation-based constructor. */
        public PromptEntry(String name, String description, Method method, Object instance,
                           List<ParamEntry> params) {
            this(name, description, method, instance, params, null);
        }

        /** Programmatic constructor. */
        public PromptEntry(String name, String description, List<ParamEntry> params,
                           PromptHandler handler) {
            this(name, description, null, null, params, handler);
        }

        public boolean isDynamic() {
            return handler != null;
        }
    }

    /**
     * Metadata for a method parameter.
     */
    public record ParamEntry(String name, String description, boolean required, Class<?> type) {}

    /**
     * MCP 2025-06-18 / 2025-11-25 spec extensions: human-friendly display
     * {@code title}, optional {@code iconUrl}, and an open {@code _meta}
     * passthrough blob. Carried as a sidecar to the {@link ToolEntry} /
     * {@link ResourceEntry} / {@link PromptEntry} records so existing
     * record signatures (and their many constructor call-sites) don't
     * have to change. Empty/absent fields are simply omitted from the
     * wire response.
     */
    public record EntryMetadata(String title, String iconUrl, Map<String, Object> meta) {
        public EntryMetadata {
            title = title == null ? "" : title;
            iconUrl = iconUrl == null ? "" : iconUrl;
            meta = meta == null ? Map.of() : Map.copyOf(meta);
        }
        public boolean isEmpty() {
            return title.isEmpty() && iconUrl.isEmpty() && meta.isEmpty();
        }
    }

    private final Map<String, ToolEntry> tools = new ConcurrentHashMap<>();
    private final Map<String, ResourceEntry> resources = new ConcurrentHashMap<>();
    private final Map<String, PromptEntry> prompts = new ConcurrentHashMap<>();
    private final Map<String, EntryMetadata> toolMetadata = new ConcurrentHashMap<>();
    private final Map<String, EntryMetadata> resourceMetadata = new ConcurrentHashMap<>();
    private final Map<String, EntryMetadata> promptMetadata = new ConcurrentHashMap<>();

    /**
     * Scan the given instance for @McpTool, @McpResource, @McpPrompt methods.
     */
    public void scan(Object instance) {
        for (var method : instance.getClass().getMethods()) {
            if (method.isAnnotationPresent(McpTool.class)) {
                var a = method.getAnnotation(McpTool.class);
                var params = extractParams(method);
                tools.put(a.name(), new ToolEntry(a.name(), a.description(), method, instance, params));
                var meta = new EntryMetadata(a.title(), a.iconUrl(), Map.of());
                if (!meta.isEmpty()) {
                    toolMetadata.put(a.name(), meta);
                }
            }
            if (method.isAnnotationPresent(McpResource.class)) {
                var a = method.getAnnotation(McpResource.class);
                var params = extractParams(method);
                resources.put(a.uri(), new ResourceEntry(a.uri(), a.name(), a.description(),
                        a.mimeType(), method, instance, params));
                var meta = new EntryMetadata(a.title(), a.iconUrl(), Map.of());
                if (!meta.isEmpty()) {
                    resourceMetadata.put(a.uri(), meta);
                }
            }
            if (method.isAnnotationPresent(McpPrompt.class)) {
                var a = method.getAnnotation(McpPrompt.class);
                var params = extractParams(method);
                prompts.put(a.name(), new PromptEntry(a.name(), a.description(), method, instance, params));
                var meta = new EntryMetadata(a.title(), a.iconUrl(), Map.of());
                if (!meta.isEmpty()) {
                    promptMetadata.put(a.name(), meta);
                }
            }
        }
    }

    // ── Programmatic Tool Registration ───────────────────────────────────

    /**
     * Register a tool with a lambda handler.
     *
     * @param name        tool name (unique identifier)
     * @param description human-readable description for AI agents
     * @param params      parameter metadata (name, description, required, type)
     * @param handler     function that receives arguments and returns the result
     */
    public void registerTool(String name, String description, List<ParamEntry> params,
                             ToolHandler handler) {
        tools.put(name, new ToolEntry(name, description, params, handler));
    }

    /**
     * Register a tool with no parameters.
     */
    public void registerTool(String name, String description, ToolHandler handler) {
        registerTool(name, description, List.of(), handler);
    }

    /**
     * Remove a previously registered tool.
     *
     * @return true if the tool was found and removed
     */
    public boolean removeTool(String name) {
        return tools.remove(name) != null;
    }

    // ── Programmatic Resource Registration ───────────────────────────────

    /**
     * Register a resource with a lambda handler.
     */
    public void registerResource(String uri, String name, String description,
                                 String mimeType, List<ParamEntry> params,
                                 ResourceHandler handler) {
        resources.put(uri, new ResourceEntry(uri, name, description, mimeType, params, handler));
    }

    /**
     * Register a resource with no parameters.
     */
    public void registerResource(String uri, String name, String description,
                                 String mimeType, ResourceHandler handler) {
        registerResource(uri, name, description, mimeType, List.of(), handler);
    }

    /**
     * Remove a previously registered resource.
     */
    public boolean removeResource(String uri) {
        return resources.remove(uri) != null;
    }

    // ── Programmatic Prompt Registration ─────────────────────────────────

    /**
     * Register a prompt with a lambda handler.
     */
    public void registerPrompt(String name, String description, List<ParamEntry> params,
                               PromptHandler handler) {
        prompts.put(name, new PromptEntry(name, description, params, handler));
    }

    /**
     * Register a prompt with no parameters.
     */
    public void registerPrompt(String name, String description, PromptHandler handler) {
        registerPrompt(name, description, List.of(), handler);
    }

    /**
     * Remove a previously registered prompt.
     */
    public boolean removePrompt(String name) {
        return prompts.remove(name) != null;
    }

    // ── Queries ──────────────────────────────────────────────────────────

    public Map<String, ToolEntry> tools() {
        return Collections.unmodifiableMap(tools);
    }

    /**
     * Spec-extension metadata for a registered tool, if any.
     * @see EntryMetadata
     */
    public Optional<EntryMetadata> toolMetadata(String name) {
        return Optional.ofNullable(toolMetadata.get(name));
    }

    public Optional<EntryMetadata> resourceMetadata(String uri) {
        return Optional.ofNullable(resourceMetadata.get(uri));
    }

    public Optional<EntryMetadata> promptMetadata(String name) {
        return Optional.ofNullable(promptMetadata.get(name));
    }

    /**
     * Programmatically attach spec-extension metadata to a registered tool.
     * Useful when registering tools via {@link #registerTool} (which doesn't
     * carry annotation defaults). Pass {@code null} or an empty
     * {@link EntryMetadata} to remove.
     */
    public void setToolMetadata(String name, EntryMetadata metadata) {
        if (metadata == null || metadata.isEmpty()) {
            toolMetadata.remove(name);
        } else {
            toolMetadata.put(name, metadata);
        }
    }

    public void setResourceMetadata(String uri, EntryMetadata metadata) {
        if (metadata == null || metadata.isEmpty()) {
            resourceMetadata.remove(uri);
        } else {
            resourceMetadata.put(uri, metadata);
        }
    }

    public void setPromptMetadata(String name, EntryMetadata metadata) {
        if (metadata == null || metadata.isEmpty()) {
            promptMetadata.remove(name);
        } else {
            promptMetadata.put(name, metadata);
        }
    }

    public Map<String, ResourceEntry> resources() {
        return Collections.unmodifiableMap(resources);
    }

    public Map<String, PromptEntry> prompts() {
        return Collections.unmodifiableMap(prompts);
    }

    public Optional<ToolEntry> tool(String name) {
        return Optional.ofNullable(tools.get(name));
    }

    public Optional<ResourceEntry> resource(String uri) {
        return Optional.ofNullable(resources.get(uri));
    }

    public Optional<PromptEntry> prompt(String name) {
        return Optional.ofNullable(prompts.get(name));
    }

    /**
     * Generate JSON Schema-like input schema for a tool's parameters.
     */
    public static Map<String, Object> inputSchema(ToolEntry tool) {
        var properties = new LinkedHashMap<String, Object>();
        var required = new ArrayList<String>();

        for (var param : tool.params()) {
            var prop = new LinkedHashMap<String, Object>();
            prop.put("type", jsonSchemaType(param.type()));
            if (!param.description().isEmpty()) {
                prop.put("description", param.description());
            }
            properties.put(param.name(), prop);
            if (param.required()) {
                required.add(param.name());
            }
        }

        var schema = new LinkedHashMap<String, Object>();
        schema.put("type", "object");
        schema.put("properties", properties);
        if (!required.isEmpty()) {
            schema.put("required", required);
        }
        return schema;
    }

    private static List<ParamEntry> extractParams(Method method) {
        var params = new ArrayList<ParamEntry>();
        for (Parameter p : method.getParameters()) {
            // Skip framework-injectable types — they're not JSON-RPC arguments
            if (isInjectableType(p.getType())) {
                continue;
            }
            var mcpParam = p.getAnnotation(McpParam.class);
            if (mcpParam != null) {
                params.add(new ParamEntry(mcpParam.name(), mcpParam.description(),
                        mcpParam.required(), p.getType()));
            } else {
                // Fall back to parameter name if -parameters compiler flag is used
                params.add(new ParamEntry(p.getName(), "", true, p.getType()));
            }
        }
        return params;
    }

    /**
     * Returns true if the given type is a framework-injectable type that should
     * not appear in the tool's JSON Schema (it's injected by the framework,
     * not supplied by the MCP client).
     */
    public static boolean isInjectableType(Class<?> type) {
        // Core Atmosphere types
        if (type == org.atmosphere.cpr.Broadcaster.class
                || type == org.atmosphere.cpr.AtmosphereConfig.class
                || type == org.atmosphere.cpr.BroadcasterFactory.class
                || type == org.atmosphere.cpr.AtmosphereFramework.class) {
            return true;
        }
        // atmosphere-ai StreamingSession (optional dependency)
        try {
            var streamingSessionClass = Class.forName("org.atmosphere.ai.StreamingSession");
            if (streamingSessionClass.isAssignableFrom(type)) {
                return true;
            }
        } catch (ClassNotFoundException ex) {
            logger.trace("atmosphere-ai not on classpath", ex);
        }
        return false;
    }

    private static String jsonSchemaType(Class<?> type) {
        if (type == String.class) return "string";
        if (type == int.class || type == Integer.class) return "integer";
        if (type == long.class || type == Long.class) return "integer";
        if (type == double.class || type == Double.class) return "number";
        if (type == float.class || type == Float.class) return "number";
        if (type == boolean.class || type == Boolean.class) return "boolean";
        return "string";
    }
}
