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

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Agents.AI.Purview.Models.Common;
using Microsoft.Agents.AI.Purview.Models.Jobs;
using Microsoft.Agents.AI.Purview.Models.Requests;
using Microsoft.Agents.AI.Purview.Models.Responses;
using Microsoft.Extensions.AI;

namespace Microsoft.Agents.AI.Purview;

/// <summary>
/// Processor class that combines protectionScopes, processContent, and contentActivities calls.
/// </summary>
internal sealed class ScopedContentProcessor : IScopedContentProcessor
{
    private readonly IPurviewClient _purviewClient;
    private readonly ICacheProvider _cacheProvider;
    private readonly IChannelHandler _channelHandler;

    /// <summary>
    /// Create a new instance of <see cref="ScopedContentProcessor"/>.
    /// </summary>
    /// <param name="purviewClient">The purview client to use for purview requests.</param>
    /// <param name="cacheProvider">The cache used to store Purview data.</param>
    /// <param name="channelHandler">The channel handler used to manage background jobs.</param>
    public ScopedContentProcessor(IPurviewClient purviewClient, ICacheProvider cacheProvider, IChannelHandler channelHandler)
    {
        this._purviewClient = purviewClient;
        this._cacheProvider = cacheProvider;
        this._channelHandler = channelHandler;
    }

    /// <inheritdoc/>
    public async Task<(bool shouldBlock, string? userId)> ProcessMessagesAsync(IEnumerable<ChatMessage> messages, string? sessionId, Activity activity, PurviewSettings purviewSettings, string? userId, CancellationToken cancellationToken)
    {
        List<ProcessContentRequest> pcRequests = await this.MapMessageToPCRequestsAsync(messages, sessionId, activity, purviewSettings, userId, cancellationToken).ConfigureAwait(false);

        bool shouldBlock = false;
        string? resolvedUserId = null;

        foreach (ProcessContentRequest pcRequest in pcRequests)
        {
            resolvedUserId = pcRequest.UserId;
            ProcessContentResponse processContentResponse = await this.ProcessContentWithProtectionScopesAsync(pcRequest, cancellationToken).ConfigureAwait(false);
            if (processContentResponse.PolicyActions?.Count > 0)
            {
                foreach (DlpActionInfo policyAction in processContentResponse.PolicyActions)
                {
                    // We need to process all data before blocking, so set the flag and return it outside of this loop.
                    if (policyAction.Action == DlpAction.BlockAccess)
                    {
                        shouldBlock = true;
                    }

                    if (policyAction.RestrictionAction == RestrictionAction.Block)
                    {
                        shouldBlock = true;
                    }
                }
            }
        }

        return (shouldBlock, resolvedUserId);
    }

    private static bool TryGetUserIdFromPayload(IEnumerable<ChatMessage> messages, out string? userId)
    {
        userId = null;

        foreach (ChatMessage message in messages)
        {
            if (message.AdditionalProperties != null &&
                message.AdditionalProperties.TryGetValue(Constants.UserId, out string? potentialUserId) &&
                Guid.TryParse(potentialUserId, out Guid _))
            {
                userId = potentialUserId;
                return true;
            }
            else if (Guid.TryParse(message.AuthorName, out Guid _))
            {
                userId = message.AuthorName;
                return true;
            }
        }

        return false;
    }

