/*
 * 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.embabel

import com.embabel.agent.api.common.Ai
import com.embabel.agent.api.common.streaming.StreamingPromptRunner
import com.embabel.agent.api.streaming.StreamingPromptRunnerBuilder
import com.embabel.agent.core.AgentPlatform
import com.embabel.agent.core.ProcessOptions
import com.embabel.agent.domain.library.HasContent
import com.embabel.agent.spi.config.spring.InfrastructureInjectionConfiguration
import com.embabel.chat.AssistantMessage
import com.embabel.chat.Message
import com.embabel.chat.SystemMessage
import com.embabel.chat.UserMessage
import org.atmosphere.ai.AiCapability
import org.atmosphere.ai.AiConfig
import org.atmosphere.ai.AgentExecutionContext
import org.atmosphere.ai.AgentRuntime
import org.atmosphere.ai.StreamingSession
import org.atmosphere.ai.TokenUsage
import org.slf4j.LoggerFactory

/**
 * [AgentRuntime] implementation backed by the Embabel Agent Platform.
 *
 * Dispatches in two modes depending on how the caller addresses the runtime:
 *
 * 1. **Deployed-agent mode** — when [AgentExecutionContext.agentId] resolves to
 *    an agent already deployed on the [AgentPlatform] (or the configured
 *    [agentName] default), the request is handed to Embabel's goal-planner via
 *    [AgentPlatform.runAgentFrom]. The deployed `@Agent` owns its own system
 *    prompt, tools, and history; Atmosphere threads only `userMessage` into
 *    the process inputs. This preserves Embabel's value proposition as a
 *    goal-driven planning runtime.
 *
 * 2. **Atmosphere-native mode** — when the caller's `agentId()` does NOT match
 *    a deployed agent, the runtime falls back to a direct LLM path via
 *    Embabel's [Ai] factory ([InfrastructureInjectionConfiguration.aiFactory]).
 *    In this mode the bridge honestly threads [AgentExecutionContext.systemPrompt]
 *    and [AgentExecutionContext.history] into the outgoing [Ai.withDefaultLlm]
 *    `PromptRunner`, making `SYSTEM_PROMPT` and `CONVERSATION_MEMORY` honest
 *    capability declarations per Correctness Invariant #5 (Runtime Truth).
 *
 * Auto-detected when `embabel-agent-api` is on the classpath. The
 * [AgentPlatform] must be configured via [setAgentPlatform] — typically done
 * by Spring auto-configuration.
 */
class EmbabelAgentRuntime : AgentRuntime {

    companion object {
        private val logger = LoggerFactory.getLogger(EmbabelAgentRuntime::class.java)

        @Volatile
        private var agentPlatform: AgentPlatform? = null

        @Volatile
        private var agentName: String = "chat-assistant"

        @JvmStatic
        fun setAgentPlatform(platform: AgentPlatform) {
            agentPlatform = platform
        }

        @JvmStatic
        fun setAgentName(name: String) {
            agentName = name
        }
    }

    override fun name(): String = "embabel"

    override fun isAvailable(): Boolean {
        return try {
            Class.forName("com.embabel.agent.core.AgentPlatform")
            true
        } catch (e: ClassNotFoundException) {
            false
        }
    }

    override fun priority(): Int = 100

    override fun configure(settings: AiConfig.LlmSettings) {
        if (agentPlatform != null) {
            return
        }
        logger.info(
            "Embabel adapter active but requires an AgentPlatform bean with deployed @Agent classes. " +
                "Auto-configuration from credentials alone is not supported. " +
                "Ensure embabel-agent-spring-boot-starter is on the classpath."
        )
    }

