using System.Collections.Generic;

using System.Diagnostics.CodeAnalysis;

using System.Diagnostics.Contracts;

using Rxx;



namespace System.Net.Sockets

{

	public static partial class SocketExtensions

	{

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginConnect(this Socket socket, EndPoint remoteEP, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				RemoteEndPoint = remoteEP

			};



			return BeginInvoke(callback, state, args, socket.ConnectAsync);

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginConnect(this Socket socket, IPAddress address, int port, AsyncCallback requestCallback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				RemoteEndPoint = new IPEndPoint(address, port)

			};



			return BeginInvoke(requestCallback, state, args, socket.ConnectAsync);

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginConnect(this Socket socket, string host, int port, AsyncCallback requestCallback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				RemoteEndPoint = new DnsEndPoint(host, port)

			};



			return BeginInvoke(requestCallback, state, args, socket.ConnectAsync);

		}



		internal static void EndConnect(this Socket socket, IAsyncResult asyncResult)

		{

			Contract.Requires(asyncResult != null);



			using (var args = EndInvoke(asyncResult))

			{

				// do nothing

			}

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginReceive(this Socket socket, IList<ArraySegment<byte>> buffers, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				BufferList = buffers

			};



			return BeginInvoke(callback, state, args, socket.ReceiveAsync);

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginReceive(this Socket socket, IList<ArraySegment<byte>> buffers, out SocketError errorCode, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				BufferList = buffers

			};



			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);



			if (result.CompletedSynchronously)

			{

				errorCode = args.SocketError;



				// Prevent the error from being thrown in EndInvoke

				result.OutOfBandData = new object();

			}

			else

			{

				errorCode = default(SocketError);

			}



			return result;

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginReceive(this Socket socket, byte[] buffer, int offset, int size, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs();



			args.SetBuffer(buffer, offset, size);



			return BeginInvoke(callback, state, args, socket.ReceiveAsync);

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginReceive(this Socket socket, byte[] buffer, int offset, int size, out SocketError errorCode, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs();



			args.SetBuffer(buffer, offset, size);



			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);



			if (result.CompletedSynchronously)

			{

				errorCode = args.SocketError;



				// Prevent the error from being thrown in EndInvoke

				result.OutOfBandData = new object();

			}

			else

			{

				errorCode = default(SocketError);

			}



			return result;

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginReceiveFrom(this Socket socket, byte[] buffer, int offset, int size, ref EndPoint remoteEP, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				RemoteEndPoint = remoteEP

			};



			args.SetBuffer(buffer, offset, size);



			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);



			if (result.CompletedSynchronously)

			{

				remoteEP = args.RemoteEndPoint;

			}



