// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Agents.AI.Workflows.Declarative.Extensions;
using Microsoft.Agents.AI.Workflows.Declarative.Interpreter;
using Microsoft.Agents.AI.Workflows.Declarative.PowerFx;
using Microsoft.Bot.ObjectModel;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.Configuration;

namespace Microsoft.Agents.AI.Workflows.Declarative.Kit;

/// <summary>
/// Base class for an entry-point workflow executor that receives the initial trigger message.
/// </summary>
/// <typeparam name="TInput">The type of the initial message that starts the workflow.</typeparam>
public abstract class RootExecutor<TInput> : Executor<TInput>, IResettableExecutor where TInput : notnull
{
    private readonly IConfiguration? _configuration;
    private readonly WorkflowAgentProvider _agentProvider;
    private readonly WorkflowFormulaState _state;
    private readonly Func<TInput, ChatMessage>? _inputTransform;

    private string? _conversationId;

    /// <summary>
    /// Get the shared formula session to provide to workflow <see cref="ActionExecutor"/> instances.
    /// </summary>
    public FormulaSession Session { get; }

    /// <summary>
    /// Initializes a new instance of the <see cref="RootExecutor{TInput}"/> class.
    /// </summary>
    /// <param name="id">An optional identifier. If omitted, an identifier is generated by the base class.</param>
    /// <param name="options">Configuration options for workflow execution.</param>
    /// <param name="inputTransform">An optional function to transform the input message into a <see cref="ChatMessage"/>.</param>
    protected RootExecutor(string id, DeclarativeWorkflowOptions options, Func<TInput, ChatMessage>? inputTransform)
        : base(id)
    {
        this._configuration = options.Configuration;
        this._agentProvider = options.AgentProvider;
        this._conversationId = options.ConversationId;
        this._inputTransform = inputTransform;
        this._state = new WorkflowFormulaState(options.CreateRecalcEngine());
        this._state.InitializeSystem();
        this.Session = new RootFormulaSession(this._state);
    }

    /// <inheritdoc/>
    public ValueTask ResetAsync()
    {
        return default;
    }

    /// <inheritdoc/>
    public override async ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken)
    {
        DeclarativeWorkflowContext declarativeContext = new(context, this._state);
        await this.ExecuteAsync(message, declarativeContext, cancellationToken).ConfigureAwait(false);

        ChatMessage input = (this._inputTransform ?? DefaultInputTransform).Invoke(message);

        if (string.IsNullOrWhiteSpace(this._conversationId))
        {
            this._conversationId = await this._agentProvider.CreateConversationAsync(cancellationToken).ConfigureAwait(false);
        }
        await declarativeContext.QueueConversationUpdateAsync(this._conversationId, isExternal: true, cancellationToken).ConfigureAwait(false);

        await this._agentProvider.CreateMessageAsync(this._conversationId, input, cancellationToken).ConfigureAwait(false);
        await declarativeContext.SetLastMessageAsync(input).ConfigureAwait(false);

        await declarativeContext.SendMessageAsync(new ActionExecutorResult(this.Id), cancellationToken: cancellationToken).ConfigureAwait(false);
    }

    /// <summary>
    /// Executes the core logic of the root workflow for the provided initial message.
    /// </summary>
    /// <param name="message">The initial input message that triggered workflow execution.</param>
    /// <param name="context">The workflow execution context providing messaging and state services.</param>
    /// <param name="cancellationToken">A token that propagates notification when operation should be canceled.</param>
    /// <returns>A <see cref="ValueTask"/> representing the asynchronous execution operation.</returns>
    protected abstract ValueTask ExecuteAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default);

    /// <summary>
    /// Initializes the specified variables from <see cref="IConfiguration"/> if available;
    /// otherwise falls back to the process environment variables.
    /// </summary>
    /// <param name="context">The workflow execution context providing messaging and state services.</param>
    /// <param name="variableNames">The set of variable names to initialize.</param>
    /// <returns>A <see cref="ValueTask"/> representing the asynchronous execution operation.</returns>
    protected async ValueTask InitializeEnvironmentAsync(IWorkflowContext context, params string[] variableNames)
    {
        foreach (string variableName in variableNames)
        {
            await context.QueueStateUpdateAsync(variableName, GetEnvironmentVariable(variableName), VariableScopeNames.Environment).ConfigureAwait(false);
        }

        string GetEnvironmentVariable(string name)
        {
            if (this._configuration is not null)
            {
                return this._configuration[name] ?? string.Empty;
            }

            return Environment.GetEnvironmentVariable(name) ?? string.Empty;
        }
    }

    /// <summary>
    /// Transforms the input message into a <see cref="ChatMessage"/>.
    /// </summary>
    /// <param name="message">The original input object.</param>
    /// <returns>A <see cref="ChatMessage"/> derived from the input.</returns>
    protected internal static ChatMessage DefaultInputTransform(TInput message) =>
        message switch
        {
            ChatMessage chatMessage => chatMessage,
            string stringMessage => new ChatMessage(ChatRole.User, stringMessage),
            _ => new(ChatRole.User, $"{message}")
        };

    private sealed class RootFormulaSession : FormulaSession
    {
        internal RootFormulaSession(WorkflowFormulaState state)
        {
            this.State = state;
        }

        internal override WorkflowFormulaState State { get; }
    }
}