Writer.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Write Vorbis/Opus tags into an OGG container.
--
-- The comment packet is rebuilt and repaginated, which can change the
-- number of header pages; every following page therefore gets its
-- sequence number rewritten and its CRC recomputed. Audio data itself is
-- copied verbatim.
module Monatone.OGG.Writer
  ( writeOGGMetadata
  , WriteError(..)
  , Writer
  ) where

import Control.Exception (catch, IOException)
import Control.Monad (unless)
import Control.Monad.Except (ExceptT, throwError, runExceptT)
import Control.Monad.IO.Class (liftIO)
import Data.Binary.Put
import Data.Bits (shiftL, shiftR, testBit, xor, (.&.), (.|.))
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base64 as B64
import qualified Data.ByteString.Lazy as L
import qualified Data.HashMap.Strict as HM
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import Data.Word
import System.IO hiding (withBinaryFile)
import System.OsPath
import System.File.OsPath (withBinaryFile)
import System.IO.Temp (withSystemTempFile)

import Monatone.Metadata
import Monatone.FLAC.Writer (generateVorbisComments, renderPictureData)

-- Re-define WriteError and Writer locally to avoid circular imports
data WriteError
  = WriteIOError Text
  | UnsupportedWriteFormat AudioFormat
  | InvalidMetadata Text
  | CorruptedWrite Text
  deriving (Show, Eq)

type Writer = ExceptT WriteError IO

-- | Write metadata to an OGG (Vorbis or Opus) file
writeOGGMetadata :: Metadata -> AlbumArtUpdate -> OsPath -> Writer ()
writeOGGMetadata metadata artUpdate filePath = do
  result <- liftIO $ tryIO $ do
    withSystemTempFile "monatone-ogg.tmp" $ \tmpPath tmpHandle -> do
      hClose tmpHandle
      tmpOsPath <- encodeFS tmpPath
      rewriteResult <- withBinaryFile filePath ReadMode $ \srcHandle ->
        withBinaryFile tmpOsPath WriteMode $ \dstHandle ->
          runExceptT $ rewriteOGGFile srcHandle dstHandle metadata artUpdate
      case rewriteResult of
        Left err -> return $ Left err
        Right () -> do
          copyFileContents tmpOsPath filePath
          return $ Right ()

  case result of
    Left (e :: IOException) -> throwError $ WriteIOError $ T.pack $ show e
    Right (Left err) -> throwError err
    Right (Right ()) -> return ()
  where
    tryIO :: IO a -> IO (Either IOException a)
    tryIO action = catch (Right <$> action) (return . Left)

-- | Copy file contents
copyFileContents :: OsPath -> OsPath -> IO ()
copyFileContents src dst =
  withBinaryFile src ReadMode $ \srcHandle ->
    withBinaryFile dst WriteMode $ \dstHandle ->
      let loop = do
            chunk <- BS.hGet srcHandle 65536
            unless (BS.null chunk) $ do
              BS.hPut dstHandle chunk
              loop
      in loop

-- | One raw OGG page, split into its parts
data RawPage = RawPage
  { rpHeader :: ByteString    -- 27 bytes
  , rpSegTable :: ByteString
  , rpData :: ByteString
  }

rpContinued :: RawPage -> Bool
rpContinued page = BS.index (rpHeader page) 5 .&. 0x01 /= 0

rpSerial :: RawPage -> ByteString
rpSerial page = BS.take 4 $ BS.drop 14 $ rpHeader page

-- | Segments with an ends-packet marker (lacing value < 255)
rpSegments :: RawPage -> [(ByteString, Bool)]
rpSegments page = go (map fromIntegral (BS.unpack (rpSegTable page))) (rpData page)
  where
    go [] _ = []
    go (lacing:rest) bs =
      let (segment, unread) = BS.splitAt lacing bs
      in (segment, lacing < 255) : go rest unread

readRawPage :: Handle -> IO (Maybe RawPage)
readRawPage handle = do
  header <- BS.hGet handle 27
  if BS.length header < 27 || BS.take 4 header /= "OggS"
    then return Nothing
    else do
      let numSegments = fromIntegral $ BS.index header 26
      segTable <- BS.hGet handle numSegments
      let dataSize = sum $ map fromIntegral $ BS.unpack segTable
      pageData <- BS.hGet handle dataSize
      if BS.length segTable < numSegments || BS.length pageData < dataSize
        then return Nothing
        else return $ Just $ RawPage header segTable pageData