    /// <summary>
    /// Transform a list of ChatMessages into a list of ProcessContentRequests.
    /// </summary>
    /// <param name="messages">The messages to transform.</param>
    /// <param name="sessionId">The id of the message session.</param>
    /// <param name="activity">The activity performed on the content.</param>
    /// <param name="settings">The settings used for purview integration.</param>
    /// <param name="userId">The entra id of the user who made the interaction.</param>
    /// <param name="cancellationToken">The cancellation token used to cancel async operations.</param>
    /// <returns>A list of process content requests.</returns>
    private async Task<List<ProcessContentRequest>> MapMessageToPCRequestsAsync(IEnumerable<ChatMessage> messages, string? sessionId, Activity activity, PurviewSettings settings, string? userId, CancellationToken cancellationToken)
    {
        List<ProcessContentRequest> pcRequests = [];
        TokenInfo? tokenInfo = await this._purviewClient.GetUserInfoFromTokenAsync(cancellationToken, settings.TenantId).ConfigureAwait(false);
        string tenantId = tokenInfo?.TenantId ?? settings.TenantId ?? throw new PurviewRequestException("No tenant id provided or inferred for Purview request. Please provide a tenant id in PurviewSettings or configure the TokenCredential to authenticate to a tenant.");
        string? resolvedUserId = !string.IsNullOrEmpty(tokenInfo?.UserId) ? tokenInfo.UserId : userId;
        if (string.IsNullOrEmpty(resolvedUserId) && TryGetUserIdFromPayload(messages, out string? payloadUserId))
        {
            resolvedUserId = payloadUserId;
        }

        foreach (ChatMessage message in messages)
        {
            string messageId = message.MessageId ?? Guid.NewGuid().ToString();
            ContentBase content = new PurviewTextContent(message.Text);
            string correlationId = (sessionId ?? Guid.NewGuid().ToString()) + "@AF";
            ProcessConversationMetadata conversationMetadata = new(content, messageId, false, $"Agent Framework Message {messageId}", correlationId)
            {
                SequenceNumber = DateTime.UtcNow.Ticks,
            };
            ActivityMetadata activityMetadata = new(activity);
            PolicyLocation policyLocation;

            if (settings.PurviewAppLocation != null)
            {
                policyLocation = settings.PurviewAppLocation.GetPolicyLocation();
            }
            else if (tokenInfo?.ClientId != null)
            {
                policyLocation = new($"{Constants.ODataGraphNamespace}.policyLocationApplication", tokenInfo.ClientId);
            }
            else
            {
                throw new PurviewRequestException("No app location provided or inferred for Purview request. Please provide an app location in PurviewSettings or configure the TokenCredential to authenticate to an entra app.");
            }

            string appVersion = !string.IsNullOrEmpty(settings.AppVersion) ? settings.AppVersion : "Unknown";

            ProtectedAppMetadata protectedAppMetadata = new(policyLocation)
            {
                Name = settings.AppName,
                Version = appVersion
            };
            IntegratedAppMetadata integratedAppMetadata = new()
            {
                Name = settings.AppName,
                Version = appVersion
            };

            DeviceMetadata deviceMetadata = new()
            {
                OperatingSystemSpecifications = new()
                {
                    OperatingSystemPlatform = "Unknown",
                    OperatingSystemVersion = "Unknown"
                }
            };
            ContentToProcess contentToProcess = new([conversationMetadata], activityMetadata, deviceMetadata, integratedAppMetadata, protectedAppMetadata);

            if (string.IsNullOrEmpty(resolvedUserId))
            {
                throw new PurviewRequestException("No user id provided or inferred for Purview request. Please provide an Entra user id in each message, pass a user id to the processor, or configure the TokenCredential to authenticate to an Entra user.");
            }

            ProcessContentRequest pcRequest = new(contentToProcess, resolvedUserId, tenantId);
            pcRequests.Add(pcRequest);
        }

        return pcRequests;
    }

