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

namespace System.ServiceModel.Diagnostics
{
    using System.Runtime;
    using System.Diagnostics;

    abstract class TraceAsyncResult : AsyncResult
    {
        static Action<AsyncCallback, IAsyncResult> waitResultCallback = new Action<AsyncCallback, IAsyncResult>(DoCallback);

        protected TraceAsyncResult(AsyncCallback callback, object state) :
            base(callback, state)
        {
            if (TraceUtility.MessageFlowTracingOnly)
            {
                this.CallbackActivity = ServiceModelActivity.CreateLightWeightAsyncActivity(Trace.CorrelationManager.ActivityId);
                base.VirtualCallback = waitResultCallback;
            } 
            else if (DiagnosticUtility.ShouldUseActivity)
            {
                this.CallbackActivity = ServiceModelActivity.Current;
                if (this.CallbackActivity != null)
                {
                    base.VirtualCallback = waitResultCallback;
                }
            }
        }

        public ServiceModelActivity CallbackActivity
        {
            get;
            private set;
        }

        static void DoCallback(AsyncCallback callback, IAsyncResult result)
        {
            if (result is TraceAsyncResult)
            {
                TraceAsyncResult thisPtr = result as TraceAsyncResult;
                Fx.Assert(thisPtr.CallbackActivity != null, "this shouldn't be hooked up if we don't have a CallbackActivity");

                if (TraceUtility.MessageFlowTracingOnly)
                {
                    Trace.CorrelationManager.ActivityId = thisPtr.CallbackActivity.Id;
                    thisPtr.CallbackActivity = null;
                }

                using (ServiceModelActivity.BoundOperation(thisPtr.CallbackActivity))
                {
                    callback(result);
                }
            }
        }
    }
}