    override fun execute(context: AgentExecutionContext, session: StreamingSession) {
        // Admit through the process-wide AiGateway before issuing the native
        // Embabel dispatch — uniform per-user rate limiting and credential
        // resolution across all seven runtimes (Correctness Invariant #3).
        org.atmosphere.ai.AbstractAgentRuntime.admitThroughGateway(name(), context)
        // Install the cross-runtime ToolLoopGuard before reading listeners so
        // the guard sees every onModelStart this dispatch fires. No-op when
        // no ToolLoopPolicy is attached. Same install pattern as the Java
        // base class (AbstractAgentRuntime.execute) — keeps the cap behaviour
        // uniform across all runtimes regardless of base class.
        val effective = org.atmosphere.ai.llm.ToolLoopGuard.installIfPresent(name(), context, session)
        val platform = agentPlatform
            ?: throw IllegalStateException(
                "EmbabelAgentRuntime: AgentPlatform not configured. " +
                    "Call EmbabelAgentRuntime.setAgentPlatform() or use Spring auto-configuration."
            )

        session.progress("Connecting to embabel...")

        val targetAgent = effective.agentId() ?: agentName
        val deployedAgent = platform.agents().firstOrNull { it.name == targetAgent }

        // Model-lifecycle hooks: fire onModelStart before dispatch,
        // onModelEnd / onModelError on terminal paths. messageCount = history
        // turns + (system prompt? 1 : 0) + 1 user prompt; toolCount =
        // context.tools(). Same posture as Spring AI / LC4j / ADK / Koog.
        // Token usage is captured inside executeDeployedAgent via
        // process.usage() and stashed in lastUsage so onModelEnd reports the
        // final aggregate.
        val listeners = effective.listeners()
        val modelName = effective.model() ?: name()
        val messageCount = effective.history().size +
            (if (effective.systemPrompt()?.isNotEmpty() == true) 1 else 0) + 1
        val toolCount = effective.tools().size
        val startNanos = System.nanoTime()
        val lastUsage = java.util.concurrent.atomic.AtomicReference<TokenUsage>()
        org.atmosphere.ai.AgentLifecycleListener.fireModelStart(
            listeners, modelName, messageCount, toolCount
        )

        try {
            executeWithOuterRetry(effective, session) {
                if (deployedAgent != null) {
                    warnIfDeployedAgentDropsRequestFeatures(deployedAgent, effective, session)
                    executeDeployedAgent(platform, deployedAgent, effective, session, lastUsage)
                } else {
                    executeAtmosphereNative(platform, effective, session, lastUsage)
                }
            }
            val durationMs = (System.nanoTime() - startNanos) / 1_000_000L
            org.atmosphere.ai.AgentLifecycleListener.fireModelEnd(
                listeners, modelName, lastUsage.get(), durationMs
            )
        } catch (t: Throwable) {
            org.atmosphere.ai.AgentLifecycleListener.fireModelError(listeners, modelName, t)
            throw t
        }
    }

    /**
     * Outer retry wrapper honouring [AgentExecutionContext.retryPolicy] on
     * pre-stream transient failures. Matches the Java base-class semantics
     * in `AbstractAgentRuntime.executeWithOuterRetry`. Embabel does not
     * extend that class (it implements [AgentRuntime] directly) so the
     * retry loop is duplicated here.
     */
    private inline fun executeWithOuterRetry(
        context: AgentExecutionContext,
        session: StreamingSession,
        crossinline body: () -> Unit
    ) {
        val policy = context.retryPolicy()
        if (policy == null || policy.isInheritSentinel || policy.maxRetries() <= 0) {
            body()
            return
        }
        var attempt = 0
        while (true) {
            try {
                body()
                return
            } catch (e: RuntimeException) {
                if (attempt >= policy.maxRetries() || session.hasErrored()) {
                    throw e
                }
                val delay = policy.delayForAttempt(attempt)
                attempt++
                logger.info(
                    "embabel outer-retry attempt {}/{} after {}ms (cause: {})",
                    attempt, policy.maxRetries(), delay.toMillis(), e.message
                )
                if (!delay.isZero && !delay.isNegative) {
                    try {
                        Thread.sleep(delay.toMillis())
                    } catch (ie: InterruptedException) {
                        Thread.currentThread().interrupt()
                        throw e
                    }
                }
            }
        }
    }

