import std/base64
import std/httpclient
import std/os
import std/json
import std/strutils

var client = newHttpClient(userAgent="Mozilla/5.0 (Windows NT 10.0; rv:91.0) Gecko/20100101 Firefox/91.0")
client.headers = newHttpHeaders({
 "Accept":"application/json",
 "Accept-Language":"en-US,en;q=0.5",
 "Accept-Encoding":"gzip, deflate, br",
 "Referer":"https://hf.space/",
 "Content-Type":"application/json",
 "Origin":"https://hf.space",
 "DNT":"1",
 "Connection":"keep-alive",
 "Sec-Fetch-Dest":"empty",
 "Sec-Fetch-Mode":"cors",
 "Sec-Fetch-Site":"cross-site"
    })


let encoded = base64.encode(readFile(os.paramStr(1)))
var filename = os.paramStr(1)
var data = %*{
    "action": "predict",
    "data": [
        "data:image/png;base64," & encoded
    ],
    "session_hash": "faq0rxk7k3"
}

echo "Starting to upscale!"
var response = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/push/", httpMethod = HttpPost, body = $data)
while response.status != "200 OK":
    echo "SwinIR: Could not get response due to " & response.status & "! Trying again in 15 seconds!"
    sleep 15000
    response = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/push/", httpMethod = HttpPost, body = $data)

var hash = parseJSON(response.body)
delete(hash, "queue_position")


var isItDoneYet = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/status/", httpMethod = HttpPost, body = $hash)
var doneStatus = parseJSON(isItDoneYet.body)
var waitingList = parseJSON(response.body)
echo "we are currently in position number " &  $waitingList["queue_position"] & " in the queue!"

while $doneStatus["status"] == "\"QUEUED\"" or $doneStatus["status"] == "\"PENDING\"":
    echo "Waiting for 3 more seconds"
    sleep 3000
    isItDoneYet = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/status/", httpMethod = HttpPost, body = $hash)
    doneStatus = parseJSON(isItDoneYet.body)

if $doneStatus["status"] == "\"FAILED\"":
    echo "ERROR WHILE GETTING STATUS:"
    echo $doneStatus
    quit 1

echo "Done!"
var outputdata = $doneStatus["data"]["data"][0]
var image = base64.decode(outputdata.replace("\"")[22..^1])
echo "writing image to " & filename
writeFile(filename, image)