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

import io.agentscope.core.ReActAgent;
import io.agentscope.core.agent.Event;
import io.agentscope.core.agent.StreamOptions;
import io.agentscope.core.message.Msg;
import io.agentscope.core.message.MsgRole;
import org.atmosphere.ai.AbstractAgentRuntime;
import org.atmosphere.ai.AgentExecutionContext;
import org.atmosphere.ai.AiCapability;
import org.atmosphere.ai.AiConfig;
import org.atmosphere.ai.ExecutionHandle;
import org.atmosphere.ai.StreamingSession;
import org.atmosphere.ai.TokenUsage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Disposable;

import java.util.ArrayList;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;

/**
 * {@link org.atmosphere.ai.AgentRuntime} adapter for the AgentScope Java
 * framework (Tongyi Lab — {@code io.agentscope:agentscope-core}). Wraps a
 * {@link ReActAgent} so its native {@code Flux<Event>} streaming surface
 * forwards token deltas through Atmosphere's {@link StreamingSession}.
 *
 * <p>Auto-detected when {@code io.agentscope.core.ReActAgent} is on the
 * classpath. The agent itself is wired via Spring (
 * {@link AtmosphereAgentScopeAutoConfiguration}) — AgentScope models
 * (DashScope, OpenAI-compatible) need provider-specific construction that
 * can't be done from {@code llm.*} properties alone, so the runtime accepts
 * a user-built {@link ReActAgent} bean rather than building one itself.</p>
 *
 * <p><b>Capabilities:</b> {@link AiCapability#TEXT_STREAMING},
 * {@link AiCapability#SYSTEM_PROMPT}, {@link AiCapability#STRUCTURED_OUTPUT}
 * (via the pipeline), {@link AiCapability#CONVERSATION_MEMORY},
 * {@link AiCapability#TOKEN_USAGE}, {@link AiCapability#TOOL_CALLING}, and
 * {@link AiCapability#TOOL_APPROVAL}. Tools are bridged per request into an
 * AgentScope {@code Toolkit} through {@link AgentScopeToolBridge}, so
 * Atmosphere's validation and HITL gate remain authoritative.</p>
 */
public class AgentScopeAgentRuntime extends AbstractAgentRuntime<ReActAgent> {

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

    private static volatile ReActAgent staticAgent;

    /**
     * Inject a pre-built {@link ReActAgent} from Spring auto-configuration.
     * Mirrors the {@code KoogAgentRuntime.setPromptExecutor} /
     * {@code LangChain4jAgentRuntime.setModel} pattern — the autoconfig
     * holds the Spring-managed agent bean and seeds the runtime statically
     * so the {@link java.util.ServiceLoader}-instantiated runtime picks it
     * up without needing constructor injection.
     */
    public static void setAgent(ReActAgent agent) {
        staticAgent = agent;
    }

    @Override
    public String name() {
        return "agentscope";
    }

    @Override
    protected String nativeClientClassName() {
        return "io.agentscope.core.ReActAgent";
    }

    @Override
    protected String clientDescription() {
        return "ReActAgent";
    }

    @Override
    protected String configurationHint() {
        return "Wire a ReActAgent @Bean (e.g. ReActAgent.builder().model(...).build()) "
                + "or call AgentScopeAgentRuntime.setAgent() before the first request.";
    }

    @Override
    protected ReActAgent createNativeClient(AiConfig.LlmSettings settings) {
        // AgentScope's Model surface (DashScope, OpenAI-compatible) needs
        // provider-specific construction the AiConfig.LlmSettings DTO can't
        // express. Defer to the user's @Bean wired via autoconfig.
        return staticAgent;
    }

    @Override
    public void configure(AiConfig.LlmSettings settings) {
        if (getNativeClient() == null && staticAgent != null) {
            setNativeClient(staticAgent);
        }
        super.configure(settings);
    }

    @Override
    protected void doExecute(ReActAgent agent,
                             AgentExecutionContext context,
                             StreamingSession session) {
        var handle = doExecuteWithHandle(agent, context, session);
        try {
            handle.whenDone().get();
        } catch (InterruptedException ie) {
            Thread.currentThread().interrupt();
        } catch (java.util.concurrent.ExecutionException ex) {
            // Already surfaced via session.error() in the subscribe handler.
        }
    }

