package com.turn.ttorrent.client.storage;

import org.jetbrains.annotations.Nullable;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.BitSet;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class PieceStorageImpl implements PieceStorage {

  private final TorrentByteStorage fileCollectionStorage;
  private final ReadWriteLock readWriteLock;

  @Nullable
  private volatile BitSet availablePieces;
  private final int piecesCount;
  private final int pieceSize;
  private final AtomicBoolean isOpen;

  public PieceStorageImpl(TorrentByteStorage fileCollectionStorage,
                          BitSet availablePieces,
                          int piecesCount,
                          int pieceSize) {
    this.fileCollectionStorage = fileCollectionStorage;
    this.readWriteLock = new ReentrantReadWriteLock();
    this.piecesCount = piecesCount;
    this.pieceSize = pieceSize;
    BitSet bitSet = new BitSet(piecesCount);
    bitSet.or(availablePieces);
    if (bitSet.cardinality() != piecesCount) {
      this.availablePieces = bitSet;
    }
    isOpen = new AtomicBoolean(false);
  }

  private void checkPieceIndex(int pieceIndex) {
    if (pieceIndex < 0 || pieceIndex >= piecesCount) {
      throw new IllegalArgumentException("Incorrect piece index " + pieceIndex + ". Piece index must be positive less than" + piecesCount);
    }
  }

  @Override
  public void savePiece(int pieceIndex, byte[] pieceData) throws IOException {
    checkPieceIndex(pieceIndex);
    try {
      readWriteLock.writeLock().lock();

      BitSet availablePieces = this.availablePieces;

      boolean isFullyDownloaded = availablePieces == null;

      if (isFullyDownloaded) return;

      if (availablePieces.get(pieceIndex)) return;

      openStorageIsNecessary(false);

      long pos = pieceIndex;
      pos = pos * pieceSize;
      ByteBuffer buffer = ByteBuffer.wrap(pieceData);
      fileCollectionStorage.write(buffer, pos);

      availablePieces.set(pieceIndex);
      boolean isFullyNow = availablePieces.cardinality() == piecesCount;
      if (isFullyNow) {
        this.availablePieces = null;
        fileCollectionStorage.finish();
        fileCollectionStorage.close();
        fileCollectionStorage.open(true);
      }
    } finally {
      readWriteLock.writeLock().unlock();
    }
  }

  private void openStorageIsNecessary(boolean onlyRead) throws IOException {
    if (!isOpen.getAndSet(true)) {
      fileCollectionStorage.open(onlyRead);
    }
  }

  @Override
  public byte[] readPiecePart(int pieceIndex, int offset, int length) throws IOException {
    checkPieceIndex(pieceIndex);
    try {
      readWriteLock.readLock().lock();

      BitSet availablePieces = this.availablePieces;
      if (availablePieces != null && !availablePieces.get(pieceIndex)) {
        throw new IllegalArgumentException("trying reading part of not available piece");
      }

      openStorageIsNecessary(availablePieces == null);

      ByteBuffer buffer = ByteBuffer.allocate(length);
      long pos = pieceIndex;
      pos = pos * pieceSize + offset;
      fileCollectionStorage.read(buffer, pos);
      return buffer.array();
    } finally {
      readWriteLock.readLock().unlock();
    }
  }

  @Override
  public boolean isFinished() {
    try {
      readWriteLock.readLock().lock();
      return availablePieces == null;
    } finally {
      readWriteLock.readLock().unlock();
    }
  }

  @Override
  public BitSet getAvailablePieces() {
    try {
      readWriteLock.readLock().lock();
      BitSet result = new BitSet(piecesCount);

      BitSet availablePieces = this.availablePieces;
      boolean isFullyDownloaded = availablePieces == null;

      if (isFullyDownloaded) {
        result.set(0, piecesCount);
        return result;
      }
      result.or(availablePieces);
      return result;
    } finally {
      readWriteLock.readLock().unlock();
    }
  }

  @Override
  public void close() throws IOException {
    try {
      readWriteLock.writeLock().lock();
      if (!isOpen.get()) return;
      fileCollectionStorage.close();
      isOpen.set(false);
    } finally {
      readWriteLock.writeLock().unlock();
    }
  }
}

@@ -165,58 +165,8 @@
 et;%0A
-import java.util.concurrent.atomic.AtomicBoolean;%0A
 impo
@@ -210,16 +210,16 @@
 teLock;%0A
+
 import j
@@ -430,16 +430,72 @@
 eLock;%0A%0A
+  private final Object openStorageLock = new Object();%0A%0A
   @Nulla
@@ -619,21 +619,18 @@
 ate 
-final AtomicB
+volatile b
 oole
@@ -1208,32 +1208,13 @@
 n = 
-new AtomicBoolean(
 false
-)
 ;%0A  
@@ -2516,24 +2516,8 @@
 Open
-.getAndSet(true)
 ) %7B%0A
@@ -2560,16 +2560,37 @@
 yRead);%0A
+      isOpen = true;%0A
     %7D%0A  
@@ -3003,16 +3003,57 @@
      %7D%0A%0A
+      synchronized (openStorageLock) %7B%0A  
       op
@@ -3098,16 +3098,24 @@
 = null);
+%0A      %7D
 %0A%0A      
@@ -4181,14 +4181,8 @@
 Open
-.get()
 ) re
@@ -4240,19 +4240,16 @@
 Open
-.set(
+ = 
 false
-)
 ;%0A  

fix possible race condition when one thread can see that storage is opened but really other thread only set flag as "open"