    /**
     * Embabel's deployed-agent dispatch reads `userMessage` from the input map
     * but ignores Atmosphere's per-request `tools()` and image `parts()` — the
     * deployed `@Agent` owns its own configuration. When the caller sends a
     * context that carries those features, log a WARN and emit a metadata event
     * so the drop is visible to operators rather than silent. The runtime's
     * `capabilities()` set advertises TOOL_CALLING / TOOL_APPROVAL / VISION /
     * MULTI_MODAL because the Atmosphere-native fallback path honors them; this
     * warning surfaces the drop on the deployed-agent path so partial-matrix
     * behavior is loud (Correctness Invariant #5 — Runtime Truth, applied to
     * dispatch-mode semantics).
     */
    private fun warnIfDeployedAgentDropsRequestFeatures(
        agent: com.embabel.agent.core.Agent,
        context: AgentExecutionContext,
        session: StreamingSession
    ) {
        val droppedTools = context.tools().isNotEmpty()
        val droppedImages = context.parts().any { it is org.atmosphere.ai.Content.Image }
        if (!droppedTools && !droppedImages) {
            return
        }
        val dropped = buildList {
            if (droppedTools) add("${context.tools().size} request-level tool(s)")
            if (droppedImages) {
                val n = context.parts().count { it is org.atmosphere.ai.Content.Image }
                add("$n image part(s)")
            }
        }.joinToString(", ")
        logger.warn(
            "Embabel deployed-agent dispatch [{}] ignores {} — deployed @Agent classes " +
                "own their own tool/image configuration. Invoke without agentId() " +
                "(Atmosphere-native path) to honor request-level features.",
            agent.name, dropped
        )
        session.sendMetadata("ai.embabel.dropped_features", dropped)
    }

    /**
     * Deployed-agent path: hand the invocation to Embabel's goal-planner
     * via [AgentPlatform.runAgentFrom]. The deployed `@Agent` owns its own
     * system prompt, tools, and conversation memory; Atmosphere only binds
     * `userMessage` into the process input map. This preserves Embabel's
     * goal-driven value proposition for callers who deploy `@Agent` classes.
     */
    private fun executeDeployedAgent(
        platform: AgentPlatform,
        agent: com.embabel.agent.core.Agent,
        context: AgentExecutionContext,
        session: StreamingSession,
        lastUsage: java.util.concurrent.atomic.AtomicReference<TokenUsage>
    ) {
        val channel = AtmosphereOutputChannel(session)
        val process = try {
            val options = ProcessOptions.DEFAULT.withOutputChannel(channel)
            platform.runAgentFrom(agent, options, mapOf("userMessage" to context.message()))
        } catch (e: Exception) {
            logger.error("Agent execution failed", e)
            session.error(e)
            return
        }

        // Structured output (e.g. promptedTransformer with typed output) doesn't
        // stream through the OutputChannel. Extract the result from the blackboard
        // and send it to the session so the console renders the response.
        if (!session.isClosed) {
            val result = process.blackboard.lastResult()
            if (result != null) {
                val text = when (result) {
                    is HasContent -> result.content
                    is String -> result
                    else -> result.toString()
                }
                if (!text.isNullOrBlank()) {
                    session.send(text)
                }
            }
        }

        // Emit typed token usage. Embabel's AgentProcess implements
        // LlmInvocationHistory.usage() which aggregates every LLM call made
        // during this process into a single Usage record. The usage is also
        // stashed in lastUsage so the caller's onModelEnd lifecycle hook
        // reports the final aggregate.
        try {
            val usage = process.usage()
            val input = usage.promptTokens?.toLong() ?: 0L
            val output = usage.completionTokens?.toLong() ?: 0L
            val total = usage.totalTokens?.toLong() ?: (input + output)
            val tokenUsage = TokenUsage(input, output, 0L, total, null)
            if (tokenUsage.hasCounts()) {
                session.usage(tokenUsage)
                lastUsage.set(tokenUsage)
            }
        } catch (t: Throwable) {
            logger.debug("Failed to extract Embabel token usage (this is non-fatal)", t)
        }

        if (!session.isClosed) {
            session.complete()
        }
    }

