{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
module Block
  ( doSeqCompress,
    doParCompress,
    combineCrc,
    combine,
    splitUp,
  )
where

import BitHelper (word16ToLBS, word32ToLBS, wordsToBits, bitsToLBS)
import qualified Control.Parallel.Strategies as S
import qualified Data.Bits as Bits
import qualified Data.ByteString.Lazy as B
import qualified Data.List as L
import qualified Data.Word as W
import Deflate (deflate)
import LZ77 (lz77Compress)
import MyCRC32 (CRC32 (crc32), crc32Combine)
import qualified GHC.Generics as S

-- Three types of data blocks
-- Uncompressed, using static tree, using dynamic tree
data BlockType
  = Uncompressed
  | Static
  | Dynamic
  deriving (Eq)

data InputBlock = InputBlock
  { iType :: BlockType,
    iLast :: Bool,
    iData :: B.ByteString
  }

data OutputBlock = OutputBlock
  { oData :: [Bool],
    oCrc :: W.Word32,
    oLen :: W.Word32
  } deriving (S.NFData, S.Generic)

splitUp :: B.ByteString -> [InputBlock]
splitUp s
  | B.length t == 0 = [InputBlock Static True h]
  | otherwise = InputBlock Static False h : splitUp t
  where
    (h, t) = B.splitAt 4096 s

combine :: [OutputBlock] -> B.ByteString
combine blocks = B.concat [bitsToLBS contents, fcrc, flen]
  where
    contents = concatMap oData blocks
    (len, pairs) =
      L.mapAccumL (\a x -> (a + oLen x, (oCrc x, oLen x))) 0 blocks
    (tcrc, tlen) = foldl1 combinePair pairs
    flen = word32ToLBS tlen
    fcrc = word32ToLBS tcrc

combinePair ::
  (W.Word32, W.Word32) -> (W.Word32, W.Word32) -> (W.Word32, W.Word32)
combinePair p1 p2 = (newCrc, newLen)
  where
    (crc1, len1) = p1
    (crc2, len2) = p2
    newCrc = combineCrc crc1 p2
    newLen = combineLen len1 len2

-- TODO: Do this math in haskell
-- CRC32(AB) = CRC32(A0) ^ CRC32(0B) = CRC32(A0) ^ CRC32(B)
-- https://stackoverflow.com/questions/23122312/crc-calculation-of-a-mostly-static-data-stream
combineCrc :: W.Word32 -> (W.Word32, W.Word32) -> W.Word32
combineCrc crc1 pair = crc32Combine crc1 crc2 (fromIntegral len2)
  where
    (crc2, len2) = pair

combineLen :: W.Word32 -> W.Word32 -> W.Word32
combineLen len1 len2 = len1 + len2

doParCompress :: B.ByteString -> B.ByteString
doParCompress s = combine blocks
  where
    chunks = splitUp s
    blocks = map doCompress' chunks `S.using` S.parList S.rdeepseq

doSeqCompress :: B.ByteString -> B.ByteString
doSeqCompress s = combine $ map doCompress' $ splitUp s

doCompress' :: InputBlock -> OutputBlock
doCompress' i
  | t == Uncompressed = getUncompressed i
  | t == Static = getStatic i
  | t == Dynamic = getDynamic i
  | otherwise = error "invalid block type"
  where
    t = iType i

getUncompressed :: InputBlock -> OutputBlock
getUncompressed iblock = OutputBlock hc crc (fromIntegral l)
  where
    content = iData iblock
    crc = S.runEval $ S.rseq $ crc32 content
    header = [iLast iblock, False, False] ++ replicate 5 False
    l = (fromIntegral $ B.length content) :: W.Word16
    cl = Bits.complement l
    hc = reverse header ++ comp 
    comp = wordsToBits (B.unpack $ B.concat [word16ToLBS l, word16ToLBS cl, content])


getStatic :: InputBlock -> OutputBlock
getStatic iblock = OutputBlock (bits ++ output) crc (fromIntegral l)
  where
    content = iData iblock
    crc = crc32 content
    isFinal = iLast iblock
    l = (fromIntegral $ B.length content) :: W.Word16
    bits = isFinal : [True, False]
    output = deflate content

getDynamic :: InputBlock -> OutputBlock
getDynamic _ = error "unimplemented"
