{- |
ToDo:

Lazy read and write.

Problem on writing:
If the length of data is computed lazily,
then you must seek back to the file position
where the size is stored.
That is for writing of lazily generated data
we need a seekable file device.
-}
module Data.IFF
   (T(..), Chunk(..),
    ChunkId, chunkIdFromString, chunkIdToString,
    fromByteString, toByteString,
    ) where

import Control.Monad.Trans.State (StateT(StateT, runStateT), get, )
import Control.Monad.Trans.Class (lift, )
import Control.Monad (guard, liftM, liftM2, liftM4, )

import Data.Bits(Bits, (.|.), (.&.), shiftL, shiftR, bitSize)
import Data.Word(Word32)
import Data.Maybe(fromMaybe)
import Data.ByteString(ByteString)
import qualified Data.ByteString as B

import Prelude hiding (getChar)

data T = Cons
   {chunkId :: ChunkId
   ,chunk   :: Chunk
   }

data ChunkId = ChunkId Char Char Char Char
   deriving Eq

data Chunk =
     Form  {parts :: [T]
           }
   | List  {props :: [T]
           ,parts :: [T]
           }
   | Cat   {parts :: [T]
           }
   | Prop  {parts :: [T]
           }
   | Chunk {content :: ByteString
           }

chunkIdFromString :: String -> ChunkId
chunkIdFromString (c0:c1:c2:c3:[]) = ChunkId c0 c1 c2 c3
chunkIdFromString _ = error "chunkId must have four characters"

chunkIdToString :: ChunkId -> String
chunkIdToString (ChunkId c0 c1 c2 c3) = c0:c1:c2:c3:[]


formId, listId, catId, propId :: ChunkId
formId = chunkIdFromString "FORM"
listId = chunkIdFromString "LIST"
catId  = chunkIdFromString "CAT "
propId = chunkIdFromString "PROP"


instance Show T where
  show (Cons { chunkId = name, chunk = chk }) =
    "chunk "++show name++" "++show chk

instance Show ChunkId where
  show cid = show (chunkIdToString cid)

instance Show Chunk where
  show (Form { parts = p }) =
    "IFF.Form {parts="++show p++"}"
  show (List { props = ps, parts = p }) =
    "IFF.List {props"++show ps++",parts="++show p++"}"
  show (Cat { parts = p }) =
    "IFF.Cat {parts="++show p++"}"
  show (Prop { parts = p }) =
    "IFF.Prop {parts="++show p++"}"
  show (Chunk { content = cont }) =
    "IFF.Chunk { binary content, size "++show (B.length cont)++" }"


{-
b2s :: ByteString -> String
b2s = map (toEnum . fromEnum) . B.unpack
-}

isProp :: T -> Bool
isProp (Cons _ (Prop _)) = True
isProp _ = False

type Parser a = StateT ByteString Maybe a

atLeast :: Int -> ByteString -> Bool
atLeast 0 _ = True
atLeast n b = not (B.null (B.drop (pred n) b))

getBytes :: Int -> Parser ByteString
getBytes n = StateT $ \b ->
   toMaybe (atLeast n b) (B.splitAt n b)

getChar :: Parser Char
getChar = liftM (toEnum . fromIntegral . B.head) (getBytes 1)

getChunkId :: Parser ChunkId
getChunkId = liftM4 ChunkId getChar getChar getChar getChar

unpackWordBE :: (Num b, Bits b) => ByteString -> b
unpackWordBE bs =
   foldl1 (\acc byte -> shiftL acc 8 .|. byte)
          (map fromIntegral (B.unpack bs))

fromByteString :: ByteString -> Maybe T
fromByteString = runParser parser

runParser :: Parser a -> ByteString -> Maybe a
runParser parse bs =
   do (tree, rest) <- runStateT parse bs
      guard (B.null rest)
      return tree

parseMany :: Parser [T]
parseMany =
   do rest <- get
      if B.null rest
        then return []
        else liftM2 (:) parser parseMany


parser :: Parser T
parser =
   do cid        <- getChunkId
      size       <- liftM unpackWordBE (getBytes 4)
      rawContent <- getBytes size
      _ <- getBytes (mod (-size) 2) -- skip padding byte
      lift $
        fromMaybe (Just (Cons cid
                     (Chunk {content = rawContent}))) $
        lookup cid $
          (formId, runParseStruct Form rawContent) :
          (listId, runParseStruct (uncurry List . span isProp) rawContent) :
          (catId,  runParseStruct Cat  rawContent) :
          (propId, runParseStruct Prop rawContent) :
          []

runParseStruct :: ([T] -> Chunk) -> ByteString -> Maybe T
runParseStruct cons bs =
   do (format, subChunks) <- runParser parseStruct bs
      return (Cons format (cons subChunks))

parseStruct :: Parser (ChunkId, [T])
parseStruct =
   liftM2 (,) getChunkId parseMany


packWordBE :: (Bits b, Integral b) => b -> ByteString
packWordBE w =
   let n = div (bitSize w) 8
   in  B.pack $ map fromIntegral $ map (255 .&.) $ reverse $
       take n $ iterate (flip shiftR 8) w

packLengthBE :: ByteString -> ByteString
packLengthBE dat =
   packWordBE (fromIntegral (B.length dat) :: Word32)

s2b :: String -> ByteString
s2b = B.pack . map (toEnum . fromEnum)

toByteString :: T -> ByteString
toByteString (Cons {chunkId = name, chunk = chk}) =
   let bid = s2b (chunkIdToString name)
   in  case chk of
          Form subChunks -> makeStructureChunk formId bid subChunks
          List ps
               subChunks -> makeStructureChunk listId bid (ps++subChunks)
          Cat  subChunks -> makeStructureChunk catId  bid subChunks
          Prop subChunks -> makeStructureChunk propId bid subChunks
          Chunk chunkData ->
             B.concat [bid, packLengthBE chunkData, padData chunkData]

makeStructureChunk :: ChunkId -> ByteString -> [T] -> ByteString
makeStructureChunk name formatName chunks =
   let cont = B.concat (formatName : map toByteString chunks)
   in  B.concat [s2b (chunkIdToString name), packLengthBE cont, cont]

padData :: ByteString -> ByteString
padData str =
   if even (B.length str)
     then str
     else B.snoc str 0



-- from Utility
toMaybe :: Bool -> a -> Maybe a
toMaybe False _ = Nothing
toMaybe True  x = Just x