			return result;

		}



		internal static int EndReceive(this Socket socket, IAsyncResult asyncResult)

		{

			Contract.Requires(asyncResult != null);



			using (var args = EndInvoke(asyncResult))

			{

				return args.BytesTransferred;

			}

		}



		[SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket",

			Justification = "Needed for parity with .NET 4.0 Framework.")]

		internal static int EndReceive(this Socket socket, IAsyncResult asyncResult, out SocketError errorCode)

		{

			Contract.Requires(asyncResult != null);



			using (var args = EndInvoke(asyncResult))

			{

				errorCode = args.SocketError;



				return args.BytesTransferred;

			}

		}



		[SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket",

			Justification = "Needed for parity with .NET 4.0 Framework.")]

		internal static int EndReceiveFrom(this Socket socket, IAsyncResult asyncResult, ref EndPoint endPoint)

		{

			Contract.Requires(asyncResult != null);



			using (var args = EndInvoke(asyncResult))

			{

				endPoint = args.RemoteEndPoint;



				return args.BytesTransferred;

			}

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginSend(this Socket socket, IList<ArraySegment<byte>> buffers, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				BufferList = buffers

			};



			return BeginInvoke(callback, state, args, socket.SendAsync);

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginSend(this Socket socket, IList<ArraySegment<byte>> buffers, out SocketError errorCode, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				BufferList = buffers

			};



			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);



			if (result.CompletedSynchronously)

			{

				errorCode = args.SocketError;



				// Prevent the error from being thrown in EndInvoke

				result.OutOfBandData = new object();

			}

			else

			{

				errorCode = default(SocketError);

			}



			return result;

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginSend(this Socket socket, byte[] buffer, int offset, int size, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs();



			args.SetBuffer(buffer, offset, size);



			return BeginInvoke(callback, state, args, socket.SendAsync);

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginSend(this Socket socket, byte[] buffer, int offset, int size, out SocketError errorCode, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs();



			args.SetBuffer(buffer, offset, size);



			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);



			if (result.CompletedSynchronously)

			{

				errorCode = args.SocketError;



				// Prevent the error from being thrown in EndInvoke

				result.OutOfBandData = new object();

			}

			else

			{

				errorCode = default(SocketError);

			}



			return result;

		}



		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",

			Justification = "The args object is disposed in the EndInvoke method.")]

		internal static IAsyncResult BeginSendTo(this Socket socket, byte[] buffer, int offset, int size, EndPoint remoteEP, AsyncCallback callback, object state)

		{

			Contract.Requires(socket != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var args = new SocketAsyncEventArgs()

			{

				RemoteEndPoint = remoteEP

			};



			args.SetBuffer(buffer, offset, size);



			return BeginInvoke(callback, state, args, socket.SendAsync);

		}



		internal static int EndSend(this Socket socket, IAsyncResult asyncResult)

		{

			Contract.Requires(asyncResult != null);



			using (var args = EndInvoke(asyncResult))

			{

				return args.BytesTransferred;

			}

		}



		[SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket",

			Justification = "Needed for parity with .NET 4.0 Framework.")]

		internal static int EndSend(this Socket socket, IAsyncResult asyncResult, out SocketError errorCode)

		{

			Contract.Requires(asyncResult != null);



			using (var args = EndInvoke(asyncResult))

			{

				errorCode = args.SocketError;



				return args.BytesTransferred;

			}

		}



		[System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket", 

			Justification = "Must have the same signature as in .NET 4.0.")]

		internal static int EndSendTo(this Socket socket, IAsyncResult asyncResult)

		{

			Contract.Requires(asyncResult != null);



			using (var args = EndInvoke(asyncResult))

			{

				return args.BytesTransferred;

			}

		}



		private static AsyncResult<SocketAsyncEventArgs> BeginInvoke(AsyncCallback callback, object state, SocketAsyncEventArgs args, Func<SocketAsyncEventArgs, bool> target)

		{

			Contract.Requires(args != null);

			Contract.Requires(target != null);

			Contract.Ensures(Contract.Result<IAsyncResult>() != null);



			var result = new AsyncResult<SocketAsyncEventArgs>(state);



			EventHandler<SocketAsyncEventArgs> completed = (sender, e) =>

			{

				if (!result.IsCompleted)

				{

					result.Complete(args, synchronous: false);

				}



				if (callback != null)

				{

					callback(result);

				}

			};



			args.Completed += completed;



			if (!target(args))

			{

				result.Complete(args, synchronous: true);



				completed(null, null);

			}



			return result;

		}



		private static SocketAsyncEventArgs EndInvoke(IAsyncResult asyncResult)

		{

			Contract.Requires(asyncResult != null);

			Contract.Ensures(Contract.Result<SocketAsyncEventArgs>() != null);



			var result = (AsyncResult<SocketAsyncEventArgs>) asyncResult;



			using (result)

			{

				result.Wait();



				var args = result.Data;



				Contract.Assume(args != null);



				if (result.OutOfBandData == null)

				{

					if (args.ConnectByNameError != null)

					{

						throw args.ConnectByNameError;

					}



					if (args.SocketError != SocketError.Success)

					{

						throw new SocketException((int) args.SocketError);

					}

				}



				return args;

			}

		}

	}

}