-- | Rewrite the OGG stream with a new comment packet
rewriteOGGFile :: Handle -> Handle -> Metadata -> AlbumArtUpdate -> Writer ()
rewriteOGGFile srcHandle dstHandle metadata artUpdate = do
  -- Page 0 carries the identification packet alone; copy it verbatim
  maybeFirst <- liftIO $ readRawPage srcHandle
  firstPage <- case maybeFirst of
    Nothing -> throwError $ CorruptedWrite "Not an OGG file"
    Just page -> return page
  let identPacket = BS.concat (map fst (rpSegments firstPage))

  -- Vorbis has a setup packet after the comment; Opus does not
  (commentMagic, trailingHeaderPackets) <-
    if "\x01vorbis" `BS.isPrefixOf` identPacket then return ("\x03vorbis", 1)
    else if "OpusHead" `BS.isPrefixOf` identPacket then return ("OpusTags", 0)
    else throwError $ UnsupportedWriteFormat OGG

  -- Recompute the CRC rather than trusting the source's
  let firstRaw = rpHeader firstPage <> rpSegTable firstPage <> rpData firstPage
  liftIO $ BS.hPut dstHandle $ patchCrc $
    BS.take 22 firstRaw <> BS.replicate 4 0 <> BS.drop 26 firstRaw

  -- Collect the remaining header packets (comment [+ setup])
  headerPackets <- collectHeaderPackets srcHandle (1 + trailingHeaderPackets)
  (oldComment, setupPackets) <- case headerPackets of
    (c:rest) -> return (c, rest)
    [] -> throwError $ CorruptedWrite "Missing comment header packet"

  -- Build the new comment packet, keeping the original vendor string
  let newComment = buildCommentPacket commentMagic (extractVendor commentMagic oldComment)
        metadata artUpdate

  -- Repaginate the header packets, then stream the audio pages through
  -- with renumbered sequence numbers and recomputed CRCs
  nextSeqno <- liftIO $ writeHeaderPages dstHandle (rpSerial firstPage) (newComment : setupPackets)
  liftIO $ rewriteAudioPages srcHandle dstHandle nextSeqno

-- | Read pages after the first, reassembling exactly n packets. The
-- header packets must end on a page boundary (they do in spec-conformant
-- streams); anything else would interleave audio into the header pages.
collectHeaderPackets :: Handle -> Int -> Writer [ByteString]
collectHeaderPackets handle n = go [] []
  where
    go :: [ByteString] -> [ByteString] -> Writer [ByteString]
    go completed partial
      | length completed >= n =
          if null partial && length completed == n
            then return $ reverse completed
            else throwError $ CorruptedWrite "Header packets do not end on a page boundary"
      | otherwise = do
          maybePage <- liftIO $ readRawPage handle
          case maybePage of
            Nothing -> throwError $ CorruptedWrite "Truncated OGG header pages"
            Just page -> do
              let partial0 = if rpContinued page then partial else []
                  (completed', partial') = foldl step (completed, partial0) (rpSegments page)
              go completed' partial'

    step (completed, partial) (segmentData, endsPacket)
      | endsPacket = (BS.concat (reverse (segmentData : partial)) : completed, [])
      | otherwise = (completed, segmentData : partial)

-- | The vendor string sits right after the packet magic
extractVendor :: ByteString -> ByteString -> ByteString
extractVendor magic packet =
  let body = BS.drop (BS.length magic) packet
      vendorLen = word32leAt body
      vendor = BS.take vendorLen (BS.drop 4 body)
  in if BS.length body >= 4 && BS.length vendor == vendorLen then vendor else ""
  where
    word32leAt bs
      | BS.length bs < 4 = 0
      | otherwise = sum [fromIntegral (BS.index bs i) `shiftL` (8 * i) | i <- [0 .. 3]]

