local version = 022
local expect = require("cc.expect")
local cha = (function()
    -- Chacha20 cipher in ComputerCraft
    -- By Anavrins
    -- For help and details, you can DM me on Discord (Anavrins#4600)
    -- MIT License
    -- Pastebin: https://pastebin.com/GPzf9JSa
    -- Last updated: March 27 2020

    local mod32 = 2^32
    local bor = bit32.bor
    local bxor = bit32.bxor
    local band = bit32.band
    local blshift = bit32.lshift
    local brshift = bit32.arshift

    local tau = {("expand 16-byte k"):byte(1,-1)}
    local sigma = {("expand 32-byte k"):byte(1,-1)}
    local null32 = {("A"):rep(32):byte(1,-1)}
    local null12 = {("A"):rep(12):byte(1,-1)}

    local function rotl(n, b)
        local s = n/(2^(32-b))
        local f = s%1
        return (s-f) + f*mod32
    end

    local function quarterRound(s, a, b, c, d)
        s[a] = (s[a]+s[b])%mod32; s[d] = rotl(bxor(s[d], s[a]), 16)
        s[c] = (s[c]+s[d])%mod32; s[b] = rotl(bxor(s[b], s[c]), 12)
        s[a] = (s[a]+s[b])%mod32; s[d] = rotl(bxor(s[d], s[a]), 8)
        s[c] = (s[c]+s[d])%mod32; s[b] = rotl(bxor(s[b], s[c]), 7)
        return s
    end

    local function hashBlock(state, rnd)
        local s = {unpack(state)}
        for i = 1, rnd do
            local r = i%2==1
            s = r and quarterRound(s, 1, 5,  9, 13) or quarterRound(s, 1, 6, 11, 16)
            s = r and quarterRound(s, 2, 6, 10, 14) or quarterRound(s, 2, 7, 12, 13)
            s = r and quarterRound(s, 3, 7, 11, 15) or quarterRound(s, 3, 8,  9, 14)
            s = r and quarterRound(s, 4, 8, 12, 16) or quarterRound(s, 4, 5, 10, 15)
        end
        for i = 1, 16 do s[i] = (s[i]+state[i])%mod32 end
        return s
    end

    local function LE_toInt(bs, i)
        return (bs[i+1] or 0)+
        blshift((bs[i+2] or 0), 8)+
        blshift((bs[i+3] or 0), 16)+
        blshift((bs[i+4] or 0), 24)
    end

    local function initState(key, nonce, counter)
        local isKey256 = #key == 32
        local const = isKey256 and sigma or tau
        local state = {}

        state[ 1] = LE_toInt(const, 0)
        state[ 2] = LE_toInt(const, 4)
        state[ 3] = LE_toInt(const, 8)
        state[ 4] = LE_toInt(const, 12)

        state[ 5] = LE_toInt(key, 0)
        state[ 6] = LE_toInt(key, 4)
        state[ 7] = LE_toInt(key, 8)
        state[ 8] = LE_toInt(key, 12)
        state[ 9] = LE_toInt(key, isKey256 and 16 or 0)
        state[10] = LE_toInt(key, isKey256 and 20 or 4)
        state[11] = LE_toInt(key, isKey256 and 24 or 8)
        state[12] = LE_toInt(key, isKey256 and 28 or 12)

        state[13] = counter
        state[14] = LE_toInt(nonce, 0)
        state[15] = LE_toInt(nonce, 4)
        state[16] = LE_toInt(nonce, 8)

        return state
    end

    local function serialize(state)
        local r = {}
        for i = 1, 16 do
            r[#r+1] = band(state[i], 0xFF)
            r[#r+1] = band(brshift(state[i], 8), 0xFF)
            r[#r+1] = band(brshift(state[i], 16), 0xFF)
            r[#r+1] = band(brshift(state[i], 24), 0xFF)
        end
        return r
    end

    local mt = {
        __tostring = function(a) return string.char(unpack(a)) end,
        __index = {
            toHex = function(self) return ("%02x"):rep(#self):format(unpack(self)) end,
            isEqual = function(self, t)
                if type(t) ~= "table" then return false end
                if #self ~= #t then return false end
                local ret = 0
                for i = 1, #self do
                    ret = bor(ret, bxor(self[i], t[i]))
                end
                return ret == 0
            end
        }
    }

    local function crypt(data, key, nonce, cntr, round)
        assert(type(key) == "table", "ChaCha20: Invalid key format ("..type(key).."), must be table")
        assert(type(nonce) == "table", "ChaCha20: Invalid nonce format ("..type(nonce).."), must be table")
        assert(#key == 16 or #key == 32, "ChaCha20: Invalid key length ("..#key.."), must be 16 or 32")
        assert(#nonce == 12, "ChaCha20: Invalid nonce length ("..#nonce.."), must be 12")

        local data = type(data) == "table" and {unpack(data)} or {tostring(data):byte(1,-1)}
        cntr = tonumber(cntr) or 1
        round = tonumber(round) or 20

        local out = {}
        local state = initState(key, nonce, cntr)
        local blockAmt = math.floor(#data/64)
        for i = 0, blockAmt do
            local ks = serialize(hashBlock(state, round))
            state[13] = (state[13]+1) % mod32

            local block = {}
            for j = 1, 64 do
                block[j] = data[((i)*64)+j]
            end
            for j = 1, #block do
                out[#out+1] = bxor(block[j], ks[j])
            end

            --[[if i % 1000 == 0 then
                os.queueEvent("")
                os.pullEvent("")
            end]]
        end
        return setmetatable(out, mt)
    end

    local function genNonce(len)
        local nonce = {}
        for i = 1, len do
            nonce[i] = math.random(0, 0xFF)
        end
        return setmetatable(nonce, mt)
    end

    local obj = {}
    local mtrng = {['__index'] = obj}
    local function newRNG(seed)
    local objVars = {}
    objVars.seed = seed
    objVars.cnt = 0
    objVars.block = {}
    return setmetatable(objVars, mtrng)
    end

    -- Specify how many bytes to return
    -- 1 Byte is 8 bits, returns int between 0 and 255
    -- 2 Bytes is 16 bits, returns int up to 65535
    -- 4 Bytes is 32 bits, returns int up to 4294967295
    -- Max of 6 Bytes is 48 bits, returns int up to 281474976710655
    function obj:nextInt(byte)
    if not byte or byte < 1 or byte > 6 then error("Can only return 1-6 bytes", 2) end
    local output = 0
    for i = 0, byte-1 do
        if #self.block == 0 then
        self.cnt = self.cnt + 1
        self.block = crypt(null32, self.seed, null12, self.cnt)
        end
        local newByte = table.remove(self.block)
        output = output + (newByte * (2^(8*i)))
    end
    return output
    end

    return {
        crypt = crypt,
        genNonce = genNonce,
        newRNG = newRNG
    }
end)()

-- START OF API --

settings.define("shironeko.pkey", {
    description = "Your Shironeko pkey, obtained from \\sn me",
    type = "string"
})

-- Wrap the modem with the given name with the given channel and privatekey
-- Creates a functionally secure modem peripheral with useful methods related to communication
local function wrap(name, channel, pkey)
    local key = table.pack(pkey:gsub("-", ""):byte(1,-1))
    local modem = assert(peripheral.wrap(name), "No such modem "..tostring(name))
    modem.open(channel)
    modem.name = name
    local transmit = modem.transmit
    modem.transmit = function(msg)
        local nonce = cha.genNonce(12)
        local data = cha.crypt(textutils.serialize(msg), key, nonce)
        transmit(channel,channel,{
            data = data,
            nonce = nonce
        })
    end
    modem.response = function(action, ok, result)
        modem.transmit({
            action = action,
            ok = ok,
            result = result
        })
    end

    modem.listen = function(callback, timeout)
        local ok = false
        local out
        parallel.waitForAny(function()
            local done = false
            while not done do
                local _, side, ch, _, msg = os.pullEvent("modem_message")
                if side == name and ch == channel then
                    local valid, response = pcall(function()
                        return textutils.unserialize(tostring(cha.crypt(msg.data, key, msg.nonce)))
                    end)
                    if valid then
                        if callback then
                            done, ok, out = callback(response)
                        else
                            done, ok, out = true, true, response
                        end
                    end
                end
            end
            return ok, out
        end,
        function()
            sleep(timeout or 5)
            out = "timeout"
        end)
        return ok, out
    end
    modem.await = function(callback)
        local done, ok = false, false
        local out
        while not done do
            local _, side, ch, _, msg = os.pullEvent("modem_message")
            if side == name and ch == channel then
                local valid, response = pcall(function()
                    return textutils.unserialize(tostring(cha.crypt(msg.data, key, msg.nonce)))
                end)
                if valid then
                    callback(response)
                end
            end
        end
    end
    return modem
end

local function responseCallback(action)
    return function(data)
        if data and data.action == action then
            return true, data.ok, data.result
        end
    end
end

-- Initialize a modem for use with Shironeko. Should be wireless.
-- Provides the run, send, and close functions.
local function init(modemName)
    expect(1, modemName, "string")
    local pkey = assert(settings.get("shironeko.pkey"), "Missing shironeko.pkey setting. See \\sn me to access it.")
    local modem = wrap(modemName, 4625, pkey)

    -- Emit events into the CC event loop when Shironeko sends them. Run this in parallel with your program.
    local function run()
        modem.await(function(data)
            if data.event then
                os.queueEvent("shironeko_"..data.event, data.data)
            end
        end)
    end

    -- Possible events emitted by run:
    -- shironeko_receive:
    -- data.items (table): table of items sent. May contain duplicate items that merged into the same stack on delivery.
    -- data.from (string): user that sent the items (the sender).
    -- data.fromLabel (string): label from which items were sent (sender's label).
    -- data.to (string): user these items were sent to (the receiver).
    -- data.toLabel (string): label to which items were sent (receiver's label).


    -- Send the provided slots to the provided user using the optional labels. fromLabel is one of the senders labels, toLabel is one of the recipients labels.
    -- Returns:
    -- boolean ok - Whether the operation was successful or not.
    -- if ok:
    -- table info - A table containing the same values as what would be provided by the recipient's shironeko_receive event, seen above.
    -- if not ok:
    -- string error - The nature of how the operation failed. See below function for possible errors.
    local function send(user, slots, fromLabel, toLabel)
        expect(1, user, "string")
        expect(2, slots, "table", "number")
        expect(3, fromLabel, "string", "nil")
        expect(4, toLabel, "string", "nil")
        modem.transmit({
            action = "send",
            user = user,
            slots = slots,
            fromLabel = fromLabel,
            toLabel = toLabel,
            version = version
        })
        return modem.listen(responseCallback("send"))
    end
    -- Possible errors for send:
    -- "timeout" - The operation timed out/Shironeko didn't respond.
    -- "version" - The API version used is out of date.
    -- "no_such_label_sender" - The sender (you) did not have the given fromLabel.
    -- "no_outbound_storage" - The sender's (you) label did not have a designated outbound storage.
    -- "no_such_user" - The recipient does not exist and/or does not have a Shironeko user.
    -- "no_such_label_recipient" - The recipient did not have the given toLabel.
    -- "no_inbound_storage" - The recipient's label did not have a designated inbound storage.
    -- "blocked" - You are blocked from performing this send operation by the recipient.
    -- "same_storage" - The fromLabel and toLabel point to the same storage.
    -- "recipient_full" - The recpient's inbound storage is full.


    -- Sends an amount of items specified by a list of tables containing item names, counts, and optional nbt values.
    -- In other words, the items in the list are identical to ones provided by generic inventory's .list() method.
    -- The only difference being that you can specify a count greater than 64 to easily send multiple stacks of an item.
    -- Returns:
    -- boolean ok - Whether the operation was successful or not.
    -- if ok:
    -- table info - A table containing the same values as what would be provided by the recipient's shironeko_receive event, seen above.
    -- if not ok:
    -- string error - The nature of how the operation failed. See below function for possible errors.
    local function sendItems(user, items, fromLabel, toLabel)
        expect(1, user, "string")
        expect(2, items, "table")
        expect(3, fromLabel, "string", "nil")
        expect(4, toLabel, "string", "nil")
        modem.transmit({
            action = "send_items",
            user = user,
            items = items,
            fromLabel = fromLabel,
            toLabel = toLabel,
            version = version
        })
        return modem.listen(responseCallback("send_items"))
    end
    -- Possible errors for sendItems:
    -- "timeout" - The operation timed out/Shironeko didn't respond.
    -- "version" - The API version used is out of date.
    -- "no_such_label_sender" - The sender (you) did not have the given fromLabel.
    -- "no_outbound_storage" - The sender's (you) label did not have a designated outbound storage.
    -- "no_such_user" - The recipient does not exist and/or does not have a Shironeko user.
    -- "no_such_label_recipient" - The recipient did not have the given toLabel.
    -- "no_inbound_storage" - The recipient's label did not have a designated inbound storage.
    -- "blocked" - You are blocked from performing this send operation by the recipient.
    -- "same_storage" - The fromLabel and toLabel point to the same storage.
    -- "recipient_full" - The recpient's inbound storage is full.


    -- Get a simple list of users who have a label "main" that can accept items.
    -- Returns:
    -- boolean ok - Whether the operation was successful or not.
    -- if ok:
    -- table users - List of users that have a "main" label that can accept items.
    -- if not ok:
    -- string error - The nature of how the operation failed. See below function for possible errors.
    local function getUsers()
        modem.transmit({
            action = "get_users",
            version = version
        })
        return modem.listen(responseCallback("get_users"))
    end
    -- Possible errors for send:
    -- "timeout" - The operation timed out/Shironeko didn't respond.
    -- "version" - The API version used is out of date.

    -- Close the modem channel used by Shironeko
    local function close()
        modem.close(4625)
    end

    return {
        run = run,
        send = send,
        sendItems = sendItems,
        getUsers = getUsers,
        close = close
    }
end

return {
    init = init,
}