namespace Owin
{
    using System;
    using System.Collections.Generic;
    using System.Text;
    using System.IO;

    public interface IResponseHandler
    {
        Type TypeToHandle { get; }
        void WriteToStream(object value, Stream stream);
    }

    public interface IResponseHandler<T> : IResponseHandler
    {
        void WriteToStream(T value, Stream stream);
    }

    public static class OWIN
    {
        static IDictionary<Type, Action<object, Stream>> responseHandlers = new Dictionary<Type, Action<object, Stream>>();

        public static void AddDefaultHandlers()
        {
            AddRequestHandler<byte[]>((x, s) => s.Write(x, 0, x.Length));

            AddRequestHandler<string>((x, s) =>
            {
                byte[] buffer = Encoding.UTF8.GetBytes(x);
                s.Write(buffer, 0, buffer.Length);
            });

            AddRequestHandler<ArraySegment<byte>>((x, s) => s.Write(x.Array, 0, xas.Array.Length));
        }

        public static void AddResponseHandler<T>(Action<T, Stream> handler)
        {
            responseHandlers[typeof(T)] = handler;
        }

        public static void AddResponseHandler<T>(IResponseHandler<T> handler)
        {
            responseHandlers[handler.TypeToHandle] = (x, s) => handler.WriteToStream(x, s);
        }

        public static void AddResponseHandler(IResponseHandler handler)
        {
            responseHandlers[handler.TypeToHandle] = (x, s) => handler.WriteToStream(x, s);
        }

        public static Action<T, Stream> GetResponseHandler<T>()
        {
            Type objectType = typeof(T);
            Action<T, Stream> handler = null;
            if (!responseHandlers.TryGetValue(objectType, out handler))
            {
                if (objectType.IsGenericType)
                    objectType = objectType.GetGenericTypeDefinition();
                responseHandlers.TryGetValue(objectType, out handler);
            }

            Action<T, Stream> action = handler as Action<T, Stream>;
            return action != null ? action : (x, s) => { };
        }


        public static void WriteBodyToStream(this IResponse response, Stream stream)
        {
            var body = response.GetBody();
            using (var enumerator = body.GetEnumerator())
            {
                while (enumerator.MoveNext())
                {
                    var value = enumerator.Current;
                    if (value != null)
                    {
                        OWIN.GetResponseHandler(value.GetType())(value, stream);
                    }
                }
            }
        }
    }
}