File.hs

{-# LANGUAGE RecordWildCards #-}

module Telescope.Asdf.Encoding.File where

import Control.Monad (replicateM_)
import Data.Binary.Get
import Data.Binary.Put
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as BC
import Data.ByteString.Lazy qualified as BL
import Data.String (IsString)
import Data.Word
import Effectful
import Effectful.Error.Static
import Effectful.Fail
import Effectful.NonDet
import Effectful.State.Static.Local
import Telescope.Asdf.Error
import Telescope.Asdf.Node (Tree)


-- | Split an encoded 'ByteString' into a 'Tree', '[Encoded Block]' and 'Encoded Index'
splitAsdfFile :: (Error AsdfError :> es) => ByteString -> Eff es AsdfFile
splitAsdfFile dat = do
  res <- runFail $ evalState dat parseAsdfFile
  case res of
    Left e -> throwError $ ParseError e
    Right a -> pure a


-- | Lightweight first-pass parsing using 'State ByteString' instead of Parsec or similar
parseAsdfFile :: (State ByteString :> es, Fail :> es) => Eff es AsdfFile
parseAsdfFile = do
  -- if it's empty, then give an error
  tree <- onEmpty "tree" parseTree
  blocks <- onEmpty "blocks" parseBlocks
  index <- parseIndex
  pure $ AsdfFile{tree, blocks, index}
 where
  onEmpty ex eff = do
    res <- runNonDet OnEmptyKeep eff
    case res of
      Left _ -> do
        inp <- get @ByteString
        fail $ "Expected " ++ ex ++ " at " ++ show (BS.take 100 inp)
      Right a -> pure a


parseTree :: (State ByteString :> es, NonDet :> es) => Eff es (Encoded Tree)
parseTree = do
  t <- parseUntil blockMagicToken <|> parseUntil blockIndexHeader <|> remainingBytes
  pure $ Encoded t


parseIndex :: (State ByteString :> es) => Eff es (Encoded Index)
parseIndex =
  Encoded <$> get


parseBlocks :: (State ByteString :> es, NonDet :> es) => Eff es [Encoded Block]
parseBlocks = many parseBlock


parseBlock :: (State ByteString :> es, NonDet :> es) => Eff es (Encoded Block)
parseBlock = do
  exactly blockMagicToken
  b <- parseUntil blockMagicToken <|> parseUntil blockIndexHeader <|> nonEmpty
  case b of
    "" -> empty
    _ -> pure $ Encoded $ blockMagicToken <> b


exactly :: (State ByteString :> es, NonDet :> es) => ByteString -> Eff es ()
exactly val = do
  inp <- get @ByteString
  if val `BS.isPrefixOf` inp
    then do
      put $ BS.drop (BS.length val) inp
      pure ()
    else empty


parseUntil :: (State ByteString :> es, NonDet :> es) => ByteString -> Eff es ByteString
parseUntil val = do
  inp <- get @ByteString
  let (before, rest) = BS.breakSubstring val inp
  case rest of
    "" -> empty
    _ -> do
      put rest
      pure before


nonEmpty :: (NonDet :> es, State ByteString :> es) => Eff es ByteString
nonEmpty = do
  b <- remainingBytes
  case b of
    "" -> empty
    _ -> pure b


remainingBytes :: (NonDet :> es, State ByteString :> es) => Eff es ByteString
remainingBytes = do
  inp <- get
  put @ByteString ""
  pure inp


breakIndex :: (State ByteString :> es) => Eff es (Encoded Index)
breakIndex = Encoded <$> get


concatAsdfFile :: AsdfFile -> ByteString
concatAsdfFile a =
  mconcat [a.tree.bytes, blocks a.blocks, index a.index]
 where
  blocks :: [Encoded Block] -> ByteString
  blocks ebks = mconcat $ fmap (.bytes) ebks
  -- case mconcat $ fmap (.bytes) ebks of
  --   "" -> ""
  --   s -> s <> "\n"
  index ix = ix.bytes


encodeTree :: ByteString -> Encoded Tree
encodeTree tr =
  Encoded $ BS.intercalate "\n" (headers <> formatDoc tr) <> "\n" -- has a trailing newline
 where
  formatDoc doc = ["--- " <> doc, "..."]
  headers = ["#ASDF 1.0.0", "#ASDF_STANDARD 1.5.0", "%YAML 1.1", tagDirective]
  tagDirective = "%TAG ! tag:stsci.edu:asdf/"


encodeBlocks :: [BlockData] -> [Encoded Block]
encodeBlocks = fmap encodeBlock


encodeBlock :: BlockData -> Encoded Block
encodeBlock b =
  Encoded $ BL.toStrict $ runPut $ putBlock b


encodeIndex :: BlockIndex -> Encoded Index
encodeIndex (BlockIndex ns) =
  Encoded $ BS.intercalate "\n" $ ["#ASDF BLOCK INDEX", "%YAML 1.1", "---"] <> fmap indexEntry ns <> ["..."]
 where
  indexEntry n = "- " <> BC.pack (show n)


data Index
data Block


-- | values that have been encoded to the file format: padding, etc
newtype Encoded a = Encoded {bytes :: ByteString}
  deriving (Show, Eq)
  deriving newtype (IsString)


-- | Decompressed block data
newtype BlockData = BlockData {bytes :: ByteString}
  deriving (Eq)


newtype BlockSource = BlockSource Int
  deriving (Eq)


instance Show BlockData where
  show (BlockData bs) = "BlockData " <> show (BS.length bs)


data BlockHeader = BlockHeader
  { headerSize :: Word16
  , flags :: Word32
  , compression :: Compression
  , allocatedSize :: Word64
  , usedSize :: Word64
  , dataSize :: Word64
  , checksum :: Checksum
  }
  deriving (Show, Eq)


newtype BlockIndex = BlockIndex [Int]
  deriving (Eq)


data Compression
  = NoCompression -- "\0\0\0\0"
  | ZLib -- "zlib"
  | BZip2 -- "bzp2"
  deriving (Show, Eq)


newtype Checksum = Checksum ByteString
  deriving (Show, Eq)


noChecksum :: Checksum
noChecksum = Checksum $ BC.replicate 16 '0'


data AsdfFile = AsdfFile
  { tree :: Encoded Tree
  , blocks :: [Encoded Block]
  , index :: Encoded Index
  }
  deriving (Show, Eq)


getBlock :: Get BlockData
getBlock = do
  h <- getBlockHeader
  getBlockData h


-- | Skip along blocks and create a list of all of them
getBlocks :: Get [BlockData]
getBlocks = do
  isBlock <- checkMagicToken
  if not isBlock
    then pure []
    else do
      b <- getBlock
      bs <- getBlocks
      pure (b : bs)


getBlockHeader :: Get BlockHeader
getBlockHeader = do
  expectMagicToken
  headerSize <- label "header_size" getWord16be
  start <- bytesRead

  -- the remainder is inside the headerSize
  flags <- label "flags" getWord32be -- 4
  compression <- label "compression" getCompression -- 4
  allocatedSize <- label "allocated_size" getWord64be -- 8
  usedSize <- label "used_size" getWord64be -- 8
  dataSize <- label "data_size" getWord64be -- 8
  checksum <- label "checksum" getChecksum -- 2
  end <- bytesRead

  -- skip until the end of headerSize
  let usedHead = fromIntegral $ end - start
  skip $ fromIntegral headerSize - usedHead

  pure $ BlockHeader{..}
 where
  getCompression = do
    val <- getByteString 4
    case val of
      "\0\0\0\0" -> pure NoCompression
      -- "zlib" -> pure ZLib
      -- "bzp2" -> pure BZip2
      _ -> fail $ "BlockHeader compression invalid, found " <> show val

  getChecksum = Checksum <$> getByteString 16

  expectMagicToken = do
    m <- getMagicToken
    case m of
      Right a -> pure a
      Left str -> fail $ "BlockHeader magic token invalid: " ++ show str


putBlockHeader :: BlockHeader -> Put
putBlockHeader h = do
  putByteString blockMagicToken
  putWord16be h.headerSize
  size <- putHeaderContent
  let emptyBytes = fromIntegral h.headerSize - size :: Int
  replicateM_ emptyBytes $ putWord8 0x0
 where
  putCompression _ = putByteString "\0\0\0\0"
  putChecksum (Checksum cs) = putByteString cs
  putHeaderContent = do
    let bs = runPut $ do
          putWord32be 0 -- flags
          putCompression NoCompression
          putWord64be h.allocatedSize
          putWord64be h.usedSize
          putWord64be h.dataSize
          putChecksum h.checksum
    putLazyByteString bs
    pure $ fromIntegral $ BL.length bs


blockIndex :: Encoded Tree -> [Encoded Block] -> BlockIndex
blockIndex (Encoded bytes) ebs =
  let ns = scanl go (BS.length bytes) ebs
   in BlockIndex $ take (length ns - 1) ns
 where
  go :: Int -> Encoded Block -> Int
  go n eb = n + BS.length eb.bytes


putBlock :: BlockData -> Put
putBlock bd@(BlockData bs) = do
  putBlockHeader $ blockHeader bd
  putByteString bs


blockHeader :: BlockData -> BlockHeader
blockHeader (BlockData bs) =
  let bytes = fromIntegral $ BS.length bs
   in BlockHeader
        { headerSize = 48 -- minimum allowed size. Our encoding uses fewer bytes than this
        , compression = NoCompression
        , allocatedSize = bytes
        , usedSize = bytes
        , dataSize = bytes
        , checksum = noChecksum
        , flags = 0
        }


getBlockData :: BlockHeader -> Get BlockData
getBlockData h = do
  -- LATER: handle compression
  bytes <- getByteString $ fromIntegral h.usedSize
  _empty <- getByteString $ fromIntegral $ h.allocatedSize - h.usedSize
  pure $ BlockData bytes


blockMagicToken :: ByteString
blockMagicToken = BS.pack [0xd3, 0x42, 0x4c, 0x4b]


blockIndexHeader :: ByteString
blockIndexHeader = "#ASDF BLOCK INDEX"


-- | Verify that the magic token is next without consuming any input
checkMagicToken :: Get Bool
checkMagicToken = do
  emp <- isEmpty
  if emp
    then pure False
    else do
      eb <- lookAhead getMagicToken
      pure $ either (const False) (const True) eb


-- | consume the magic token if available
getMagicToken :: Get (Either ByteString ())
getMagicToken = do
  -- this still fails if it is empty
  str <- getByteString 4
  pure $
    if str == blockMagicToken
      then Right ()
      else Left str