    /**
     * Atmosphere-native path: when no deployed agent matches the requested
     * `agentId()`, fall back to Embabel's [Ai] factory and drive a direct
     * LLM call via [Ai.withDefaultLlm]. This path honestly threads
     * [AgentExecutionContext.systemPrompt] and [AgentExecutionContext.history]
     * into the outgoing request through
     * [com.embabel.agent.api.common.PromptRunner.withSystemPrompt] and
     * [com.embabel.agent.api.common.PromptRunner.withMessages] respectively,
     * which is the mechanism that makes `SYSTEM_PROMPT` and
     * `CONVERSATION_MEMORY` honest capability declarations.
     *
     * **Streaming**: When Embabel's underlying LLM service implements the
     * [com.embabel.agent.spi.streaming.StreamingLlmOperations] contract
     * (true for the OpenAI / Gemini chat-client backends shipped with
     * embabel-agent-spring-boot-starter), this path uses
     * [StreamingPromptRunnerBuilder] to acquire a Reactor `Flux<String>`
     * of token-level chunks and forwards each chunk to [StreamingSession.send]
     * — making `TEXT_STREAMING` an honest capability declaration on the
     * Atmosphere-native path. When streaming is unavailable on the
     * configured runner, the path falls back to the blocking
     * [com.embabel.agent.api.common.PromptRunnerOperations.generateText]
     * call so the dispatch still completes (no silent capability drop).
     *
     * Token usage reporting is not available on this path — Embabel's
     * PromptRunner surface does not expose an aggregated usage record the
     * way `AgentProcess.usage()` does on the runAgentFrom path. Callers who
     * require `ai.tokens.*` telemetry should either (a) use a deployed
     * `@Agent` class (which routes through runAgentFrom) or (b) select a
     * non-Embabel runtime.
     */
    @Suppress("UNUSED_PARAMETER")
    private fun executeAtmosphereNative(
        platform: AgentPlatform,
        context: AgentExecutionContext,
        session: StreamingSession,
        lastUsage: java.util.concurrent.atomic.AtomicReference<TokenUsage>
    ) {
        // lastUsage is intentionally unused on this path: the Atmosphere-native
        // path drives Ai.withDefaultLlm directly without an AgentProcess, so
        // no LlmInvocationHistory.usage() aggregate is available — the per-
        // streamed-call usage holder stays null. This is documented in the
        // function-level KDoc above.
        val ai: Ai = try {
            InfrastructureInjectionConfiguration().aiFactory(platform)
        } catch (t: Throwable) {
            logger.error("Failed to obtain Embabel Ai factory from AgentPlatform", t)
            session.error(t)
            return
        }

        // Embabel's PromptRunner contract: generateText(text) carries the
        // current user prompt; withMessages(...) carries prior history only.
        // Calling generateText("") makes Embabel reject the request with
        // "Text content cannot be empty" before it consults withMessages,
        // so the current message MUST go to generateText, not into messages.
        val historyMessages = buildList<Message> {
            context.history().forEach { msg ->
                when (msg.role()) {
                    "system" -> add(SystemMessage(msg.content()))
                    "assistant" -> add(AssistantMessage(msg.content()))
                    else -> add(UserMessage(msg.content()))
                }
            }
        }
        val currentMessage = context.message().orEmpty().ifBlank { " " }

        // Translate Atmosphere ToolDefinitions and multi-modal Content parts
        // into Embabel-native surfaces so PromptRunner's withTools(...) /
        // withImages(...) surfaces unlock TOOL_CALLING, TOOL_APPROVAL,
        // VISION, and MULTI_MODAL on this path. Every tool invocation
        // routes through ToolExecutionHelper.executeWithApproval so
        // @RequiresApproval gates fire uniformly across runtimes.
        val embabelTools = EmbabelToolBridge.toEmbabelTools(
            context.tools(), session, context.approvalStrategy(), context.approvalPolicy()
        )
        val embabelImages = EmbabelToolBridge.toEmbabelImages(context.parts())

        // Build the configured PromptRunner once — both the streaming and
        // blocking dispatch paths reuse the same withSystemPrompt /
        // withMessages / withTools / withImages wiring, so the only branch
        // is the terminal call (generateStream() vs generateText()).
        // Per-request customizer (EmbabelPromptRunner.attach) applies AFTER
        // the runtime's default wiring so callers can override temperature /
        // model / etc. without losing the system prompt + history we just set.
        val runner = try {
            var r = ai.withDefaultLlm()
            if (!context.systemPrompt().isNullOrBlank()) {
                r = r.withSystemPrompt(context.systemPrompt())
            }
            if (historyMessages.isNotEmpty()) {
                r = r.withMessages(historyMessages)
            }
            if (embabelTools.isNotEmpty()) {
                r = r.withTools(embabelTools)
            }
            if (embabelImages.isNotEmpty()) {
                r = r.withImages(embabelImages)
            }
            // Native ToolLoopPolicy enforcement is NOT available on Embabel
            // 0.3.4 (the pinned version): PromptRunner does not yet expose
            // withToolLoopInspectors / ToolLoopInspector. The cross-runtime
            // ToolLoopGuard installed in execute() provides strict() cap
            // semantics at the wire layer. Native upstream enforcement
            // awaits Embabel 0.3.5 release (tracked in modules/embabel/README.md).
            val customizer = EmbabelPromptRunner.from(context)
            if (customizer != null) {
                r = customizer.apply(r)
                logger.debug("Applied per-request EmbabelPromptRunner customizer")
            }
            r
        } catch (t: Throwable) {
            logger.error("Failed to configure Embabel PromptRunner", t)
            session.error(t)
            return
        }

        // Streaming dispatch path: when Embabel's underlying LLM service
        // implements the StreamingLlmOperations contract,
        // StreamingPromptRunnerBuilder exposes a Flux<String> of token-level
        // chunks. Each chunk lands on session.send() so browser clients see
        // tokens as they arrive instead of a single end-of-call burst.
        // Falls back to the blocking generateText() path when streaming is
        // not supported by the configured model — preserves the
        // pre-streaming behavior for environments where Embabel's
        // StreamingChatClientOperations cannot be wired (e.g. some
        // tool-only or non-OpenAI-compatible providers).
        val streamingRunner = try {
            StreamingPromptRunnerBuilder(runner).streaming() as StreamingPromptRunner.Streaming
        } catch (t: Throwable) {
            logger.debug(
                "Embabel StreamingPromptRunnerBuilder unavailable for this PromptRunner; " +
                    "falling back to blocking generateText() path", t
            )
            null
        }

        if (streamingRunner != null) {
            try {
                streamingRunner
                    .withPrompt(currentMessage)
                    .generateStream()
                    .doOnNext { chunk ->
                        if (!session.isClosed && chunk.isNotEmpty()) {
                            session.send(chunk)
                        }
                    }
                    .doOnError { t ->
                        logger.error("Embabel streaming dispatch failed", t)
                        if (!session.isClosed) {
                            session.error(t)
                        }
                    }
                    .doOnComplete {
                        if (!session.isClosed) {
                            session.complete()
                        }
                    }
                    .blockLast()
                return
            } catch (t: Throwable) {
                // Surface the streaming failure rather than silently retrying
                // through the blocking path — a mid-stream failure has
                // already emitted partial content, and re-dispatching as a
                // fresh blocking call would duplicate output. The doOnError
                // above already wrote session.error; this catch covers
                // pre-subscribe construction errors.
                if (!session.hasErrored() && !session.isClosed) {
                    session.error(t)
                }
                return
            }
        }

        // Blocking fallback — used only when StreamingPromptRunnerBuilder
        // could not build a Streaming instance for this runner (e.g. the
        // model service does not implement StreamingLlmOperations).
        val result: String = try {
            runner.generateText(currentMessage)
        } catch (t: Throwable) {
            logger.error("Embabel Ai dispatch failed", t)
            session.error(t)
            return
        }

        if (!session.isClosed) {
            if (result.isNotBlank()) {
                session.send(result)
            }
            session.complete()
        }
    }

