import argparse
import std/httpclient
import std/base64
import std/oids
import std/json
import std/strutils

var parser = newParser:
    help("Abuse the dalle-mini API from the cli!")
    option("-o", "--output", help="Write output to this folder", default=some("/tmp/"))
    option("-n", "--number", help="the number of times you'd like to request for each prompt. Expect 9 per request.", default=some("1"))
    flag("-u", "--upscale", help="Use SwiniR to upscale the image")
    arg("prompt", nargs = -1)


try:
    let arguments = parser.parse()
    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"
        })

    for parsedPrompt in arguments.prompt:
        echo "Now generating: " & parsedPrompt
        var data = "{\"prompt\":\"" & parsedPrompt & "\"}"
        var response: Response
        var max = (parseInt(arguments.number)-1)
        for i in countup(0,max):
            echo "New run..."

            echo "Waiting for response..."
            response = client.request("https://bf.dallemini.ai/generate", httpMethod = HttpPost, body=data)
            while response.status != "200 OK":
                echo "DALLE-Mini: Could not get response due to " & response.status & "! Trying again in 15 seconds!"
                sleep 15000
                response = client.request("https://bf.dallemini.ai/generate", httpMethod = HttpPost, body=data)
    
            var JSONifiedResponse = parseJSON(response.body)
            for encodedImage in JSONifiedResponse["images"]:
                var stringImage = $encodedImage
                
                if arguments.upscale:
                    var encoded = stringImage.replace("\\n").replace("\"")
                    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])
                    var filename = arguments.output & parsedPrompt & " - " & $oids.genOid() & ".png"
                    echo "writing image to " & filename
                    writeFile(filename, image)

                else:
                    var image = base64.decode(stringImage.replace("\\n").replace("\""))
                    var filename = arguments.output & parsedPrompt & " - " & $oids.genOid() & ".jpg"
                    echo "writing image to " & filename
                    writeFile(filename, image)


except ShortCircuit as e:
  if e.flag == "argparse_help":
    echo parser.help
    quit(1)
except UsageError:
  stderr.writeLine getCurrentExceptionMsg()
  quit(1)