/*
 * 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.ai.processor;

import org.atmosphere.ai.AiInterceptor;
import org.atmosphere.ai.AiRequest;
import org.atmosphere.ai.PromptLoader;
import org.atmosphere.ai.StreamingSession;
import org.atmosphere.ai.annotation.AiEndpoint;
import org.atmosphere.ai.annotation.Prompt;
import org.atmosphere.cpr.AtmosphereFramework;
import org.atmosphere.cpr.AtmosphereHandler;
import org.atmosphere.cpr.AtmosphereResource;
import org.mockito.ArgumentCaptor;

import java.util.List;

import static org.mockito.Mockito.*;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.*;

@SuppressWarnings({"unchecked", "rawtypes"})
public class AiEndpointProcessorTest {

    private AiEndpointProcessor processor;
    private AtmosphereFramework framework;

    @BeforeEach
    public void setUp() throws Exception {
        processor = new AiEndpointProcessor();
        framework = mock(AtmosphereFramework.class);
    }

    @AfterEach
    public void tearDown() {
        PromptLoader.clearCache();
    }

    @Test
    public void testRegistersHandlerAtConfiguredPath() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        var pathCaptor = ArgumentCaptor.forClass(String.class);
        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        var interceptorsCaptor = ArgumentCaptor.forClass(List.class);
        verify(framework).addAtmosphereHandler(pathCaptor.capture(), handlerCaptor.capture(), interceptorsCaptor.capture());

        assertEquals("/atmosphere/test-ai", pathCaptor.getValue());
        assertNotNull(handlerCaptor.getValue());
        assertTrue(handlerCaptor.getValue() instanceof AiEndpointHandler);
    }

    @Test
    public void testConfiguresPerResourceTimeout() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new CustomTimeoutEndpoint());

        processor.handle(framework, (Class) CustomTimeoutEndpoint.class);

        // Timeout should NOT be set globally
        verify(framework, never()).addInitParameter(
                eq(org.atmosphere.cpr.ApplicationConfig.MAX_INACTIVE), anyString());

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertEquals(60_000L, handler.suspendTimeout());
    }

    @Test
    public void testDefaultTimeout() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        // Timeout should NOT be set globally
        verify(framework, never()).addInitParameter(
                eq(org.atmosphere.cpr.ApplicationConfig.MAX_INACTIVE), anyString());

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertEquals(120_000L, handler.suspendTimeout());
    }

    @Test
    public void testRejectsClassWithoutPromptMethod() throws Exception {
        processor.handle(framework, (Class) NoPromptEndpoint.class);

        verify(framework, never()).addAtmosphereHandler(anyString(), any(AtmosphereHandler.class), any(List.class));
    }

    @Test
    public void testRejectsInvalidPromptSignature() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new InvalidSignatureEndpoint());

        processor.handle(framework, (Class) InvalidSignatureEndpoint.class);

        // Invalid signature should prevent registration
        verify(framework, never()).addAtmosphereHandler(anyString(), any(AtmosphereHandler.class), any(List.class));
    }

    @Test
    public void testAcceptsThreeParamPrompt() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ThreeParamEndpoint());

        processor.handle(framework, (Class) ThreeParamEndpoint.class);

        verify(framework).addAtmosphereHandler(eq("/atmosphere/three-param"), any(AtmosphereHandler.class), any(List.class));
    }

    @Test
    public void testHandlerTargetIsCorrectInstance() throws Exception {
        var instance = new ValidEndpoint();
        when(framework.newClassInstance(eq(Object.class), any())).thenReturn(instance);

        processor.handle(framework, (Class) ValidEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));

        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertSame(handler.target(), instance);
        assertEquals("onPrompt", handler.promptMethod().getName());
    }

    @Test
    public void testSkipsClassWithoutAnnotation() throws Exception {
        processor.handle(framework, (Class) NotAnnotated.class);

        verify(framework, never()).addAtmosphereHandler(anyString(), any(AtmosphereHandler.class), any(List.class));
    }

    @Test
    public void testSystemPromptPassedToHandler() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new SystemPromptEndpoint());

        processor.handle(framework, (Class) SystemPromptEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertEquals("You are a helpful assistant.", handler.systemPrompt());
    }

    @Test
    public void testDefaultSystemPromptIsEmpty() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertEquals("", handler.systemPrompt());
    }

    @Test
    public void testSystemPromptFromResource() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ResourcePromptEndpoint());

        processor.handle(framework, (Class) ResourcePromptEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertEquals("You are a helpful assistant.", handler.systemPrompt());
    }

    @Test
    public void testResourceTakesPrecedenceOverInline() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new BothPromptsEndpoint());

        processor.handle(framework, (Class) BothPromptsEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        // Resource content should win over the inline systemPrompt
        assertEquals("You are a helpful assistant.", handler.systemPrompt());
    }

    @Test
    public void testRejectsMultiplePromptMethods() throws Exception {
        processor.handle(framework, (Class) MultiplePromptEndpoint.class);

        // Multiple @Prompt methods should prevent registration
        verify(framework, never()).addAtmosphereHandler(anyString(), any(AtmosphereHandler.class), any(List.class));
    }

    @Test
    public void testHandlerHasRuntime() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertNotNull(handler.runtime());
    }

    @Test
    public void testHandlerHasInterceptors() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new InterceptorEndpoint());
        when(framework.newClassInstance(eq(AiInterceptor.class), eq(TestInterceptor.class)))
                .thenReturn(new TestInterceptor());

        processor.handle(framework, (Class) InterceptorEndpoint.class);

        // Verify DI was used instead of raw reflection
        verify(framework).newClassInstance(AiInterceptor.class, TestInterceptor.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertEquals(1, handler.interceptors().size());
        assertTrue(handler.interceptors().get(0) instanceof TestInterceptor);
    }

    @Test
    public void testHandlerHasEmptyInterceptorsWhenNoneSpecified() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertTrue(handler.interceptors().isEmpty());
    }

    @Test
    public void testModelOverridePassedToHandler() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ModelOverrideEndpoint());

        processor.handle(framework, (Class) ModelOverrideEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertEquals("gpt-4o", handler.endpointModel());
    }

    @Test
    public void testDefaultModelIsNull() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        assertNull(handler.endpointModel());
    }

    @Test
    public void testDefaultEndpointDoesNotAutoDiscoverContextProviders() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        // autoDiscoverContextProviders defaults to false, so no ServiceLoader discovery
        assertTrue(handler.contextProviders().isEmpty());
    }

    @Test
    public void testStreamCacheAppliedToFramework() throws Exception {
        // @AiEndpoint(streamCache = UUIDBroadcasterCache.class) must call
        // setBroadcasterCacheClassName before addAtmosphereHandler so the
        // newly-created broadcaster picks up the cache class.
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new StreamCacheEndpoint());

        processor.handle(framework, (Class) StreamCacheEndpoint.class);

        verify(framework).setBroadcasterCacheClassName(
                eq("org.atmosphere.cache.UUIDBroadcasterCache"));
    }

    @Test
    public void testStreamCacheDefaultDoesNotMutateFramework() throws Exception {
        // @AiEndpoint without streamCache (or with the default) must NOT touch
        // setBroadcasterCacheClassName so the framework-wide default is preserved.
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());

        processor.handle(framework, (Class) ValidEndpoint.class);

        verify(framework, never()).setBroadcasterCacheClassName(anyString());
    }

    @Test
    public void testHeartbeatOverrideAppliedToInterceptor() throws Exception {
        // Verify @AiEndpoint.heartbeatSeconds reconfigures the per-endpoint
        // HeartbeatInterceptor instance produced by defaultManagedServiceInterceptors.
        // Closes the long-tool-approval gap where intermediate proxies close
        // idle streams while a parked virtual thread waits on @RequiresApproval.
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new HeartbeatOverrideEndpoint());
        when(framework.excludedInterceptors()).thenReturn(java.util.Collections.emptyList());
        when(framework.newClassInstance(eq(org.atmosphere.cpr.AtmosphereInterceptor.class), any()))
                .thenAnswer(inv -> {
                    Class<?> clazz = inv.getArgument(1);
                    return clazz.getDeclaredConstructor().newInstance();
                });

        processor.handle(framework, (Class) HeartbeatOverrideEndpoint.class);

        var interceptorsCaptor = ArgumentCaptor.forClass(List.class);
        verify(framework).addAtmosphereHandler(anyString(), any(AtmosphereHandler.class),
                interceptorsCaptor.capture());

        var heartbeat = ((List<?>) interceptorsCaptor.getValue()).stream()
                .filter(i -> i instanceof org.atmosphere.interceptor.HeartbeatInterceptor)
                .map(i -> (org.atmosphere.interceptor.HeartbeatInterceptor) i)
                .findFirst()
                .orElseThrow(() ->
                        new AssertionError("HeartbeatInterceptor expected in default chain"));

        assertEquals(45, heartbeat.heartbeatFrequencyInSeconds(),
                "@AiEndpoint(heartbeatSeconds=45) must reconfigure the HeartbeatInterceptor");
    }

    @Test
    public void testHeartbeatDisabledViaNegativeOne() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new HeartbeatDisabledEndpoint());
        when(framework.excludedInterceptors()).thenReturn(java.util.Collections.emptyList());
        when(framework.newClassInstance(eq(org.atmosphere.cpr.AtmosphereInterceptor.class), any()))
                .thenAnswer(inv -> {
                    Class<?> clazz = inv.getArgument(1);
                    return clazz.getDeclaredConstructor().newInstance();
                });

        processor.handle(framework, (Class) HeartbeatDisabledEndpoint.class);

        var interceptorsCaptor = ArgumentCaptor.forClass(List.class);
        verify(framework).addAtmosphereHandler(anyString(), any(AtmosphereHandler.class),
                interceptorsCaptor.capture());

        var heartbeats = ((List<?>) interceptorsCaptor.getValue()).stream()
                .filter(i -> i instanceof org.atmosphere.interceptor.HeartbeatInterceptor)
                .toList();
        assertTrue(heartbeats.isEmpty(),
                "@AiEndpoint(heartbeatSeconds=-1) must NOT attach a per-endpoint HeartbeatInterceptor");
    }

    @Test
    public void testHeartbeatZeroDoesNotAttachInterceptor() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new ValidEndpoint());
        when(framework.excludedInterceptors()).thenReturn(java.util.Collections.emptyList());
        when(framework.newClassInstance(eq(org.atmosphere.cpr.AtmosphereInterceptor.class), any()))
                .thenAnswer(inv -> {
                    Class<?> clazz = inv.getArgument(1);
                    return clazz.getDeclaredConstructor().newInstance();
                });

        processor.handle(framework, (Class) ValidEndpoint.class);

        var interceptorsCaptor = ArgumentCaptor.forClass(List.class);
        verify(framework).addAtmosphereHandler(anyString(), any(AtmosphereHandler.class),
                interceptorsCaptor.capture());

        var heartbeats = ((List<?>) interceptorsCaptor.getValue()).stream()
                .filter(i -> i instanceof org.atmosphere.interceptor.HeartbeatInterceptor)
                .toList();
        // Default (heartbeatSeconds=0) inherits the framework-wide configuration.
        // No per-endpoint HeartbeatInterceptor is added; the endpoint either uses
        // the global one (if HEARTBEAT_INTERVAL_IN_SECONDS init param is set) or
        // has no heartbeat at all.
        assertTrue(heartbeats.isEmpty(),
                "default heartbeatSeconds=0 must not attach a per-endpoint HeartbeatInterceptor");
    }

    @Test
    public void testAutoDiscoverContextProvidersOptIn() throws Exception {
        when(framework.newClassInstance(eq(Object.class), any()))
                .thenReturn(new AutoDiscoverEndpoint());

        processor.handle(framework, (Class) AutoDiscoverEndpoint.class);

        var handlerCaptor = ArgumentCaptor.forClass(AtmosphereHandler.class);
        verify(framework).addAtmosphereHandler(anyString(), handlerCaptor.capture(), any(List.class));
        var handler = (AiEndpointHandler) handlerCaptor.getValue();
        // ServiceLoader will run; whether providers are found depends on classpath,
        // but the flag should not cause errors
        assertNotNull(handler.contextProviders());
    }

    // ---- Test fixture classes ----

    /**
     * Test interceptor that augments the message.
     */
    public static class TestInterceptor implements AiInterceptor {
        @Override
        public AiRequest preProcess(AiRequest request, AtmosphereResource resource) {
            return request.withMessage("[augmented] " + request.message());
        }
    }

    @AiEndpoint(path = "/atmosphere/interceptor-test", interceptors = {TestInterceptor.class})
    public static class InterceptorEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/test-ai")
    public static class ValidEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/custom-timeout", timeout = 60_000L)
    public static class CustomTimeoutEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/no-prompt")
    public static class NoPromptEndpoint {
        public void notAPrompt(String message) {
        }
    }

    @AiEndpoint(path = "/atmosphere/invalid-sig")
    public static class InvalidSignatureEndpoint {
        @Prompt
        public void onPrompt(int notAString) {
        }
    }

    @AiEndpoint(path = "/atmosphere/three-param")
    public static class ThreeParamEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session, AtmosphereResource resource) {
        }
    }

    @AiEndpoint(path = "/atmosphere/system-prompt", systemPrompt = "You are a helpful assistant.")
    public static class SystemPromptEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    public static class NotAnnotated {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/resource-prompt",
                systemPromptResource = "prompts/test-system-prompt.md")
    public static class ResourcePromptEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/both-prompts",
                systemPrompt = "This should be ignored.",
                systemPromptResource = "prompts/test-system-prompt.md")
    public static class BothPromptsEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/model-override", model = "gpt-4o")
    public static class ModelOverrideEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/auto-discover", autoDiscoverContextProviders = true)
    public static class AutoDiscoverEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/stream-cache",
                streamCache = org.atmosphere.cache.UUIDBroadcasterCache.class)
    public static class StreamCacheEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/heartbeat-override", heartbeatSeconds = 45)
    public static class HeartbeatOverrideEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/heartbeat-disabled", heartbeatSeconds = -1)
    public static class HeartbeatDisabledEndpoint {
        @Prompt
        public void onPrompt(String message, StreamingSession session) {
        }
    }

    @AiEndpoint(path = "/atmosphere/multi-prompt")
    public static class MultiplePromptEndpoint {
        @Prompt
        public void firstPrompt(String message, StreamingSession session) {
        }

        @Prompt
        public void secondPrompt(String message, StreamingSession session) {
        }
    }
}
