{-# LANGUAGE LambdaCase #-}

module Main where

import           Codec.Picture
import           Data.Bits                 (xor)
import qualified Data.ByteString.Lazy      as BL
import qualified Data.ByteString.Lazy.UTF8 as BLU
import           Data.Char                 (chr, ord)
import           Data.Word                 (Word8)
import           DecodeSteg
import           EncodeSteg
import           System.Environment        (getArgs)
import           System.Random             (randomRIO)

generateKey :: Int -> IO [ Int ]
generateKey n = sequence $ replicate n $ randomRIO ( 0, 65535 )

encodeText :: String -> [ Word8 ] -> IO [ Word8 ]
encodeText keypath content = do
    key <- generateKey . length $ content
    BL.writeFile keypath $ BL.pack (map fromIntegral key)
    return $ map (\( a, b ) -> a `xor` b) $ zip content (map fromIntegral key)

decodeText :: [ Word8 ] -> [ Word8 ] -> [ Word8 ]
decodeText encodedText key = map (\( a, b ) -> a `xor` b) $ zip encodedText key

main :: IO ()
main = getArgs >>= \case
    [ "encode", "file", filepath, encpath, "--img", imgpath ] -> do
        let keypath = encpath ++ ".key"
        filecontent <- BL.unpack <$> BL.readFile filepath
        encoded <- encodeText keypath filecontent
        readImage imgpath >>= \case
            Left err -> putStrLn err
            Right image -> do
                let conv = (convertRGB8 image)
                    len = (length filepath) + 1 + (length encoded)
                    bitsPerPixel = getMinBits conv len
                    message = (BL.unpack . BLU.fromString $ filepath)
                        ++ [ fromIntegral 0 ] ++ encoded
                    finalimg = (encodeImage conv bitsPerPixel message)
                case finalimg of
                    Left errorStr -> putStrLn errorStr
                    Right encrypted -> do
                        savePngImage encpath $ ImageRGB8 encrypted
                putStrLn ("Done")
    [ "encode", "file", filepath, encpath ] -> do
        let keypath = encpath ++ ".key"
        filecontent <- BL.unpack <$> BL.readFile filepath
        encoded <- encodeText keypath filecontent
        BL.writeFile encpath $ BL.pack $ encoded
    [ "encode", text, encpath, "--img", imgpath ] -> do
        let keypath = encpath ++ ".key"
        encoded <- encodeText keypath $ BL.unpack . BLU.fromString $ text
        readImage imgpath >>= \case
            Left err -> putStrLn err
            Right image -> do
                let conv = (convertRGB8 image)
                    len = 1 + (length encoded)
                    bitsPerPixel = getMinBits conv len
                    message = [ fromIntegral 0 ] ++ encoded
                    finalimg = (encodeImage conv bitsPerPixel message)
                case finalimg of
                    Left errorStr -> putStrLn errorStr
                    Right encrypted -> do
                        savePngImage encpath $ ImageRGB8 encrypted
                putStrLn ("Done")
    [ "encode", text, encpath ] -> do
        let keypath = encpath ++ ".key"
        encoded <- encodeText keypath $ BL.unpack . BLU.fromString $ text
        BL.writeFile encpath $ BL.pack $ encoded
    [ "decode", "--img", imgpath, keypath, decodepath ] -> do
        readImage imgpath >>= \case
            Left err -> putStrLn err
            Right image -> do
                encodedText <- decodeImg . convertRGB8 $ image
                key <- BL.unpack <$> BL.readFile keypath
                printOrWriteFile decodepath $ BL.pack
                    $ decodeText encodedText key
    [ "decode", encpath, keypath, decodepath ] -> do
        encodedText <- BL.unpack <$> BL.readFile encpath
        key <- BL.unpack <$> BL.readFile keypath
        printOrWriteFile decodepath $ BL.pack $ decodeText encodedText key
    _ -> help
  where
    printOrWriteFile filepath
        = if filepath == "-" then print else BL.writeFile filepath

help :: IO ()
help = do
  putStrLn "Name"
  putStrLn ""
  putStrLn "  picsg - a tool for steganographing information in a picture encoded using the Vernam cipher."
  putStrLn ""
  putStrLn "Synopsis"
  putStrLn ""
  putStrLn "  picsg [subcommand]"
  putStrLn ""
  putStrLn "  where subcommand is one of the following:"
  putStrLn ""
  putStrLn "  - encode <text> <output-path> [--img <image-path>]"
  putStrLn "  - encode file <file-path> <output-path> [--img <image-path>]"
  putStrLn "  - decode [--img] <encoded-path> <key-path> <output-path>"
  putStrLn ""
  putStrLn "Examples"
  putStrLn ""
  putStrLn "  - Encode text using the Vernam cipher to the file hello.enc and create a file hello.enc.key with a key to decode the hello.enc file."
  putStrLn ""
  putStrLn "    | picsg encode \"ABC\" hello.enc"
  putStrLn ""
  putStrLn "  - Decode the hello.enc file using key and print to the standart output"
  putStrLn ""
  putStrLn "    | picsg decode hello.enc hello.enc.key -"