solutions/CCTF_Solutions_main/src/main/kotlin/hu/bmeta/cctf/ZoKrates.kt

package hu.bmeta.cctf


import java.io.File
import java.lang.RuntimeException
import java.net.URL

object Zokrates {
    enum class ProvingScheme{
        G16,GM17
    }

    private var isInitDone = false
    private lateinit var zokratesPath: String

    private fun checkInit() {
        if (isInitDone.not()) {
            throw RuntimeException("")
        }
    }

    private fun installZokrates() {
        val tmp = File(System.getProperty("user.dir") + "/tmp").also { file ->
            file.mkdir()
        }
        val dlFile = File(tmp, "zokrates.sh").also {
            it.writeText(URL("https://raw.githubusercontent.com/ZoKrates/ZoKrates/master/scripts/one_liner.sh").readText())
            it.setExecutable(true)
        }

        val pb = ProcessBuilder(dlFile.absoluteFile.absolutePath.toString(), "--to", ".zokrates")
        val process = pb.start()

        val exitCode = process.waitFor()
        if (exitCode != 0) {
            throw RuntimeException("")
        }

        dlFile.delete()
        tmp.delete()
    }

    fun initZokrates() {
        if (File(System.getProperty("user.dir") + "/.zokrates/bin/zokrates").exists().not()) {
            installZokrates()
        }

        zokratesPath = System.getProperty("user.dir") + "/.zokrates/bin/zokrates"
        isInitDone = true
    }

    fun compile(path: String, outPut: String = "out") {
        checkInit()
        compile(File(path), outPut)
    }

    fun compile(file: File, outPut: String = "out") {
        checkInit()
        val pb = ProcessBuilder(zokratesPath, "compile", "-i", file.absoluteFile.absolutePath.toString(), "-o", outPut)

        val process = pb.start()

        val exitCode = process.waitFor()
        if (exitCode != 0) {
            val output = process.inputStream.bufferedReader().use { it.readText() }
            throw throw RuntimeException(output)
        }
    }

    fun computeWithness(args: List<String>, input: String = "out", output: String = "witness"): List<String> {
        checkInit()
        val sb = StringBuilder()
        var prefix = ""
        args.forEach {
            sb.append(prefix)
            prefix = " "
            sb.append(it)
        }
        println(sb.toString())
        val command = mutableListOf(zokratesPath, "compute-witness", "-i", input, "-o", output, "-a").apply {
            addAll(args)
        }
        val processBuilder = ProcessBuilder(command)

        val process = processBuilder.start()

        val exitCode = process.waitFor()
        if (exitCode != 0) {
            val output1 = process.inputStream.bufferedReader().use { it.readText() }
            println("Error $output1")
            throw throw RuntimeException(output1)
        }
        return readWitness(output)
    }

    fun readWitness(witness: String = "witness"): List<String> {
        return File(witness).readText().split("\n").filter { it.startsWith("~out_") }.sorted().map { it.split(' ')[1] }
    }

    fun setup(input: String = "out") {
        val pb = ProcessBuilder(zokratesPath, "setup", "-i", input)
        val process = pb.start()
        val exitCode = process.waitFor()
        if (exitCode != 0) {
            val output = process.inputStream.bufferedReader().use { it.readText() }
            println("Error $output")
            throw RuntimeException(output)
        }
    }

    fun generateProof(input: String = "out", witness: String = "witness", proof: String = "proof.json", scheme: ProvingScheme = ProvingScheme.G16) {
        val ps =  if (scheme == ProvingScheme.G16) "g16" else "gm17"
        val pb = ProcessBuilder(zokratesPath, "generate-proof", "-i", input, "-w", witness, "-j", proof,"-s",ps)
        val process = pb.start()
        val exitCode = process.waitFor()
        if (exitCode != 0) {
            val output = process.inputStream.bufferedReader().use { it.readText() }
            println("Error $output")
            throw RuntimeException(output)
        }
    }
    fun verify() {
        val pb = ProcessBuilder(zokratesPath, "verify")
        val process = pb.start()
        val exitCode = process.waitFor()
        if (exitCode != 0) {
            val output = process.inputStream.bufferedReader().use { it.readText() }
            println("Error $output")
            throw RuntimeException(output)
        }
    }


    fun exportVerifier(input: String = "out", output: String = "verifier.sol") {
        val pb = ProcessBuilder(zokratesPath, "export-verifier")
        val process = pb.start()
        val exitCode = process.waitFor()
        if (exitCode != 0) {
            val output = process.inputStream.bufferedReader().use { it.readText() }
            println("Error $output")
            throw RuntimeException(output)
        }
    }
}