Vorbis.hs

module Codec.Audio.Vorbis (readAudioInfoAndComment, AudioInfo(..)) where

import qualified Codec.Binary.UTF8.String as UTF8
--import Control.Monad (when)
import Control.Monad.Error -- For Either monad instance
--import qualified Data.ByteString.Lazy as B
import qualified Data.ByteString as B
--import qualified Data.ByteString.Lazy.Char8 as Char8
import qualified Data.ByteString.Char8 as Char8
import Data.Binary.Strict.Get
import Data.Char (toLower)
import Data.Word

data AudioInfo = AudioInfo {channels :: Word8, sampleRate :: Word32, maxBitrate :: Word32,
                            nominalBitrate :: Word32, minBitrate :: Word32,
                            playtime :: Maybe Double}
  deriving Show

readAudioInfoAndComment :: FilePath -> IO (Either String (AudioInfo, [(String, String)]))
readAudioInfoAndComment fname = do
  content <- B.readFile fname
  return $ case readVorbisAudioInfo content of
             Left e -> Left e
             Right audioInfo -> case readVorbisCommentFromOgg content of
                                  Left e -> Left e
                                  Right comment -> Right (audioInfo, comment)

readVorbisCommentFromOgg :: Char8.ByteString -> Either String [(String, String)]
readVorbisCommentFromOgg fileContent = case runGet readVorbisCommentOggPage fileContent of
    (Left e, _) -> Left e
    (Right p2, _) ->
      -- XXX: This won't work right if packet is spanned across multiple pages...
      case runGet (readVorbisComment False) commentData of
        (Left e, _) -> Left e
        (Right comment, _) -> Right comment
      where commentData = B.drop 7 $ packetData p2 -- 'drop 7' strips off "\x03vorbis"

readVorbisCommentOggPage = do
  p1 <- readOggPage
  p2 <- readOggPage
  when (not $ isVorbisPage p1) $ fail "First page is not a vorbis page"
  when (serialNum p1 /= serialNum p2) $ fail "First two pages do not have same serial number"
  --when (not $ isComplete p2) $ fail "Second vorbis page is not a complete packet"
  return p2

{-
readVorbisCommentFromOgg :: Get [(String, String)]
readVorbisCommentFromOgg = do
  p1 <- readOggPage
  p2 <- readOggPage
  when (not $ isVorbisPage p1) $ fail "First page is not a vorbis page"
  when (serialNum p1 /= serialNum p2) $ fail "First two pages do not have same serial number"
  when (not $ isComplete p2) $ fail "Second vorbis page is not a complete packet"
  -- XXX: This won't work right if packet is spanned across multiple pages...
  let commentData = B.drop 7 $ packetData p2 -- 'drop 7' strips off "\x03vorbis"
  let (comment, _) = runGet (readVorbisComment False) commentData
  return comment
-}



{-
readVorbisCommentFromOgg = do
  allPages <- readOggPages
  -- XXX: This could be wrong, if the first page is not a "\x01vorbis" one...
  let serial = serialNum $ head allPages
  let relevantPages = filter (\p -> serialNum p == serial) allPages
  let secondVorbisPage = head $ tail relevantPages
  -- XXX: This won't work right if packet is spanned across multiple pages...
  let commentData = B.drop 7 $ packetData secondVorbisPage -- 'drop 7' strips off "\x03vorbis"
  return $ runGet (readVorbisComment False) commentData
-}

readVorbisComment framingXXX = do
  vendorLength <- getWord32le
  vendor <- getAByteString $ fromIntegral vendorLength -- XXX: Need to decode from utf-8
  numItems <- getWord32le
  readEntries numItems
  --if framing and not ord(fileobj.read(1)) & 0x01:
  --   raise VorbisUnsetFrameError("framing bit was unset")
  where
    readEntries 0 = return []
    readEntries numItems = do
      len <- getWord32le
      rawStr <- getAByteString $ fromIntegral len -- XXX: Need to decode from utf-8
      let ustr = UTF8.decode $ B.unpack rawStr
      restOfEntries <- readEntries (numItems - 1)
      return $ (downcase $ takeWhile (/= '=') ustr, tail $ dropWhile (/= '=') ustr) : restOfEntries
      -- XXX: tag, value = string.split('=', 1)
      -- XXX: tag = tag.encode('ascii', errors)
    downcase = map toLower

