import android.util.AtomicFile
import okio.Buffer
import okio.Sink
import okio.sink
import okio.source
import java.io.FileOutputStream
import java.io.IOException

fun AtomicFile.source() = openRead().source()

fun AtomicFile.sink(): Sink {
  return object : Sink {
    val fos: FileOutputStream = startWrite()
    val sink: Sink = fos.sink()
    var terminated = false

    override fun write(source: Buffer, byteCount: Long) {
      guardedIo {
        sink.write(source, byteCount)
      }
    }

    override fun flush() {
      guardedIo {
        sink.flush()
      }
    }

    override fun timeout() = sink.timeout()

    override fun close() {
      guardedIo {
        finishWrite(fos)
      }
      terminated = true
    }

    private fun <T> guardedIo(block: () -> T): T {
      if (terminated) throw IOException("terminated")
      try {
        return block()
      } catch (e: IOException) {
        terminated = true
        failWrite(fos)
        throw e
      }
    }
  }
}