    @Override
    protected ExecutionHandle doExecuteWithHandle(ReActAgent agent,
                                                  AgentExecutionContext context,
                                                  StreamingSession session) {
        admitThroughGateway(context);

        // Per-request agent override: when the caller attached a different
        // ReActAgent via AgentScopeAgent.attach(...), dispatch against it
        // instead of the runtime's installed default. Useful when different
        // prompts need different agent topologies (e.g. planner vs. quick
        // lookup) without re-installing the runtime client globally. Bind to
        // a fresh effectively-final local so the cancel handler's anonymous
        // ExecutionHandle inner class can capture the resolved agent.
        var perRequestAgent = AgentScopeAgent.from(context);
        ReActAgent resolvedAgent;
        if (perRequestAgent != null) {
            resolvedAgent = perRequestAgent;
            logger.debug("Dispatching against per-request AgentScope ReActAgent override");
        } else {
            resolvedAgent = agent;
        }
        if (!context.tools().isEmpty()) {
            resolvedAgent = agentWithTools(resolvedAgent, context, session);
        }
        final ReActAgent activeAgent = resolvedAgent;

        var msgs = new ArrayList<Msg>();
        for (var chat : assembleMessages(context)) {
            msgs.add(Msg.builder()
                    .role(toRole(chat.role()))
                    .textContent(chat.content())
                    .build());
        }

        var done = new CompletableFuture<Void>();
        var cancelled = new java.util.concurrent.atomic.AtomicBoolean();

        // Model-lifecycle hooks: same posture as Spring AI / LC4j / ADK / Koog
        // / Embabel / SK. fireModelStart synchronously before subscribe;
        // fireModelEnd on completion (with the last captured TokenUsage from
        // handleEvent); fireModelError on the error callback.
        var listeners = context.listeners();
        var modelName = context.model() != null ? context.model() : name();
        var messageCount = msgs.size();
        var toolCount = context.tools().size();
        var startNanos = System.nanoTime();
        var lastUsage = new java.util.concurrent.atomic.AtomicReference<TokenUsage>();
        org.atmosphere.ai.AgentLifecycleListener.fireModelStart(
                listeners, modelName, messageCount, toolCount);

        Disposable subscription = activeAgent.stream(msgs, StreamOptions.defaults())
                .subscribe(
                        event -> handleEvent(event, session, lastUsage),
                        error -> {
                            // Boundary safety (Invariant #4): every terminal
                            // path must close the session. error() is
                            // idempotent on Atmosphere's StreamingSession.
                            org.atmosphere.ai.AgentLifecycleListener.fireModelError(
                                    listeners, modelName, error);
                            session.error(error);
                            done.completeExceptionally(error);
                        },
                        () -> {
                            long durationMs = (System.nanoTime() - startNanos) / 1_000_000L;
                            org.atmosphere.ai.AgentLifecycleListener.fireModelEnd(
                                    listeners, modelName, lastUsage.get(), durationMs);
                            if (!session.isClosed()) {
                                session.complete();
                            }
                            done.complete(null);
                        });

        return new ExecutionHandle() {
            @Override
            public void cancel() {
                if (cancelled.compareAndSet(false, true)) {
                    // AgentScope exposes interrupt() for cooperative cancel
                    // and the Reactor Disposable for unsubscription. Fire
                    // both: interrupt() unwinds the agent's internal loop;
                    // dispose() drops any in-flight emission to the session.
                    try {
                        activeAgent.interrupt();
                    } catch (RuntimeException re) {
                        logger.trace("ReActAgent.interrupt() threw on cancel", re);
                    }
                    subscription.dispose();
                    if (!done.isDone()) {
                        done.completeExceptionally(new CancellationException("cancelled"));
                    }
                }
            }

            @Override
            public boolean isDone() {
                return done.isDone();
            }

            @Override
            public CompletableFuture<Void> whenDone() {
                return done;
            }
        };
    }

    private static void handleEvent(
            Event event,
            StreamingSession session,
            java.util.concurrent.atomic.AtomicReference<TokenUsage> lastUsage) {
        if (event == null || event.getMessage() == null) {
            return;
        }
        var msg = event.getMessage();
        var text = msg.getTextContent();
        if (text != null && !text.isEmpty()) {
            session.send(text);
        }
        // Token usage rides on the terminal Msg. Forwarding TOKEN_USAGE on
        // every chunk would double-count; AgentScope sets ChatUsage only on
        // the message that completes the round. Stash to lastUsage so the
        // caller's onModelEnd lifecycle hook reports it.
        var usage = msg.getChatUsage();
        if (usage != null && event.isLast()) {
            var tokenUsage = new TokenUsage(
                    usage.getInputTokens(),
                    usage.getOutputTokens(),
                    0L,
                    usage.getTotalTokens(),
                    null);
            session.usage(tokenUsage);
            lastUsage.set(tokenUsage);
        }
    }

