import com.google.common.collect.ImmutableSet;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import org.jspecify.annotations.NullUnmarked;

/**
 * A byte sink for testing that has configurable behavior.
 *
 * @author Colin Decker
 */
@NullUnmarked
public class TestByteSink extends ByteSink implements TestStreamSupplier {

  @Test fun doesNotReadEntireSource() {
    val unlimitedSource =
      object : Source {
        override fun read(
          sink: Buffer,
          byteCount: Long,
        ): Long {
          sink.writeUtf8("a".repeat(byteCount.toInt()))
          return byteCount
        }

        override fun close() {
        }

        override fun timeout() = Timeout.NONE
      }

        .method("POST", AsyncRequestBody())
        .build()
    val call = client.newCall(request)
    call.execute().use { response ->
      val sink = (request.body as AsyncRequestBody?)!!.takeSink()
      sink.writeUtf8("hey\n")
      sink.writeUtf8("whats going on\n")
      sink.close()
    }
    body.awaitSuccess()
  }

  @Test
  fun requestBodyEndsAfterResponseBody() {
    enableProtocol(Protocol.HTTP_2)

    class StreamingBody(
      private val body: String,
    ) : RequestBody() {
      override fun contentType(): MediaType? = null

      @Throws(IOException::class)
      override fun writeTo(sink: BufferedSink) {
        sink.writeUtf8(body)
      }
    }

    val expected =
      """
      |--123
      |
      |Quick
      |--123
      |
      |Brown
      |--123
      |
      |Fox

        private var cacheRequestClosed = false

        @Throws(IOException::class)
        override fun read(
          sink: Buffer,
          byteCount: Long,
        ): Long {
          val bytesRead: Long
          try {
            bytesRead = source.read(sink, byteCount)
          } catch (e: IOException) {
            if (!cacheRequestClosed) {
              cacheRequestClosed = true

/**
 * Utility methods and classes for I/O; for example input streams, output streams, readers, writers,
 * and files.
 *
 * <p>At the core of this package are the Source/Sink types: {@link ByteSource ByteSource}, {@link
 * CharSource CharSource}, {@link ByteSink ByteSink} and {@link CharSink CharSink}. They are
 * factories for I/O streams that provide many convenience methods that handle both opening and
 * closing streams for you.
 *

          .post(
            object : RequestBody() {
              override fun contentType() = null

              override fun writeTo(sink: BufferedSink) {
                for (i in 0..9) {
                  sink.writeByte(0)
                  sink.flush()
                  sleep(100)
                }
                fail("Expected connection to be closed")
              }
            },

) : ResponseBody(),
  Source {
  override fun contentType() = mediaType

  override fun contentLength() = contentLength

  override fun source() = buffer()

  override fun read(
    sink: Buffer,
    byteCount: Long,
  ): Long =
    throw IllegalStateException(
      """
      |Unreadable ResponseBody! These Response objects have bodies that are stripped:
      | * Response.cacheResponse

      while (accumulatorBitCount >= 8) {
        accumulatorBitCount -= 8
        sink.writeByte((accumulator shr accumulatorBitCount).toInt())
      }
    }

    if (accumulatorBitCount > 0) {
      accumulator = accumulator shl (8 - accumulatorBitCount)
      accumulator = accumulator or (0xffL ushr accumulatorBitCount)
      sink.writeByte(accumulator.toInt())
    }
  }

        val socket = response.socket!!
        socket.sink.buffer().use { sink ->
          socket.source.buffer().use { source ->
            sink.writeUtf8("client says hello\n")
            sink.flush()

            assertThat(source.readUtf8Line()).isEqualTo("server says hello")

            sink.writeUtf8("client says goodbye\n")
            sink.flush()