    /// <summary>
    /// Orchestrates process content and protection scopes calls.
    /// </summary>
    /// <param name="pcRequest">The process content request.</param>
    /// <param name="cancellationToken">The cancellation token used to cancel async operations.</param>
    /// <returns>A process content response. This could be a response from the process content API or a response generated from a content activities call.</returns>
    private async Task<ProcessContentResponse> ProcessContentWithProtectionScopesAsync(ProcessContentRequest pcRequest, CancellationToken cancellationToken)
    {
        ProtectionScopesRequest psRequest = CreateProtectionScopesRequest(pcRequest, pcRequest.UserId, pcRequest.TenantId, pcRequest.CorrelationId);

        PaymentRequiredCacheEntry? cachedPaymentRequired = await this._cacheProvider.GetAsync<PaymentRequiredCacheKey, PaymentRequiredCacheEntry>(
            new PaymentRequiredCacheKey(pcRequest.TenantId),
            cancellationToken).ConfigureAwait(false);

        if (cachedPaymentRequired != null)
        {
            throw new PurviewPaymentRequiredException(cachedPaymentRequired.Message ?? "Payment required");
        }

        ProtectionScopesCacheKey cacheKey = new(psRequest);

        ProtectionScopesResponse? cacheResponse = await this._cacheProvider.GetAsync<ProtectionScopesCacheKey, ProtectionScopesResponse>(cacheKey, cancellationToken).ConfigureAwait(false);

        if (cacheResponse != null)
        {
            return await this.ProcessWithCachedScopesAsync(pcRequest, cacheResponse, cacheKey, cancellationToken).ConfigureAwait(false);
        }

        try
        {
            this._channelHandler.QueueJob(new ScopeRetrievalJob(psRequest, cacheKey, pcRequest));
        }
        catch (PurviewJobException)
        {
            // QueueJob already logs failures. Scope warmup is best effort; don't block ProcessContent.
        }

        return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions: null, cancellationToken).ConfigureAwait(false);
    }

    /// <summary>
    /// Apply locally-cached protection scopes to the request and dispatch ProcessContent appropriately.
    /// </summary>
    private async Task<ProcessContentResponse> ProcessWithCachedScopesAsync(
        ProcessContentRequest pcRequest,
        ProtectionScopesResponse psResponse,
        ProtectionScopesCacheKey cacheKey,
        CancellationToken cancellationToken)
    {
        pcRequest.ScopeIdentifier = psResponse.ScopeIdentifier;

        (bool shouldProcess, List<DlpActionInfo> dlpActions, ExecutionMode executionMode) = CheckApplicableScopes(pcRequest, psResponse);

        if (shouldProcess)
        {
            pcRequest.ProcessInline = executionMode == ExecutionMode.EvaluateInline;

            if (executionMode == ExecutionMode.EvaluateOffline)
            {
                this._channelHandler.QueueJob(new ProcessContentJob(pcRequest));
                return new ProcessContentResponse();
            }

            return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions, cancellationToken).ConfigureAwait(false);
        }

        ContentActivitiesRequest caRequest = new(pcRequest.UserId, pcRequest.TenantId, pcRequest.ContentToProcess, pcRequest.CorrelationId);
        this._channelHandler.QueueJob(new ContentActivityJob(caRequest));

        return new ProcessContentResponse();
    }

    /// <summary>
    /// Call ProcessContent and invalidate the protection scopes cache when the response indicates the cached scopes are stale.
    /// </summary>
    private async Task<ProcessContentResponse> CallProcessContentAsync(
        ProcessContentRequest pcRequest,
        ProtectionScopesCacheKey cacheKey,
        List<DlpActionInfo>? dlpActions,
        CancellationToken cancellationToken)
    {
        ProcessContentResponse pcResponse = await this._purviewClient.ProcessContentAsync(pcRequest, cancellationToken).ConfigureAwait(false);

        if (pcRequest.ScopeIdentifier != null && pcResponse.ProtectionScopeState == ProtectionScopeState.Modified)
        {
            await this._cacheProvider.RemoveAsync(cacheKey, cancellationToken).ConfigureAwait(false);
        }

        if (dlpActions?.Count > 0)
        {
            pcResponse = CombinePolicyActions(pcResponse, dlpActions);
        }

        return pcResponse;
    }

    /// <summary>
    /// Dedupe policy actions received from the service.
    /// </summary>
    /// <param name="pcResponse">The process content response which may contain DLP actions.</param>
    /// <param name="actionInfos">DLP actions returned from protection scopes.</param>
    /// <returns>The process content response with the protection scopes DLP actions added.</returns>
    private static ProcessContentResponse CombinePolicyActions(ProcessContentResponse pcResponse, List<DlpActionInfo>? actionInfos)
    {
        if (actionInfos?.Count > 0)
        {
            List<DlpActionInfo> combinedActions = [];
            HashSet<(DlpAction Action, RestrictionAction? RestrictionAction)> seenActions = [];
            IEnumerable<DlpActionInfo> allActions = pcResponse.PolicyActions is null
                ? actionInfos
                : pcResponse.PolicyActions.Concat(actionInfos);

            foreach (DlpActionInfo actionInfo in allActions)
            {
                if (seenActions.Add((actionInfo.Action, actionInfo.RestrictionAction)))
                {
                    combinedActions.Add(actionInfo);
                }
            }

            pcResponse.PolicyActions = combinedActions;
        }

        return pcResponse;
    }

    /// <summary>
    /// Check if any scopes are applicable to the request.
    /// </summary>
    /// <param name="pcRequest">The process content request.</param>
    /// <param name="psResponse">The protection scopes response that was returned for the process content request.</param>
    /// <returns>A bool indicating if the content needs to be processed. A list of applicable actions from the scopes response, and the execution mode for the process content request.</returns>
    internal static (bool shouldProcess, List<DlpActionInfo> dlpActions, ExecutionMode executionMode) CheckApplicableScopes(ProcessContentRequest pcRequest, ProtectionScopesResponse psResponse)
    {
        ProtectionScopeActivities requestActivity = TranslateActivity(pcRequest.ContentToProcess.ActivityMetadata.Activity);

        // The location data type is formatted as microsoft.graph.{locationType}
        // Sometimes a '#' gets appended by graph during responses, so for the sake of simplicity,
        // Split it by '.' and take the last segment. We'll do a case-insensitive endsWith later.
        string[] locationSegments = pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation.DataType.Split('.');
        string locationType = locationSegments.Length > 0 ? locationSegments[locationSegments.Length - 1] : pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation.Value;

        string locationValue = pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation.Value;
        List<DlpActionInfo> dlpActions = [];
        bool shouldProcess = false;
        ExecutionMode executionMode = ExecutionMode.EvaluateOffline;

        foreach (var scope in psResponse.Scopes ?? Array.Empty<PolicyScopeBase>())
        {
            bool activityMatch = scope.Activities.HasFlag(requestActivity);
            bool locationMatch = false;

            foreach (var location in scope.Locations ?? Array.Empty<PolicyLocation>())
            {
                if (location.DataType.EndsWith(locationType, StringComparison.OrdinalIgnoreCase) && location.Value.Equals(locationValue, StringComparison.OrdinalIgnoreCase))
                {
                    locationMatch = true;
                    break;
                }
            }

            if (activityMatch && locationMatch)
            {
                shouldProcess = true;

                if (scope.ExecutionMode == ExecutionMode.EvaluateInline)
                {
                    executionMode = ExecutionMode.EvaluateInline;
                }

                if (scope.PolicyActions != null)
                {
                    dlpActions.AddRange(scope.PolicyActions);
                }
            }
        }

        return (shouldProcess, dlpActions, executionMode);
    }

    /// <summary>
    /// Create a ProtectionScopesRequest for the given content ProcessContentRequest.
    /// </summary>
    /// <param name="pcRequest">The process content request.</param>
    /// <param name="userId">The entra user id of the user who sent the data.</param>
    /// <param name="tenantId">The tenant id of the user who sent the data.</param>
    /// <param name="correlationId">The correlation id of the request.</param>
    /// <returns>The protection scopes request generated from the process content request.</returns>
    private static ProtectionScopesRequest CreateProtectionScopesRequest(ProcessContentRequest pcRequest, string userId, string tenantId, Guid correlationId)
    {
        return new ProtectionScopesRequest(userId, tenantId)
        {
            Activities = TranslateActivity(pcRequest.ContentToProcess.ActivityMetadata.Activity),
            Locations = [pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation],
            DeviceMetadata = pcRequest.ContentToProcess.DeviceMetadata,
            IntegratedAppMetadata = pcRequest.ContentToProcess.IntegratedAppMetadata,
            CorrelationId = correlationId
        };
    }

    /// <summary>
    /// Map process content activity to protection scope activity.
    /// </summary>
    /// <param name="activity">The process content activity.</param>
    /// <returns>The protection scopes activity.</returns>
    private static ProtectionScopeActivities TranslateActivity(Activity activity)
    {
        return activity switch
        {
            Activity.Unknown => ProtectionScopeActivities.None,
            Activity.UploadText => ProtectionScopeActivities.UploadText,
            Activity.UploadFile => ProtectionScopeActivities.UploadFile,
            Activity.DownloadText => ProtectionScopeActivities.DownloadText,
            Activity.DownloadFile => ProtectionScopeActivities.DownloadFile,
            _ => ProtectionScopeActivities.UnknownFutureValue,
        };
    }
}
