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.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 volatile boolean isOpen;

  public PieceStorageImpl(TorrentByteStorage fileCollectionStorage,
                          BitSet availablePieces,
                          int piecesCount,
                          int pieceSize) throws IOException {
    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;
      this.fileCollectionStorage.open(false);
    } else {
      this.fileCollectionStorage.open(true);
    }
    isOpen = true;
  }

  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 synchronized 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 {
    try {
      readWriteLock.writeLock().lock();
      if (!isOpen) {
        fileCollectionStorage.open(onlyRead);
        isOpen = true;
      }
    } finally {
      readWriteLock.writeLock().unlock();
    }
  }

  @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();
      fileCollectionStorage.close();
      isOpen = false;
    } finally {
      readWriteLock.writeLock().unlock();
    }
  }

  @Override
  public void delete() throws IOException {
    close();
    fileCollectionStorage.delete();
  }
}

@@ -161,16 +161,66 @@
 BitSet;%0A
+import java.util.concurrent.atomic.AtomicBoolean;%0A
 import j
@@ -260,16 +260,16 @@
 teLock;%0A
-
 import j
@@ -613,18 +613,21 @@
 ate 
-volatile b
+final AtomicB
 oole
@@ -1320,28 +1320,47 @@
    isOpen = 
+new AtomicBoolean(
 true
+)
 ;%0A  %7D%0A%0A  pri
@@ -2656,77 +2656,39 @@
     
-try %7B%0A      readWriteLock.writeLock().lock();%0A      if (!isOpen
+if (!isOpen.getAndSet(true)
 ) %7B%0A
-  
     
@@ -2731,97 +2731,8 @@
 d);%0A
-        isOpen = true;%0A      %7D%0A    %7D finally %7B%0A      readWriteLock.writeLock().unlock();%0A
     
@@ -4310,16 +4310,19 @@
 Open
- = 
+.set(
 false
+)
 ;%0A  

fix possible race condition