package de.openknowledge.jaxrs.reactive;

import javax.servlet.AsyncContext;
import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.MultivaluedMap;
import javax.ws.rs.ext.MessageBodyWriter;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.reflect.Type;
import java.util.concurrent.Flow;
import de.openknowledge.reactive.commons.AbstractSubscriber;

/**
 * @author Christian Schulz - open knowledge GmbH
 */
public class AsyncContextMessageBodyWriterSubscriber extends AbstractSubscriber<Object> {

  private boolean first = true;
  private AsyncContext asyncContext;
  private final MessageBodyWriter entityWriter;
  private final Class<?> targetClass;
  private final Type targetType;
  private final Annotation[] annotations;
  private final MediaType mediaType;
  private final MultivaluedMap<String, Object> httpHeaders;

  public AsyncContextMessageBodyWriterSubscriber(AsyncContext asyncContext, MessageBodyWriter entityWriter, Class<?> targetClass,
                                                 Type targetType, Annotation[] annotations, MediaType mediaType,
                                                 MultivaluedMap<String, Object> httpHeaders) {
    this.asyncContext = asyncContext;
    this.entityWriter = entityWriter;
    this.targetClass = targetClass;
    this.targetType = targetType;
    this.annotations = annotations;
    this.mediaType = mediaType;
    this.httpHeaders = httpHeaders;
  }

  @Override
  public void onSubscribe(Flow.Subscription subscription) {
    super.onSubscribe(subscription);
    try {
      asyncContext.getResponse().getOutputStream().setWriteListener(new WriteListener() {

        @Override
        public void onWritePossible() throws IOException {
          request(Long.MAX_VALUE);
        }

        @Override
        public void onError(Throwable error) {
          AsyncContextMessageBodyWriterSubscriber.this.onError(error);
        }
      });
    } catch (IOException e) {
      onError(e);
    }
  }

  @Override
  public void onNext(Object item) {
    try {
      ServletOutputStream outputStream = asyncContext.getResponse().getOutputStream();
      if (first) {
        outputStream.print('[');
        first = false;
      } else {
        outputStream.print(',');
      }
      entityWriter.writeTo(item, targetClass, targetType, annotations, mediaType, httpHeaders, new NonClosableOutputStream(outputStream));
      outputStream.flush();

    } catch (IOException e) {
      onError(e);
    }
  }

  @Override
  public void onComplete() {
    try {
      ServletOutputStream outputStream = asyncContext.getResponse().getOutputStream();
      if (first) {
        outputStream.print('[');
      }
      outputStream.println(']');
      outputStream.flush();
      asyncContext.complete();
    } catch (IOException e) {
      // we are finished, ignoring errors
    }
  }
}

@@ -935,16 +935,75 @@
 Headers;
+%0A  private NonClosableOutputStream nonClosableOutputStream;
 %0A%0A  publ
@@ -1709,24 +1709,59 @@
  try %7B%0A     
+ ServletOutputStream outputStream =
  asyncContex
@@ -1793,16 +1793,36 @@
 Stream()
+;%0A      outputStream
 .setWrit
@@ -2133,16 +2133,91 @@
     %7D);%0A
+      nonClosableOutputStream = new NonClosableOutputStream(outputStream);%0A
     %7D ca
@@ -2639,20 +2639,16 @@
 aders, n
-ew N
 onClosab
@@ -2661,30 +2661,16 @@
 utStream
-(outputStream)
 );%0A     
@@ -3007,24 +3007,24 @@
 am.flush();%0A
-
       asyncC
@@ -3038,24 +3038,68 @@
 complete();%0A
+      nonClosableOutputStream.closeAsync();%0A
     %7D catch 

create stream once and close it