--readVorbisAudioInfo :: Get (Word8, Word32, Word32, Word32, Word32)
readVorbisAudioInfo :: Char8.ByteString -> Either String AudioInfo
readVorbisAudioInfo fileContent = do
  firstPage <- readOggPageOrBust fileContent "Failed to parse first Ogg page"
  let reversed = B.reverse fileContent
  let oggs = Char8.pack "OggS"
  let lastPageData = B.concat [oggs, B.reverse $ fst $ B.breakSubstring (B.reverse oggs) reversed]
  {- lastPageData <- case findSubstring (B.reverse $ Char8.pack "OggS") reversed of
                    Nothing -> Left "Could not find any Ogg pages in file"
                    --Just pos -> Left $ (show $ snd $ B.breakSubstring (B.reverse $ Char8.pack "OggS") reversed) 
                    Just pos -> Right $ B.reverse $ B.take (B.length fileContent - pos) reversed -}
  lastPage <- readOggPageOrBust lastPageData "Failed to parse last Ogg page"
  let (maybeInfo, _) = runGet (readFromVorbisPacket (granulePosition lastPage)) (packetData firstPage)
  maybeInfo

  --pages <- readOggPages
  --let vorbisPages = filter isVorbisPage pages
  --when (length vorbisPages < 1) $ fail "No Vorbis pages found in file"
  --when (length vorbisPages > 1) $ fail "More than one Vorbis page found!"
  
  --vorbisPage <- readOggPage
  --return $ case runGet readFromVorbisPacket (packetData vorbisPage) of
  --  (Left e, _) -> e
  --  (Right audioInfo, _) -> audioInfo

  where
    readFromVorbisPacket totalNumSamples = do
      skip 11
      channels <- getWord8
      sampleRate <- getWord32le
      maxBitrate <- getWord32le
      nominalBitrate <- getWord32le
      minBitrate <- getWord32le
      let playtime = (toDouble totalNumSamples) / (toDouble sampleRate) :: Double
      return $ AudioInfo channels sampleRate maxBitrate nominalBitrate minBitrate (Just playtime)
    toDouble :: Integral a => a -> Double
    toDouble = fromInteger . toInteger

readOggPageOrBust :: Char8.ByteString -> String -> Either String OggPage
readOggPageOrBust content err = case maybePage of
                                  Left e -> Left $ err ++ ": " ++ e
                                  Right p -> Right p
  where (maybePage, _) = runGet readOggPage content

{- findSubstring s l =
    if B.null s then Just 0
              else case B.breakSubstring s l of
                     (x,y) | B.null y  -> Nothing
                           | otherwise -> Just (B.length x) -}

--serialNum (_, _, _, _, sn, _, _, _, _) = sn
--packetData (_, _, _, _, _, _, _, _, d) = d

{- readOggVorbisPages = do
  p <- readOggPage
  if isVorbisPage p
    then do
      otherPages <- readPagesWithSerialNum $ serialNum p
      return $ p : otherPages
    else fail "first page wasn't a vorbis page"
      --readOggVorbisPages
  where
    readPagesWithSerialNum sn = do
      p <- readOggPage
      if serialNum p == sn
        then do
          others <- readPagesWithSerialNum sn
          return $ p : others
        else return [] -}

isVorbisPage page = B.take 7 (packetData page) == Char8.pack "\x01vorbis"

{- readOggPages = do
  p <- readOggPage
  rest <- if lastPage p then (return []) else readOggPages
  return $ p : rest
  where lastPage (_, _, flags, _, _, _, _, _, _) = flags == 4 -}

data OggPage = OggPage { version :: Word8, typeFlags :: Word8, granulePosition :: Word64,
                         serialNum :: Word32, pageNum :: Word32, checkSum :: Word32,
                         segmentsCount :: Word8, packetData :: B.ByteString
                         {- isComplete :: Bool -} }

--readOggPage :: Get (ByteString, Word8, xxx)
readOggPage = do
  oggs <- getAByteString 4
  version <- getWord8
  typeFlags <- getWord8
  granulePosition <- getWord64le
  serialNum <- getWord32le
  pageNum <- getWord32le
  checkSum <- getWord32le
  segmentsCount <- getWord8
  lacingBytes <- getAByteString $ fromIntegral $ toInteger segmentsCount
  when (B.length lacingBytes /= (fromIntegral $ toInteger segmentsCount)) $
    fail "Unable to read sufficient \"lacing\" bytes"
  --let isCompletePacket = B.last lacingBytes < 255
  packetData <- getAByteString $ fromIntegral $ byteSum lacingBytes
  return $ OggPage version typeFlags granulePosition serialNum pageNum checkSum segmentsCount
                   packetData {-(isCompletePacket)-}

byteSum :: B.ByteString -> Integer
byteSum bs = B.foldl (\total byte -> total + (toInteger byte)) 0 bs

getAByteString = getByteString