Multi Thread Http Downloader implemented in kotlin script

ktd.kts

#! /usr/bin/env kscript


@file:DependsOn("com.xenomachina:kotlin-argparser:2.0.7")
@file:DependsOn("me.tongfei:progressbar:0.7.1")

import com.xenomachina.argparser.ArgParser
import com.xenomachina.argparser.InvalidArgumentException
import com.xenomachina.argparser.SystemExitException
import com.xenomachina.argparser.default
import me.tongfei.progressbar.ProgressBar
import me.tongfei.progressbar.ProgressBarStyle
import java.io.RandomAccessFile
import java.net.HttpURLConnection
import java.net.URL
import java.util.*
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicLong
import kotlin.system.exitProcess


val param = parseArgument()

// curl -I http://i.imgur.com/z4d4kWk.jpg
val totalReadCount = AtomicLong()
val speedMeter = AtomicLong()
val url = URL(param.url)

val conn = url.openConnection() as HttpURLConnection
val threadPool = Executors.newCachedThreadPool()
val outputFile = RandomAccessFile(param.file, "rw")

conn.doInput = true
conn.doOutput = true
conn.connect()

if (conn.responseCode !in 200..299) {
    conn.disconnect()
    println("connection to url failed with code ${conn.responseCode}, url:${param.url}")
    exitProcess(-1)
}

println("url:     ${param.url}")
println("file:    ${param.file}")
println("threads: ${param.threads}")

/*
Content-Length: 146515
Accept-Ranges: bytes
 */
val contentLength = try {
    conn.getHeaderField("Content-Length").toLong()
} catch (e: NumberFormatException) {
    -1L
}
val acceptRange = conn.getHeaderField("Accept-Ranges") == "bytes"

conn.disconnect()

if (acceptRange && contentLength != -1L) {
    val segmentSize = contentLength / param.threads
    for (job in 0 until param.threads) {
        threadPool.submit {
            downloadSegment(
                    segmentSize * job, if (job != param.threads - 1) {
                segmentSize
            } else {
                contentLength - segmentSize * (param.threads - 1)
            }
            )
        }
    }
    println("size:    ${formatSize(contentLength)}")
    println("segment: ${formatSize(segmentSize)}")
} else {
    println("server doesn't support http-range, download with single thread!")

    threadPool.submit {
        downloadSegment(0, Long.MAX_VALUE, false)
    }
}

threadPool.shutdown()

println()
val bar = ProgressBar("ktd", 100)
while (!threadPool.isTerminated) {
    val speed = speedMeter.getAndSet(0L)

    if (contentLength != -1L) {
        val progress = totalReadCount.get() * 100L / contentLength
        bar.stepTo(progress)
    } else {
        bar.maxHint(-1)
    }
    bar.extraMessage = "${formatSize(speed)}/s"

    Thread.sleep(1000L)
}
bar.stepTo(bar.max)
bar.close()

// impl

fun downloadSegment(startByte: Long, size: Long, range: Boolean = true) {
    val conn = url.openConnection() as HttpURLConnection

    // Range: bytes=200-1000, 2000-6576, 19000-
    if (range) {
        conn.setRequestProperty("Range", "bytes=${startByte}-${startByte + size}")
    }
    conn.doInput = true
    conn.doOutput = true
    conn.connect()

    val buffer = ByteArray(64 * 1024)
    val input = conn.inputStream

    var allBytes = 0L
    while (allBytes < size) {
        val count = input.read(buffer)
        if (count == -1) {
            break
        }

        synchronized(outputFile) {
            outputFile.seek(startByte + allBytes)
            outputFile.write(buffer, 0, count)
        }

        totalReadCount.addAndGet(count.toLong())
        speedMeter.addAndGet(count.toLong())
        allBytes += count
    }

    conn.disconnect()
}

fun formatSize(size: Long): String {
    val K = 1024
    val M = K * K
    val G = M * K

    return when {
        size >= G -> String.format("%.2f GB", size.toDouble() / G)
        size >= M -> String.format("%.2f MB", size.toDouble() / M)
        size >= K -> String.format("%.2f KB", size.toDouble() / K)
        else -> String.format("%d B", size)
    }
}

class CLIArguments(parser: ArgParser) {
    val threads by parser.storing(
            "-t",
            "--threads",
            argName = "THREADS_COUNT",
            help = "threads used to download") {
        try {
            toInt()
        } catch (e: NumberFormatException) {
            throw InvalidArgumentException("invalid thread count [$this]")
        }
    }
            .default(4)
            .addValidator {
                if (this.value <= 0) {
                    throw InvalidArgumentException("thread count must be positive")
                }
            }

    val url by parser.positional("URL", "url to download")

    val file by parser.positional("FILE", "destination filename").default {
        val u = URL(url)
        val result = when {
            u.path.isNotBlank() -> {
                var path = u.path
                var idx = path.lastIndexOf("/")
                if (idx != -1 && idx != path.length - 1) {
                    path = path.substring(idx + 1)
                }
                idx = path.indexOf('?')
                if (idx != -1) {
                    path.substring(0, idx)
                } else {
                    path
                }
            }
            u.host.isNotBlank() -> u.host
            else -> ""
        }
        if (result.isBlank()) "ktd-file" else result
    }
}

fun parseArgument(): CLIArguments {
    try {
        return ArgParser(args).parseInto(::CLIArguments)
    } catch (e: SystemExitException) {
        e.printAndExit("ktd")
    }
}