/*
 * Copyright 2022-2026 Revetware LLC.
 *
 * 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 com.soklet;

import org.jspecify.annotations.NonNull;
import org.jspecify.annotations.Nullable;

import javax.annotation.concurrent.NotThreadSafe;
import javax.annotation.concurrent.ThreadSafe;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static java.util.Objects.requireNonNull;

/**
 * Immutable programmatic registration for one MCP prompt.
 *
 * <p>Registration is staged so a handler must be supplied before optional
 * descriptor metadata can be configured and the registration built. Prompt
 * arguments are flat strings and do not use JSON Schema.
 *
 * @author <a href="https://www.revetkn.com">Mark Allen</a>
 */
@ThreadSafe
public final class McpPromptRegistration {
        @NonNull
        private final String name;
        @NonNull
        private final Set<@NonNull McpProtocolVersion> protocolVersions;
        @Nullable
        private final String title;
        @Nullable
        private final String description;
        @NonNull
        private final List<@NonNull McpIcon> icons;
        @NonNull
        private final List<@NonNull McpPromptArgumentDeclaration> arguments;
        @NonNull
        private final List<@NonNull McpInputRequestDeclaration> inputRequestDeclarations;
        @NonNull
        private final McpRequestStateMode requestStateMode;
        @NonNull
        private final McpJsonObject metadata;
        @NonNull
        private final McpPromptHandler handler;
        @Nullable
        private final McpCompletionHandler completionHandler;
        @NonNull
        private final Set<@NonNull McpProtocolVersion> completionProtocolVersions;

        /**
         * Begins a staged registration for a named prompt.
         *
         * @param name nonblank prompt name published to MCP clients
         * @param protocolVersions nonempty exact revisions serving this prompt
         * @return required-handler stage
         * @throws IllegalArgumentException if {@code name} is blank
         */
        @NonNull
        public static HandlerStage withName(@NonNull String name,
                        @NonNull Set<@NonNull McpProtocolVersion> protocolVersions) {
                return new HandlerStage(requireName(name),
                                McpProtocolVersion.requiredSet(protocolVersions));
        }

        private McpPromptRegistration(@NonNull Builder builder) {
                this.name = builder.name;
                this.protocolVersions = builder.protocolVersions;
                this.title = builder.title;
                this.description = builder.description;
                this.icons = List.copyOf(builder.icons);
                this.arguments = immutableArguments(builder.arguments);
                this.inputRequestDeclarations =
                                List.copyOf(builder.inputRequestDeclarations);
                this.requestStateMode = builder.requestStateMode;
                this.metadata = builder.metadata;
                this.handler = builder.handler;
                this.completionHandler = builder.completionHandler;
                this.completionProtocolVersions = builder.completionProtocolVersions;
        }

        /** @return published prompt name */
        @NonNull
        public String getName() {
                return this.name;
        }

        /** @return exact protocol revisions on which this prompt is available */
        @NonNull
        public Set<@NonNull McpProtocolVersion> getProtocolVersions() {
                return this.protocolVersions;
        }

        /** @return human-readable title, if configured */
        @NonNull
        public Optional<@NonNull String> getTitle() {
                return Optional.ofNullable(this.title);
        }

        /** @return human-readable description, if configured */
        @NonNull
        public Optional<@NonNull String> getDescription() {
                return Optional.ofNullable(this.description);
        }

        /** @return immutable icon list in registration order */
        @NonNull
        public List<@NonNull McpIcon> getIcons() {
                return this.icons;
        }

        /** @return immutable argument declarations in registration order */
        @NonNull
        public List<@NonNull McpPromptArgumentDeclaration> getArguments() {
                return this.arguments;
        }

        /**
         * Returns the input requests this prompt operation may emit.
         *
         * @return immutable declarations in registration order
         */
        @NonNull
        public List<@NonNull McpInputRequestDeclaration>
                        getInputRequestDeclarations() {
                return this.inputRequestDeclarations;
        }

        /**
         * Returns the request-state contract for this prompt operation.
         *
         * @return request-state mode
         */
        @NonNull
        public McpRequestStateMode getRequestStateMode() {
                return this.requestStateMode;
        }

        /** @return immutable protocol extension metadata */
        @NonNull
        public McpJsonObject getMetadata() {
                return this.metadata;
        }

        /** @return prompt handler */
        @NonNull
        public McpPromptHandler getHandler() {
                return this.handler;
        }

        /** @return configured argument completer, or empty when omitted */
        @NonNull
        public Optional<@NonNull McpCompletionHandler> getCompletionHandler() {
                return Optional.ofNullable(this.completionHandler);
        }

        /** @return exact revisions on which the completion handler is available */
        @NonNull
        public Set<@NonNull McpProtocolVersion> getCompletionProtocolVersions() {
                return this.completionProtocolVersions;
        }