    override fun capabilities(): Set<AiCapability> = setOf(
        AiCapability.TEXT_STREAMING,
        AiCapability.STRUCTURED_OUTPUT,
        AiCapability.AGENT_ORCHESTRATION,
        // SYSTEM_PROMPT is honest via the Atmosphere-native dispatch path
        // (executeAtmosphereNative): Ai.withDefaultLlm().withSystemPrompt(...)
        // threads context.systemPrompt() into the outgoing PromptRunner.
        // Deployed-agent callers own their own prompts via @Agent classes,
        // which is the Embabel-idiomatic way to express a system prompt.
        AiCapability.SYSTEM_PROMPT,
        // CONVERSATION_MEMORY is honest via the Atmosphere-native dispatch
        // path: PromptRunner.withMessages(...) receives the full
        // context.history() translated into embabel.chat.Message variants.
        AiCapability.CONVERSATION_MEMORY,
        // TOKEN_USAGE is honest on the deployed-agent path:
        // executeDeployedAgent reads AgentProcess.usage() (Embabel's
        // LlmInvocationHistory aggregate) and emits a typed TokenUsage
        // record via session.usage(). The Atmosphere-native path does not
        // expose usage — see executeAtmosphereNative Javadoc.
        AiCapability.TOKEN_USAGE,
        // PER_REQUEST_RETRY: honored via executeWithOuterRetry which
        // wraps either dispatch branch (deployed-agent or Atmosphere-native)
        // in a retry loop respecting context.retryPolicy(). Pre-stream
        // transient failures retry on top of Embabel's own retry layer.
        AiCapability.PER_REQUEST_RETRY,
        // TOOL_CALLING / TOOL_APPROVAL: EmbabelToolBridge translates
        // Atmosphere ToolDefinitions into Embabel com.embabel.agent.api.tool.Tool
        // instances wrapped around a Tool.Handler that routes through
        // ToolExecutionHelper.executeWithApproval. PromptRunner.withTools(...)
        // attaches them for SK-style auto-invoke by Embabel's LLM layer.
        AiCapability.TOOL_CALLING,
        AiCapability.TOOL_APPROVAL,
        // VISION / MULTI_MODAL: EmbabelToolBridge.toEmbabelImages maps
        // Atmosphere Content.Image parts into Embabel AgentImage instances,
        // attached via PromptRunner.withImages(...). Only the
        // Atmosphere-native dispatch path (executeAtmosphereNative)
        // reaches the PromptRunner surface; the deployed-agent path
        // (executeDeployedAgent) ignores parts, matching Embabel's own
        // semantics where a deployed @Agent owns its own multi-modal
        // handling.
        AiCapability.VISION,
        AiCapability.MULTI_MODAL,
        // BUDGET_ENFORCEMENT: framework-level circuit breaker via the
        // AiPipeline BudgetCapturingSession decorator — honest because
        // the deployed-agent dispatch path emits typed TokenUsage through
        // session.usage(), the signal BudgetCapturingSession taps for
        // token / step abort. The Atmosphere-native dispatch path does
        // not surface usage (Embabel limitation), so token / step limits
        // are effective only on the deployed-agent path; wall-clock
        // limits trip on either path.
        AiCapability.BUDGET_ENFORCEMENT,
        // CONFIDENCE_SCORES: framework-level — AiPipeline's
        // ConfidenceCapturingSession parses the model-reported
        // confidence field on stream completion. Honest on both
        // dispatch paths (deployed-agent and Atmosphere-native) because
        // both honor SYSTEM_PROMPT and stream response text.
        AiCapability.CONFIDENCE_SCORES,
        // PASSIVATION: AgentPassivation snapshots context.history() into
        // a CheckpointStore. Honest because the Atmosphere-native dispatch
        // path threads history into PromptRunner.withMessages — a resumed
        // call observes the same conversation the paused call saw.
        AiCapability.PASSIVATION
    )
}
