//----------------------------------------------------------------
// Copyright (c) Microsoft Corporation.  All rights reserved.
//----------------------------------------------------------------
namespace System.Activities.Statements
{
    using System;
    using System.Collections.Generic;
    using System.Runtime;
    using System.Threading;
    using System.Collections.ObjectModel;

    sealed class WorkflowCompensationBehavior : NativeActivity
    {
        Variable<CompensationToken> currentCompensationToken;

        public WorkflowCompensationBehavior()
            : base()
        {
            currentCompensationToken = new Variable<CompensationToken>
                {
                    Name = "currentCompensationToken",
                };

            DefaultCompensation = new DefaultCompensation()
                {
                    Target = new InArgument<CompensationToken>(this.currentCompensationToken),
                };

            DefaultConfirmation = new DefaultConfirmation()
                {
                    Target = new InArgument<CompensationToken>(this.currentCompensationToken),
                };
        }

        Activity DefaultCompensation
        {
            get;
            set;
        }

        Activity DefaultConfirmation
        {
            get;
            set;
        }

        protected override bool CanInduceIdle
        {
            get
            {
                return true;
            }
        }

        protected override void CacheMetadata(NativeActivityMetadata metadata)
        {
            Fx.Assert(this.DefaultCompensation != null, "DefaultCompensation must be valid");
            Fx.Assert(this.DefaultConfirmation != null, "DefaultConfirmation must be valid");
            metadata.SetImplementationChildrenCollection(
                new Collection<Activity>
                {
                    this.DefaultCompensation, 
                    this.DefaultConfirmation
                });

            metadata.SetImplementationVariablesCollection(new Collection<Variable> { this.currentCompensationToken });
        }

        protected override void Execute(NativeActivityContext context)
        {
            Bookmark mainRootCompleteBookmark = context.CreateBookmark(OnMainRootComplete, BookmarkOptions.NonBlocking);
            context.RegisterMainRootCompleteCallback(mainRootCompleteBookmark);

            CompensationExtension compensationExtension = context.GetExtension<CompensationExtension>();
            Fx.Assert(compensationExtension != null, "CompensationExtension must be valid");

            compensationExtension.WorkflowCompensation = context.CreateBookmark(new BookmarkCallback(OnCompensate));
            compensationExtension.WorkflowConfirmation = context.CreateBookmark(new BookmarkCallback(OnConfirm));

            Fx.Assert(compensationExtension.WorkflowCompensationScheduled != null, "compensationExtension.WorkflowCompensationScheduled bookmark must be setup by now");
            context.ResumeBookmark(compensationExtension.WorkflowCompensationScheduled, null);
        }

        protected override void Cancel(NativeActivityContext context)
        {
            context.CancelChildren();
        }

        void OnMainRootComplete(NativeActivityContext context, Bookmark bookmark, object value)
        {
            CompensationExtension compensationExtension = context.GetExtension<CompensationExtension>();
            Fx.Assert(compensationExtension != null, "CompensationExtension must be valid");

            CompensationTokenData rootHandle = compensationExtension.Get(CompensationToken.RootCompensationId);
            Fx.Assert(rootHandle != null, "rootToken must be valid");

            ActivityInstanceState completionState = (ActivityInstanceState)value;

            if (completionState == ActivityInstanceState.Closed)
            {
                context.ResumeBookmark(compensationExtension.WorkflowConfirmation, new CompensationToken(rootHandle));
            }
            else if (completionState == ActivityInstanceState.Canceled)
            {
                context.ResumeBookmark(compensationExtension.WorkflowCompensation, new CompensationToken(rootHandle));
            }
            else if (completionState == ActivityInstanceState.Faulted)
            {
                // Do nothing. Neither Compensate nor Confirm.
                // Remove the bookmark to complete the WorkflowCompensationBehavior execution. 
               context.RemoveBookmark(compensationExtension.WorkflowConfirmation); 
               context.RemoveBookmark(compensationExtension.WorkflowCompensation);
            }
        }

        void OnCompensate(NativeActivityContext context, Bookmark bookmark, object value)
        {
            CompensationExtension compensationExtension = context.GetExtension<CompensationExtension>();
            Fx.Assert(compensationExtension != null, "CompensationExtension must be valid");

            CompensationToken rootToken = (CompensationToken)value;
            Fx.Assert(rootToken != null, "rootToken must be passed");

            this.currentCompensationToken.Set(context, rootToken);

            CompensationTokenData rootTokenData = compensationExtension.Get(rootToken.CompensationId);
            if (rootTokenData.ExecutionTracker.Count > 0)
            {
                context.ScheduleActivity(DefaultCompensation, new CompletionCallback(OnCompensationComplete));
            }
            else
            {
                OnCompensationComplete(context, null);
            }
        }

        void OnCompensationComplete(NativeActivityContext context, ActivityInstance completedInstance)
        {
            // Remove bookmark.... have a cleanup book mark method...
            CompensationExtension compensationExtension = context.GetExtension<CompensationExtension>();
            Fx.Assert(compensationExtension != null, "CompensationExtension must be valid");

            context.RemoveBookmark(compensationExtension.WorkflowConfirmation);
        }

        void OnConfirm(NativeActivityContext context, Bookmark bookmark, object value)
        {
            CompensationExtension compensationExtension = context.GetExtension<CompensationExtension>();
            Fx.Assert(compensationExtension != null, "CompensationExtension must be valid");

            CompensationToken rootToken = (CompensationToken)value;
            Fx.Assert(rootToken != null, "rootToken must be passed");

            this.currentCompensationToken.Set(context, rootToken);

            CompensationTokenData rootTokenData = compensationExtension.Get(rootToken.CompensationId);
            if (rootTokenData.ExecutionTracker.Count > 0)
            {
                context.ScheduleActivity(DefaultConfirmation, new CompletionCallback(OnConfirmationComplete));
            }
            else
            {
                OnConfirmationComplete(context, null);
            }
        }

        void OnConfirmationComplete(NativeActivityContext context, ActivityInstance completedInstance)
        {
            CompensationExtension compensationExtension = context.GetExtension<CompensationExtension>();
            Fx.Assert(compensationExtension != null, "CompensationExtension must be valid");

            context.RemoveBookmark(compensationExtension.WorkflowCompensation);
        }
    }
}