        /**
         * Validates raw string arguments and invokes the prompt handler.
         */
        @NonNull
        McpOperationResult invoke(@NonNull McpRequestContext request,
                        @NonNull McpJsonObject rawArguments,
                        @NonNull McpInvocationFeatures features) throws Exception {
                requireNonNull(request);
                requireNonNull(rawArguments);
                requireNonNull(features);

                Map<String, McpPromptArgumentDeclaration> definitions =
                                new LinkedHashMap<>();
                for (McpPromptArgumentDeclaration argument : this.arguments)
                        definitions.put(argument.getName(), argument);

                Map<String, String> values = new LinkedHashMap<>();
                for (Map.Entry<String, McpJsonValue> entry
                                : rawArguments.getMembers().entrySet()) {
                        if (!definitions.containsKey(entry.getKey())
                                        || !(entry.getValue() instanceof McpJsonString string))
                                throw new McpInvalidPromptArgumentsException();
                        values.put(entry.getKey(), string.getValue());
                }
                for (McpPromptArgumentDeclaration argument : this.arguments) {
                        if (argument.isRequired() && !values.containsKey(argument.getName()))
                                throw new McpInvalidPromptArgumentsException();
                }

                McpPromptGetContext prompt = new DefaultPromptGetContext(values);
                return requireNonNull(this.handler.handle(request, prompt, features),
                                "The MCP prompt handler returned null.");
        }

        @NonNull
        private static String requireName(@NonNull String name) {
                McpContentValueSupport.requireWellFormedString(name);
                if (name.isBlank())
                        throw new IllegalArgumentException(
                                        "MCP prompt names must not be blank.");
                return name;
        }

        @NonNull
        private static List<@NonNull McpPromptArgumentDeclaration> immutableArguments(
                        @NonNull List<@NonNull McpPromptArgumentDeclaration> arguments) {
                List<McpPromptArgumentDeclaration> copied = List.copyOf(arguments);
                Set<String> names = new LinkedHashSet<>();
                for (McpPromptArgumentDeclaration argument : copied) {
                        if (!names.add(argument.getName()))
                                throw new IllegalStateException(
                                                "Duplicate MCP prompt argument name: "
                                                                + argument.getName());
                }
                return copied;
        }

        /**
         * Required-handler stage for a named prompt.
         *
         * @author <a href="https://www.revetkn.com">Mark Allen</a>
         */
        @NotThreadSafe
        public static final class HandlerStage {
                @NonNull
                private final String name;
                @NonNull
                private final Set<@NonNull McpProtocolVersion> protocolVersions;

                private HandlerStage(@NonNull String name,
                                @NonNull Set<@NonNull McpProtocolVersion> protocolVersions) {
                        this.name = McpContentValueSupport.requireWellFormedString(name);
                        this.protocolVersions = requireNonNull(protocolVersions);
                }

                /**
                 * Supplies the required prompt handler.
                 *
                 * @param handler prompt handler
                 * @return optional-metadata builder
                 */
                @NonNull
                public Builder handler(@NonNull McpPromptHandler handler) {
                        return new Builder(this.name, this.protocolVersions,
                                        requireNonNull(handler));
                }
        }

        /**
         * Mutable builder for an immutable prompt registration.
         *
         * @author <a href="https://www.revetkn.com">Mark Allen</a>
         */
        @NotThreadSafe
        public static final class Builder {
                @NonNull
                private final String name;
                @NonNull
                private final Set<@NonNull McpProtocolVersion> protocolVersions;
                @NonNull
                private final McpPromptHandler handler;
                @Nullable
                private String title;
                @Nullable
                private String description;
                @NonNull
                private List<@NonNull McpIcon> icons = List.of();
                @NonNull
                private List<@NonNull McpPromptArgumentDeclaration> arguments =
                                List.of();
                @NonNull
                private List<@NonNull McpInputRequestDeclaration>
                                inputRequestDeclarations = List.of();
                @NonNull
                private McpRequestStateMode requestStateMode = McpRequestStateMode.NONE;
                @NonNull
                private McpJsonObject metadata = McpJsonObject.emptyInstance();
                @Nullable
                private McpCompletionHandler completionHandler;
                @NonNull
                private Set<@NonNull McpProtocolVersion> completionProtocolVersions = Set.of();

                private Builder(@NonNull String name,
                                @NonNull Set<@NonNull McpProtocolVersion> protocolVersions,
                                @NonNull McpPromptHandler handler) {
                        this.name = McpContentValueSupport.requireWellFormedString(name);
                        this.protocolVersions = requireNonNull(protocolVersions);
                        this.handler = requireNonNull(handler);
                }

                /** @param title human-readable title
                 * @return this builder */
                @NonNull
                public Builder title(@NonNull String title) {
                        this.title = McpContentValueSupport.requireWellFormedString(title);
                        return this;
                }

