import java.*
import java.lang.*
import java.awt.*
import java.awt.image.*
import java.io.*
import javax.imageio.*
import java.nio.ByteBuffer
import java.nio.file.*
import kotlin.experimental.and
import kotlin.experimental.or

object Lab2 {
    @JvmStatic
    fun main(args: Array<String>) {

        if (args.size < 3 || args.size > 4) {
            System.exit(1)
        }

        if (args[0] == "encrypt" && args.size == 4) {

            val inputFile = args[1]
            val outputFile = args[2]

            var src: BufferedImage? = null
            try {
                src = ImageIO.read(File(inputFile))
            } catch (e: IOException) {
                println("Lab2: could not open file $inputFile")
                System.exit(1)
            }

            val img = BufferedImage(src!!.width, src.height, BufferedImage.TYPE_4BYTE_ABGR)
            img.graphics.drawImage(src, 0, 0, null)

            val width = img.width
            val height = img.height
            val max_size = width * height - 4

            println("Image width: $width")
            println("Image height: $height")
            println("Max encryption length: $max_size")

            var to_encrypt_bytes: ByteArray? = null
            try {
                val fileToEncrypt = Paths.get(args[3])
                to_encrypt_bytes = Files.readAllBytes(fileToEncrypt)
            } catch (e: IOException) {
                println("Could not open file " + args[3])
                System.exit(1)
            }

            val encrypt_length = to_encrypt_bytes!!.size
            println("Encryption length: $encrypt_length")
            if (encrypt_length >= 2147483647) {
                println("Encryption data is too large")
                System.exit(1)
            }
            if (encrypt_length >= max_size) {
                println("Encryption data is too large for selected image")
                System.exit(1)
            }

            val encrypt_length_bytes = ByteBuffer.allocate(4).putInt(encrypt_length).array()
            var cur_byte = 0
            // Первые 4 пикселя содержат длинну сообщения
            for (y in 0 until height) {
                for (x in 0 until width) {
                    if (cur_byte < 4) {
                        val rgb_int = img.getRGB(x, y)
                        val color = Color(rgb_int, true)
                        val cur_encrypt_byte = encrypt_length_bytes[cur_byte]
                        // Первые 2 бита в красном
                        var red = color.red
                        val new_red = cur_encrypt_byte.toInt().ushr(6) and 0x03 or 0xfc
                        red = red or 0x03 and new_red
                        // Вторые 2 бита в зеленом
                        var green = color.green
                        val new_green = cur_encrypt_byte.toInt().ushr(4) and 0x03 or 0xfc
                        green = green or 0x03 and new_green
                        // Третьи два бита в синем
                        var blue = color.blue
                        val new_blue = cur_encrypt_byte.toInt().ushr(2) and 0x03 or 0xfc
                        blue = blue or 0x03 and new_blue
                        // Последние 2 бита в альфа
                        var alpha = color.alpha
                        val new_alpha = cur_encrypt_byte and 0x03 or 0xfc.toByte()
                        alpha = alpha or 0x03 and new_alpha.toInt()
                        // Новый цвет
                        val new_col = Color(red, green, blue, alpha)
                        img.setRGB(x, y, new_col.rgb)
                        cur_byte++
                    } else {
                        break
                    }
                }
                if (cur_byte >= 4) {
                    break
                }
            }

            cur_byte = 0
            for (y in 4 until height) {
                for (x in 4 until width) {
                    if (cur_byte < encrypt_length) {
                        val rgb_int = img.getRGB(x, y)
                        val color = Color(rgb_int, true)
                        val cur_encrypt_byte = to_encrypt_bytes[cur_byte]
                        // Первые 2 бита в красном
                        var red = color.red
                        val new_red = cur_encrypt_byte.toInt().ushr(6) and 0x03 or 0xfc
                        red = red or 0x03 and new_red
                        // Вторые два бита в зеленом
                        var green = color.green
                        val new_green = cur_encrypt_byte.toInt().ushr(4) and 0x03 or 0xfc
                        green = green or 0x03 and new_green
                        // Третьи 2 бита в синем
                        var blue = color.blue
                        val new_blue = cur_encrypt_byte.toInt().ushr(2) and 0x03 or 0xfc
                        blue = blue or 0x03 and new_blue
                        // Последние 2 бита в альфа
                        var alpha = color.alpha
                        val new_alpha = cur_encrypt_byte and 0x03 or 0xfc.toByte()
                        alpha = alpha or 0x03 and new_alpha.toInt()
                        // make new color
                        val new_col = Color(red, green, blue, alpha)
                        img.setRGB(x, y, new_col.rgb)
                        cur_byte++
                    } else {
                        break
                    }
                }
                if (cur_byte >= encrypt_length) {
                    break
                }
            }

            try {
                val output = File(outputFile)
                ImageIO.write(img, "png", output)
            } catch (e: IOException) {
                println("Could not save file $outputFile")
            }

            println("Success")

        } else if (args[0] == "decrypt" && args.size == 3) {

            var img: BufferedImage? = null

            try {
                img = ImageIO.read(File(args[1]))
            } catch (e: IOException) {
                println("Could not open file")
                System.exit(1)
            }

            val width = img!!.width
            val height = img.height

            val length_bytes = ByteArray(4)
            var cur_byte = 0
            // First 4 pixels contain length of message
            for (y in 0 until height) {
                for (x in 0 until width) {
                    if (cur_byte < 4) {
                        val rgb_int = img.getRGB(x, y)
                        val color = Color(rgb_int, true)
                        var encrypted_info: Byte = 0x0
                        // first two bits in RED
                        val red = color.red
                        encrypted_info = (encrypted_info or ((red and 0x03 shl 6).toByte()))
                        // second two bits in GREEN
                        val green = color.green
                        encrypted_info = (encrypted_info or ((green and 0x03 shl 4).toByte()))
                        // third two bits in BLUE
                        val blue = color.blue
                        encrypted_info = (encrypted_info or ((blue and 0x03 shl 2).toByte()))
                        // last two bits in ALPHA
                        val alpha = color.alpha
                        encrypted_info = (encrypted_info or ((alpha and 0x03).toByte()))
                        // add byte
                        length_bytes[cur_byte] = encrypted_info
                        cur_byte++
                    } else {
                        break
                    }
                }
                if (cur_byte >= 4) {
                    break
                }
            }

            val encrypt_length = ByteBuffer.wrap(length_bytes).int

            val decrypted_bytes = ByteArray(encrypt_length)
            cur_byte = 0
            for (y in 4 until height) {
                for (x in 4 until width) {
                    if (cur_byte < encrypt_length) {
                        val rgb_int = img.getRGB(x, y)
                        val color = Color(rgb_int, true)
                        var encrypted_info: Byte = 0x0
                        // Первые 2 бита в красном
                        val red = color.red
                        encrypted_info = (encrypted_info or ((red and 0x03 shl 6).toByte()))
                        // Вторые 2 бита в зеленом
                        val green = color.green
                        encrypted_info = (encrypted_info or ((green and 0x03 shl 4).toByte()))
                        // Третьи 2 бита в синем
                        val blue = color.blue
                        encrypted_info = (encrypted_info or ((blue and 0x03 shl 2).toByte()))
                        // Последние 2 бита в альфа
                        val alpha = color.alpha
                        encrypted_info = (encrypted_info or ((alpha and 0x03).toByte()))
                        // Добавляем байт
                        decrypted_bytes[cur_byte] = encrypted_info
                        cur_byte++
                    } else {
                        break
                    }
                }
                if (cur_byte >= encrypt_length) {
                    break
                }
            }

            try {
                val decoded = String(decrypted_bytes, "UTF-8")
                println(decoded)
                val decrypted_output = FileOutputStream(args[2])
                decrypted_output.write(decrypted_bytes)
                decrypted_output.close()
            } catch (e: UnsupportedEncodingException) {
                println("Сould not decoded bytes")
                System.exit(1)
            } catch (e: IOException) {
                println("Сould not save file " + args[2])
                System.exit(1)
            }

        } else {
            System.exit(1)
        }
    }
}