namespace System.Data.Services.Client.Async
{
    using System;
    using System.Collections.Generic;
    using System.Data.Services.Client;
    using System.Threading.Tasks;

    public static class DataServiceContextAsyncExtensions
    {
        public static async Task<IEnumerable<TResult>> ExecuteAsync<TResult>(this DataServiceQuery<TResult> query)
        {
            var queryTask = Task.Factory.FromAsync<IEnumerable<TResult>>(query.BeginExecute(null, null),
                (queryAsyncResult) =>
                {
                    var results = query.EndExecute(queryAsyncResult);
                    return results;
                });

            return await queryTask;
        }

        public static async Task<IEnumerable<TResult>> ExecuteAsync<TResult>(this DataServiceContext context, Uri requestUri)
        {
            var queryTask = Task.Factory.FromAsync<IEnumerable<TResult>>(context.BeginExecute<TResult>(requestUri, null, null),
                   (queryAsyncResult) =>
                   {
                       var results = context.EndExecute<TResult>(queryAsyncResult);
                       return results;
                   });

            return await queryTask;
        }

        public static async Task<IEnumerable<TResult>> ExecuteAsync<TResult>(this DataServiceContext context, DataServiceQueryContinuation<TResult> queryContinuationToken)
        {
            var queryTask = Task.Factory.FromAsync<IEnumerable<TResult>>(context.BeginExecute<TResult>(queryContinuationToken, null, null),
                   (queryAsyncResult) =>
                   {
                       var results = context.EndExecute<TResult>(queryAsyncResult);
                       return results;
                   });

            return await queryTask;
        }

        public static async Task<IEnumerable<TResult>> LoadPropertyAsync<TResult>(this DataServiceContext context, object entity, string propertyName)
        {
            var queryTask = Task.Factory.FromAsync<IEnumerable<TResult>>(context.BeginLoadProperty(entity, propertyName, null, null),
                   (loadPropertyAsyncResult) =>
                   {
                       var results = context.EndLoadProperty(loadPropertyAsyncResult);
                       return (IEnumerable<TResult>)results;
                   });
                                                                                        
            return await queryTask;
        }

        public static async Task<IEnumerable<TResult>> LoadPropertyAsync<TResult>(this DataServiceContext context, object entity, string propertyName, DataServiceQueryContinuation continuation)
        {
            var queryTask = Task.Factory.FromAsync<IEnumerable<TResult>>(context.BeginLoadProperty(entity, propertyName, continuation, null, null),
                  (loadPropertyAsyncResult) =>
                  {
                      var results = context.EndLoadProperty(loadPropertyAsyncResult);
                      return (IEnumerable<TResult>)results;
                  });

            return await queryTask;
        }

        public static async Task<IEnumerable<TResult>> LoadPropertyAsync<TResult>(this DataServiceContext context, object entity, string propertyName, Uri nextLinkUri)
        {
            var queryTask = Task.Factory.FromAsync<IEnumerable<TResult>>(context.BeginLoadProperty(entity, propertyName, nextLinkUri, null, null),
                    (loadPropertyAsyncResult) =>
                    {
                        var results = context.EndLoadProperty(loadPropertyAsyncResult);
                        return (IEnumerable<TResult>)results;
                    });

            return await queryTask;
        }

        public static async Task<DataServiceResponse> ExecuteBatchAsync(this DataServiceContext context, params DataServiceRequest[] requests)
        {
            var queryTask = Task.Factory.FromAsync<DataServiceResponse>(context.BeginExecuteBatch(null, null, requests),
                   (queryAsyncResult) =>
                   {
                       var results = context.EndExecuteBatch(queryAsyncResult);
                       return results;
                   });

            return await queryTask;
        }

        public static async Task<DataServiceResponse> SaveChangesAsync(this DataServiceContext context)
        {
            return await SaveChangesAsync(context, SaveChangesOptions.None);
        }

        public static async Task<DataServiceResponse> SaveChangesAsync(this DataServiceContext context, SaveChangesOptions options)
        {
            var queryTask = Task.Factory.FromAsync<DataServiceResponse>(context.BeginSaveChanges(options, null, null),
                   (queryAsyncResult) =>
                   {
                       var results = context.EndSaveChanges(queryAsyncResult);
                       return results;
                   });

            return await queryTask;
        }
    }
}