    private static ReActAgent agentWithTools(
            ReActAgent baseAgent, AgentExecutionContext context, StreamingSession session) {
        var toolkit = AgentScopeToolBridge.toToolkit(
                context.tools(), session, context.approvalStrategy(), context.approvalPolicy());
        return ReActAgent.builder()
                .name(baseAgent.getClass().getSimpleName())
                .sysPrompt(baseAgent.getSysPrompt())
                .model(baseAgent.getModel())
                .toolkit(toolkit)
                .memory(baseAgent.getMemory())
                .maxIters(baseAgent.getMaxIters())
                .generateOptions(baseAgent.getGenerateOptions())
                .build();
    }

    private static MsgRole toRole(String role) {
        if (role == null) {
            return MsgRole.USER;
        }
        return switch (role) {
            case "system" -> MsgRole.SYSTEM;
            case "assistant" -> MsgRole.ASSISTANT;
            case "tool" -> MsgRole.TOOL;
            default -> MsgRole.USER;
        };
    }

    @Override
    public Set<AiCapability> capabilities() {
        return Set.of(
                AiCapability.TEXT_STREAMING,
                AiCapability.SYSTEM_PROMPT,
                // AiPipeline wraps the session in StructuredOutputCapturingSession
                // and augments the system prompt with schema instructions for
                // every SYSTEM_PROMPT-capable runtime — the contract test
                // pins this pairing.
                AiCapability.STRUCTURED_OUTPUT,
                // assembleMessages threads context.history() into the Msg
                // list the agent receives, so prior turns are honored.
                AiCapability.CONVERSATION_MEMORY,
                // handleEvent forwards Msg.getChatUsage() on isLast().
                AiCapability.TOKEN_USAGE,
                // AgentScopeToolBridge registers Atmosphere ToolDefinition
                // instances as AgentScope AgentTool objects on a per-request
                // Toolkit, then routes each invocation through the shared
                // ToolExecutionHelper approval/validation seam.
                AiCapability.TOOL_CALLING,
                AiCapability.TOOL_APPROVAL,
                // executeWithHandle returns an ExecutionHandle whose cancel()
                // calls activeAgent.interrupt() (cooperative unwind of the
                // ReAct loop), disposes the Reactor subscription (drops
                // in-flight emission), and resolves the done-future
                // exceptionally with CancellationException — terminal-path
                // closure per Correctness Invariant #2.
                AiCapability.CANCELLATION,
                // Inherits AbstractAgentRuntime.executeWithOuterRetry — does
                // not override ownsPerRequestRetry(), so context.retryPolicy()
                // with maxRetries > 0 retries doExecute on pre-stream
                // RuntimeException. See modules/ai/README.md
                // "Per-Request Retry Architecture".
                AiCapability.PER_REQUEST_RETRY,
                // BUDGET_ENFORCEMENT: framework-level circuit breaker via the
                // AiPipeline BudgetCapturingSession decorator — honest because
                // handleEvent forwards Msg.getChatUsage() through session.usage()
                // which is the signal the decorator taps for token / step abort.
                AiCapability.BUDGET_ENFORCEMENT,
                // CONFIDENCE_SCORES: framework-level — AiPipeline's
                // ConfidenceCapturingSession parses the model-reported
                // confidence field on stream completion. Honest because
                // AgentScope honors SYSTEM_PROMPT and streams response text.
                AiCapability.CONFIDENCE_SCORES,
                // PASSIVATION: AgentPassivation snapshots context.history()
                // into a CheckpointStore. Honest because assembleMessages
                // threads history into the Msg list AgentScope receives.
                AiCapability.PASSIVATION);
    }

    @Override
    public java.util.List<String> models() {
        var settings = AiConfig.get();
        if (settings == null || settings.model() == null || settings.model().isBlank()) {
            return java.util.List.of();
        }
        return java.util.List.of(settings.model());
    }
}
