module Lib
    ( 
        createAlphabetFromText,
        encryptCaesar,
        decryptCaesar,
        textToBits,
        bitsToText,
        encodePixel,
        extractShift,
        extractBitsFromImage
    ) 
    where

import Codec.Picture

import Text.Read (readMaybe)
import Data.Word (Word8)
import Data.Char (ord, chr)
import Data.Bits (testBit, shiftL, complement, (.|.), (.&.))
import qualified Data.Vector.Unboxed as VU

createAlphabetFromText :: String -> [Char]
createAlphabetFromText [] = []
createAlphabetFromText (x:xs)
    | x `elem` alphabet = alphabet
    | otherwise     = x : alphabet
  where
    alphabet = createAlphabetFromText xs

indexOf :: (Eq t) => [t] -> t -> Int
indexOf [] _ = -1
indexOf (x : xs) target
    | x == target = 0
    | otherwise = 1 + indexOf xs target

encryptCaesar :: [Char] -> Int -> String -> String
encryptCaesar alphabet shift text = map caesarChar text
    where
        caesarChar c = alphabet !! ((indexOf alphabet c + shift) `mod` length alphabet)

decryptCaesar :: [Char] -> Int -> String -> String
decryptCaesar alphabet shift = 
    encryptCaesar alphabet (alphabetLength - (shift `mod` alphabetLength)) 
    where
        alphabetLength = length alphabet

textToBits :: String -> VU.Vector Int
textToBits text = VU.fromList $ concatMap charToBits text

charToBits :: Char -> [Int]
charToBits c = [if testBit (ord c) i then 1 else 0 | i <- [7,6..0]]

-- intToBits :: Int -> [Int]
-- intToBits n = [if testBit n i then 1 else 0 | i <- [31,30..0]]

bitsToText :: VU.Vector Int -> String
bitsToText bits
    | VU.null bits = []
    | otherwise = (chr $ bitsToInt (VU.take 8 bits)) : bitsToText (VU.drop 8 bits)

bitsToInt :: VU.Vector Int -> Int
bitsToInt bits = 
    sum [bit * (2 ^ index) | (bit, index) <- zip (VU.toList bits) [len,(len - 1)..0]]
    where
        len = VU.length bits - 1

intToWord8 :: Int -> Word8
intToWord8 x = fromIntegral x

word8ToInt :: Word8 -> Int
word8ToInt x = fromIntegral x

createMask :: Int -> Int
createMask shift = shiftL (complement 0) shift .&. complement 0

encodePixel :: Int -> Image PixelRGB8 -> VU.Vector Int -> Int -> Int -> PixelRGB8
encodePixel bitsPerByte img bits x y = PixelRGB8 newR newG newB
    where
        width = imageWidth img

        index = x + y * width
        startPos = index * 3 * bitsPerByte
        pixelBits = VU.slice startPos (3 * bitsPerByte) bits

        bitsIntR = bitsToInt $ VU.slice 0 bitsPerByte pixelBits
        bitsIntG = bitsToInt $ VU.slice bitsPerByte bitsPerByte pixelBits
        bitsIntB = bitsToInt $ VU.slice (2 * bitsPerByte) bitsPerByte pixelBits

        mask = createMask bitsPerByte

        PixelRGB8 r g b = pixelAt img x y
        newR = intToWord8 $ ((word8ToInt r) .&. mask) .|. bitsIntR
        newG = intToWord8 $ ((word8ToInt g) .&. mask) .|. bitsIntG
        newB = intToWord8 $ ((word8ToInt b) .&. mask) .|. bitsIntB

extractBits :: Int -> Pixel8 -> [Int]
extractBits bitsPerByte pixelByte = 
    [ if testBit pixelByte i then 1 else 0 | i <- [bitsPerByte-1, bitsPerByte-2..0] ] 

extractBitsFromPixel :: Int -> PixelRGB8 -> [Int]
extractBitsFromPixel bitsPerByte (PixelRGB8 r g b) =
    let bitsR = extractBits bitsPerByte r
        bitsG = extractBits bitsPerByte g
        bitsB = extractBits bitsPerByte b
    in bitsR ++ bitsG ++ bitsB

extractBitsFromImage :: Int -> Image PixelRGB8 -> [Int]
extractBitsFromImage bitsPerByte img = 
    let width = imageWidth img
        height = imageHeight img
        pixels = [ pixelAt img x y | y <- [0..height - 1], x <- [0..width - 1] ]
    in concatMap (extractBitsFromPixel bitsPerByte) pixels

extractShift :: String -> Maybe Int
extractShift path = 
    let shift = takeWhile (`elem` ['0'..'9']) (reverse $ takeWhile (/= '_') (reverse path))
    in readMaybe shift