module BitHelper where

import Data.Bits (FiniteBits)
import qualified Data.ByteString.Builder as Builder
import qualified Data.ByteString.Lazy as B
import qualified Data.List.Split as LS
import qualified Data.Word as W
import Table (distTable, litLenTable, llHuffmanCode)

word32ToLBS :: W.Word32 -> B.ByteString
word32ToLBS w = Builder.toLazyByteString $ Builder.word32LE w

word16ToLBS :: W.Word16 -> B.ByteString
word16ToLBS w = Builder.toLazyByteString $ Builder.word16LE w

bitsToLBS :: [Bool] -> B.ByteString
bitsToLBS bits = B.pack $ map bitsToWord8 (LS.chunksOf 8 bits)

{-
Given a list of bits (at most 8 bits),
convert to word8
-}
bitsToWord8 :: [Bool] -> W.Word8
bitsToWord8 bits
  | n > 8 = error "too long"
  | otherwise = foldl (\a f -> 2 * a + if f then 1 else 0) 0 (reverse padded)
  where
    n = length bits
    padded = replicate (8 - n) False ++ bits

wordsToBits :: [W.Word8] -> [Bool]
wordsToBits = concatMap (\f -> wordToBits (f, 8))

{-
Given the code and expected length in bits, convert into bits
of expected length
-}
wordToBits :: (Integral a) => (a, W.Word8) -> [Bool]
wordToBits (code, len) = replicate pLen False ++ bits
  where
    bits = wordToBits' code
    pLen = fromIntegral len - length bits

wordToBits' :: (Integral a) => a -> [Bool]
wordToBits' 0 = []
wordToBits' x = wordToBits' (x `div` 2) ++ [x `rem` 2 == 1]

{-
Given a ll code, we need to convert it into bits
if less than 256, convert to literal
if more than 256, convert to length + extra bits
-}
litToBits :: W.Word16 -> [Bool]
litToBits = wordToBits . llHuffmanCode

lenToBits :: W.Word16 -> [Bool]
lenToBits l = litToBits code ++ reverse (wordToBits (l - base, extra))
  where
    (code, extra, base) = litLenTable l

distToBits :: W.Word16 -> [Bool]
distToBits d = wordToBits (code, 5) ++ reverse (wordToBits (d - base, extra))
  where
    (code, extra, base) = distTable d