                /** @param description human-readable description
                 * @return this builder */
                @NonNull
                public Builder description(@NonNull String description) {
                        this.description = McpContentValueSupport.requireWellFormedString(description);
                        return this;
                }

                /**
                 * Replaces icon descriptors in supplied order.
                 * Null or empty clears the property. The complete list is validated and
                 * snapshotted before replacing the prior value.
                 *
                 * @param icons icon descriptors, or null to clear
                 * @return this builder
                 * @throws NullPointerException if a list element is null
                 */
                @NonNull
                public Builder icons(
                                @Nullable List<@NonNull McpIcon> icons) {
                        this.icons = icons == null ? List.of()
                                        : List.copyOf(icons);
                        return this;
                }

                /**
                 * Replaces prompt-argument declarations in supplied order.
                 * Null or empty clears the property. The complete list is validated and
                 * snapshotted before replacing the prior value.
                 *
                 * @param argumentDeclarations prompt-argument declarations, or null to clear
                 * @return this builder
                 * @throws NullPointerException if a list element is null
                 */
                @NonNull
                public Builder arguments(
                                @Nullable List<@NonNull McpPromptArgumentDeclaration> argumentDeclarations) {
                        this.arguments = argumentDeclarations == null ? List.of()
                                        : List.copyOf(argumentDeclarations);
                        return this;
                }

                /**
                 * Replaces input-request declarations in supplied order.
                 * Null or empty clears the property. The complete list is validated and
                 * snapshotted before replacing the prior value.
                 *
                 * @param inputRequestDeclarations input-request declarations, or null to clear
                 * @return this builder
                 * @throws NullPointerException if a list element is null
                 */
                @NonNull
                public Builder inputRequestDeclarations(
                                @Nullable List<@NonNull McpInputRequestDeclaration> inputRequestDeclarations) {
                        this.inputRequestDeclarations = inputRequestDeclarations == null ? List.of()
                                        : List.copyOf(inputRequestDeclarations);
                        return this;
                }

                /**
                 * Sets the request-state contract for this prompt operation.
                 *
                 * @param requestStateMode request-state mode
                 * @return this builder
                 */
                @NonNull
                public Builder requestStateMode(
                                @NonNull McpRequestStateMode requestStateMode) {
                        this.requestStateMode = requireNonNull(requestStateMode);
                        return this;
                }

                /** @param metadata protocol extension metadata
                 * @return this builder */
                @NonNull
                public Builder metadata(@NonNull McpJsonObject metadata) {
                        this.metadata = requireNonNull(metadata);
                        return this;
                }

                /**
                 * Configures the handler for partial prompt-argument suggestions.
                 * Completion is supported for MCP 2025-06-18, 2025-11-25, and
                 * 2026-07-28. Repeated calls replace the previous handler and revision
                 * selection.
                 *
                 * @param completionHandler non-null argument completer
                 * @param protocolVersions nonempty subset of this prompt's revisions
                 * @return this builder
                 */
                @NonNull
                public Builder completionHandler(
                                @NonNull McpCompletionHandler completionHandler,
                                @NonNull Set<@NonNull McpProtocolVersion> protocolVersions) {
                        requireNonNull(completionHandler);
                        Set<McpProtocolVersion> versions = McpProtocolVersion.requiredSet(protocolVersions);
                        if (!this.protocolVersions.containsAll(versions))
                                throw new IllegalArgumentException(
                                                "MCP prompt completion revisions must be a subset of prompt revisions.");
                        this.completionHandler = completionHandler;
                        this.completionProtocolVersions = versions;
                        return this;
                }

                /**
                 * Builds an immutable prompt registration.
                 *
                 * @return prompt registration
                 * @throws IllegalStateException if argument names are duplicated
                 */
                @NonNull
                public McpPromptRegistration build() {
                        return new McpPromptRegistration(this);
                }
        }

        /**
         * Immutable prompt argument projection.
         *
         * @author <a href="https://www.revetkn.com">Mark Allen</a>
         */
        @ThreadSafe
        private static final class DefaultPromptGetContext
                        implements McpPromptGetContext {
                @NonNull
                private final Map<@NonNull String, @NonNull String> arguments;

                private DefaultPromptGetContext(
                                @NonNull Map<@NonNull String, @NonNull String> arguments) {
                        this.arguments = Collections.unmodifiableMap(
                                        new LinkedHashMap<>(requireNonNull(arguments)));
                }

                @Override
                @NonNull
                public Map<@NonNull String, @NonNull String> getArguments() {
                        return this.arguments;
                }

                @Override
                @NonNull
                public Optional<@NonNull String> findArgument(@NonNull String name) {
                        return Optional.ofNullable(this.arguments.get(McpContentValueSupport.requireWellFormedString(name)));
                }
        }
}