-- | Build the comment packet: magic, vendor, comments (art included as a
-- base64 METADATA_BLOCK_PICTURE comment), and Vorbis's framing bit
buildCommentPacket :: ByteString -> ByteString -> Metadata -> AlbumArtUpdate -> ByteString
buildCommentPacket magic vendor metadata artUpdate =
  let comments = generateVorbisComments metadata ++ artComments
      renderComment (key, value) =
        let bytes = TE.encodeUtf8 (key <> "=" <> value)
        in w32le (BS.length bytes) <> bytes
      framing = if magic == "\x03vorbis" then BS.singleton 0x01 else BS.empty
  in BS.concat $
       [magic, w32le (BS.length vendor), vendor, w32le (length comments)]
       ++ map renderComment comments
       ++ [framing]
  where
    artComments :: [(Text, Text)]
    artComments = case artUpdate of
      KeepExistingArt ->
        [ ("METADATA_BLOCK_PICTURE", value)
        | value <- HM.lookupDefault [] "METADATA_BLOCK_PICTURE" (rawTags metadata)
        ]
      SetAlbumArt art ->
        [("METADATA_BLOCK_PICTURE", TE.decodeUtf8 (B64.encode (L.toStrict (renderPictureData art))))]
      RemoveAlbumArt -> []

-- | Paginate packets into header pages (granule 0) starting at sequence
-- number 1; returns the next free sequence number
writeHeaderPages :: Handle -> ByteString -> [ByteString] -> IO Word32
writeHeaderPages handle serial packets = go 1 False (pagesOf (concatMap packetSegments packets))
  where
    go seqno _ [] = return seqno
    go seqno continued (pageSegments : rest) = do
      BS.hPut handle $ renderPage continued seqno pageSegments
      let nextContinued = case reverse pageSegments of
            (lastSegment:_) -> BS.length lastSegment == 255
            [] -> False
      go (seqno + 1) nextContinued rest

    -- Segments of one packet: 255-byte chunks with a final short one
    -- (empty when the length is a multiple of 255)
    packetSegments packet =
      let (chunk, rest) = BS.splitAt 255 packet
      in if BS.length chunk == 255
           then chunk : packetSegments rest
           else [chunk]

    pagesOf [] = []
    pagesOf segments = let (page, rest) = splitAt 255 segments in page : pagesOf rest

    renderPage continued seqno segments =
      let flags = if continued then 0x01 else 0x00
          headerNoCrc = BS.concat
            [ "OggS"
            , BS.singleton 0                                   -- version
            , BS.singleton flags
            , BS.replicate 8 0                                 -- granule position
            , serial
            , w32le (fromIntegral seqno)
            , BS.replicate 4 0                                 -- CRC placeholder
            , BS.singleton (fromIntegral (length segments))
            , BS.pack (map (fromIntegral . BS.length) segments)
            ]
          page = headerNoCrc <> BS.concat segments
      in patchCrc page

-- | Copy the remaining pages, renumbering sequence numbers (the header
-- page count may have changed) and recomputing CRCs. Trailing bytes that
-- do not parse as a page are copied verbatim.
rewriteAudioPages :: Handle -> Handle -> Word32 -> IO ()
rewriteAudioPages srcHandle dstHandle seqno = do
  pos <- hTell srcHandle
  maybePage <- readRawPage srcHandle
  case maybePage of
    Just page -> do
      let raw = rpHeader page <> rpSegTable page <> rpData page
          renumbered = BS.take 18 raw <> w32le (fromIntegral seqno)
            <> BS.replicate 4 0 <> BS.drop 26 raw
      BS.hPut dstHandle $ patchCrc renumbered
      rewriteAudioPages srcHandle dstHandle (seqno + 1)
    Nothing -> do
      -- copy whatever is left verbatim
      hSeek srcHandle AbsoluteSeek pos
      let loop = do
            chunk <- BS.hGet srcHandle 65536
            unless (BS.null chunk) $ do
              BS.hPut dstHandle chunk
              loop
      loop

-- | Compute the page CRC (bytes 22-25 must be zero) and patch it in
patchCrc :: ByteString -> ByteString
patchCrc page =
  BS.take 22 page <> w32le (fromIntegral (oggCrc page)) <> BS.drop 26 page

-- | The OGG page checksum: CRC-32 with polynomial 0x04c11db7, no
-- reflection, zero initial value and no final xor
oggCrc :: ByteString -> Word32
oggCrc = BS.foldl' step 0
  where
    step crc byte = go (8 :: Int) (crc `xor` (fromIntegral byte `shiftL` 24))
    go 0 crc = crc
    go n crc = go (n - 1) $
      if testBit crc 31 then (crc `shiftL` 1) `xor` 0x04c11db7 else crc `shiftL` 1

w32le :: Int -> ByteString
w32le n = BS.pack [fromIntegral (n `shiftR` s) | s <- [0, 8, 16, 24]]