import com.azure.core.util.BinaryData;
import com.azure.storage.blob.BlobServiceClient;
import com.azure.storage.blob.BlobServiceClientBuilder;
import com.azure.storage.common.policy.RequestRetryOptions;
import com.azure.storage.common.policy.RetryPolicyType;
import com.sun.net.httpserver.HttpServer;
import org.junit.jupiter.api.Test;

import java.io.*;
import java.net.InetSocketAddress;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicInteger;

import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
 * Reproduces a bug in azure-core (verified in 1.58.0):
 * <p>
 * When BinaryData.fromStream(is, length) is called with a mark/reset-capable InputStream and
 * the first HTTP attempt fails (triggering a retry), the retry throws
 * UncheckedIOException("Stream closed") instead of re-sending the body.
 * <p>
 * Root cause — RestProxyUtils.java:66:
 * <p>
 * RestProxyUtils.validateLength() unwraps the user's InputStreamContent via toStream(),
 * wraps the result in LengthValidatingInputStream, and builds a NEW BinaryData from it.
 * The new InputStreamContent captures LengthValidatingInputStream as its replayable stream.
 * <p>
 * After the first failed attempt, Reactor's MonoUsing disposes the stream obtained from
 * toStream(). This calls LengthValidatingInputStream.close(), which delegates unconditionally
 * to inner.close() — closing the user's original InputStream.
 * <p>
 * On retry, the new InputStreamContent calls toStream() again:
 * resettableContent(LengthValidatingInputStream)
 * → LengthValidatingInputStream.reset()
 * → originalStream.reset()
 * → throws IOException("Stream closed")  ← bug
 * <p>
 * The standard Java InputStream contract does NOT guarantee that reset() works after close().
 * Most InputStream implementations throw after being closed. The Azure SDK creates an
 * impossible expectation: it builds a replayable wrapper that requires reset() to work, but
 * its own internal infrastructure calls close() on the stream between retry attempts.
 * <p>
 * Fix: LengthValidatingInputStream should not close the inner stream (or RestProxyUtils should
 * shield the caller's stream from closure between retries). Workaround on the caller side:
 * wrap the InputStream in a non-closing decorator before passing to BinaryData.fromStream().
 */
class AzureSdkInputStreamRetryBugTest {

  /**
   * Demonstrates the bug: the retry fails with "Stream closed" even though the InputStream
   * correctly implements mark/reset. The server is reachable and would return 201 on retry.
   */
  @Test
  void stageBlock_retryThrowsStreamClosed_whenLengthValidatingInputStreamClosesOriginalStream()
      throws Exception {
    try (StubHttpServer server = new StubHttpServer()) {
      BlobServiceClient client = buildClient(server.port() /*maxTries=*/);
      String blockId = Base64.getEncoder().encodeToString("block-0001".getBytes());

      byte[] data = new byte[64];
      InputStream is = new MarkResetInputStream(data);

      // BUG: should succeed on retry (server returns 201 on second request), but throws
      // UncheckedIOException("Stream closed") because LengthValidatingInputStream.close()
      // was called on `is` after the first attempt ended with 500.
      Throwable thrown = null;
      try {
        client.getBlobContainerClient("container")
            .getBlobClient("blob")
            .getBlockBlobClient()
            .stageBlock(blockId, BinaryData.fromStream(is, (long) data.length));
      } catch (Exception e) {
        thrown = e;
      }

      assertNotNull(thrown, "Expected an exception but the call succeeded");
      assertTrue(
          containsMessage(thrown, "Stream closed"),
          "Expected 'Stream closed' in exception chain, but got: " + thrown
      );
    }
  }

  /**
   * Demonstrates the workaround: wrapping the InputStream in a non-closing decorator prevents
   * LengthValidatingInputStream.close() from closing the original stream, so reset() works on
   * retry and the call succeeds.
   */
  @Test
  void stageBlock_retrySucceeds_whenStreamIsWrappedInNonClosingDecorator() throws Exception {
    try (StubHttpServer server = new StubHttpServer()) {
      BlobServiceClient client = buildClient(server.port() /*maxTries=*/);
      String blockId = Base64.getEncoder().encodeToString("block-0001".getBytes());

      byte[] data = new byte[64];
      // Workaround: non-closing wrapper prevents LengthValidatingInputStream from closing `is`.
      InputStream is = new NonClosingInputStream(new MarkResetInputStream(data));

      client.getBlobContainerClient("container")
          .getBlobClient("blob")
          .getBlockBlobClient()
          .stageBlock(blockId, BinaryData.fromStream(is, (long) data.length));

      assertTrue(server.requestCount() >= 2, "Expected retry to reach the server");
    }
  }

  // --- infrastructure ---

  private static BlobServiceClient buildClient(int port) {
    // Well-known Azurite development storage key — valid base64, accepted by the SDK for local testing.
    String connectionString =
        "DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;" +
            "AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;" +
            "BlobEndpoint=http://localhost:" + port + "/devstoreaccount1";

    RequestRetryOptions retryOptions = new RequestRetryOptions(
        RetryPolicyType.FIXED, 2, /*tryTimeoutSeconds=*/ 30,
        /*retryDelayMs=*/ 1L, /*maxRetryDelayMs=*/ 1L, null
    );

    return new BlobServiceClientBuilder()
        .connectionString(connectionString)
        .retryOptions(retryOptions)
        .buildClient();
  }

  private static boolean containsMessage(Throwable t, String fragment) {
    for (Throwable current = t; current != null; current = current.getCause()) {
      if (current.getMessage() != null && current.getMessage().contains(fragment)) {
        return true;
      }
    }
    return false;
  }

  // --- minimal HTTP stub server ---

  /**
   * Returns 500 InternalServerError on the first request, 201 Created on all subsequent ones.
   */
  private static class StubHttpServer implements AutoCloseable {

    private static final String BODY_500 =
        "<?xml version=\"1.0\" encoding=\"utf-8\"?>" +
            "<Error><Code>InternalError</Code><Message>Internal server error.</Message></Error>";

    private final HttpServer server;
    private final AtomicInteger requestCount = new AtomicInteger(0);

    StubHttpServer() throws IOException {
      server = HttpServer.create(new InetSocketAddress(0), 0);
      server.createContext(
          "/", exchange -> {
            exchange.getRequestBody().transferTo(OutputStream.nullOutputStream());
            int count = requestCount.incrementAndGet();
            if (count == 1) {
              byte[] body = BODY_500.getBytes();
              exchange.getResponseHeaders().add("Content-Type", "application/xml");
              exchange.sendResponseHeaders(500, body.length);
              exchange.getResponseBody().write(body);
            } else {
              exchange.sendResponseHeaders(201, -1);
            }
            exchange.close();
          }
      );
      server.start();
    }

    int port() {
      return server.getAddress().getPort();
    }

    int requestCount() {
      return requestCount.get();
    }

    @Override
    public void close() {
      server.stop(0);
    }
  }

  // --- stream helpers ---

  /**
   * An InputStream backed by a byte array that supports mark/reset but throws
   * IOException("Stream closed") after close() — the standard Java InputStream contract,
   * followed by most real implementations (e.g. streams backed by file descriptors or
   * off-heap buffers).
   */
  private static class MarkResetInputStream extends InputStream {

    private final byte[] data;
    private int pos = 0;
    private int markPos = 0;
    private boolean closed = false;

    MarkResetInputStream(byte[] data) {
      this.data = data;
    }

    @Override
    public int read() throws IOException {
      checkOpen();
      return pos < data.length ? data[pos++] & 0xFF : -1;
    }

    @Override
    public int read(byte[] b, int off, int len) throws IOException {
      checkOpen();
      int available = data.length - pos;
      if (available <= 0) {
        return -1;
      }
      int n = Math.min(len, available);
      System.arraycopy(data, pos, b, off, n);
      pos += n;
      return n;
    }

    @Override
    public boolean markSupported() {
      return true;
    }

    @Override
    public synchronized void mark(int limit) {
      markPos = pos;
    }

    @Override
    public synchronized void reset() throws IOException {
      checkOpen();
      pos = markPos;
    }

    @Override
    public void close() {
      closed = true;
    }

    private void checkOpen() throws IOException {
      if (closed) {
        throw new IOException("Stream closed");
      }
    }
  }

  /**
   * Delegates all InputStream operations to the wrapped stream, but close() is a no-op.
   * This prevents LengthValidatingInputStream from closing the underlying stream between
   * retry attempts — a workaround for the Azure SDK bug described above.
   */
  private static class NonClosingInputStream extends InputStream {

    private final InputStream inner;

    NonClosingInputStream(InputStream inner) {
      this.inner = inner;
    }

    @Override
    public int read() throws IOException {
      return inner.read();
    }

    @Override
    public int read(byte[] b, int off, int len) throws IOException {
      return inner.read(b, off, len);
    }

    @Override
    public boolean markSupported() {
      return inner.markSupported();
    }

    @Override
    public synchronized void mark(int limit) {
      inner.mark(limit);
    }

    @Override
    public synchronized void reset() throws IOException {
      inner.reset();
    }

    @Override
    public void close() {
      // intentionally does not close the inner stream
    }
  }
}