Writer.hs

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

module Monatone.FLAC.Writer
  ( writeFLACMetadata
  , WriteError(..)
  , Writer
    -- * Shared with the OGG writer (Vorbis comments are common to both)
  , generateVorbisComments
  , renderPictureData
  ) where

import Control.Applicative ((<|>))
import Control.Exception (catch, IOException)
import Control.Monad.Except (ExceptT, throwError, runExceptT)
import Control.Monad.IO.Class (liftIO)
import Data.Binary.Get
import Data.Binary.Put
import Data.Bits ((.|.), shiftL, shiftR, (.&.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as L
import qualified Data.HashMap.Strict as HM
import Data.Maybe (fromMaybe, listToMaybe)
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 Monatone.Metadata
import Paths_monatone (version)
import Data.Version (showVersion)

-- 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

-- | Buffer size for file operations (64KB)
bufferSize :: Int
bufferSize = 65536

-- | Write metadata to FLAC file incrementally
writeFLACMetadata :: Metadata -> AlbumArtUpdate -> OsPath -> Writer ()
writeFLACMetadata metadata artUpdate filePath = do
  -- Open file in read/write mode
  result <- liftIO $ tryIO $ withBinaryFile filePath ReadWriteMode $ \handle -> do
    runExceptT $ writeFLACHandleIncremental metadata artUpdate handle
  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)

-- | Write FLAC metadata using a file handle incrementally
writeFLACHandleIncremental :: Metadata -> AlbumArtUpdate -> Handle -> Writer ()
writeFLACHandleIncremental metadata artUpdate handle = do
  -- Verify FLAC signature
  liftIO $ hSeek handle AbsoluteSeek 0
  sig <- liftIO $ BS.hGet handle 4
  case BS.unpack sig of
    [0x66, 0x4C, 0x61, 0x43] -> pure ()  -- "fLaC"
    _ -> throwError $ CorruptedWrite "Invalid FLAC signature"

  -- Extract original STREAMINFO block for preservation (it's always first, 34 bytes)
  streamInfoHeader <- liftIO $ BS.hGet handle 4
  streamInfoData <- liftIO $ BS.hGet handle 34
  if BS.length streamInfoHeader < 4 || BS.length streamInfoData < 34
    then throwError $ CorruptedWrite "File too small for STREAMINFO block"
    else pure ()
  let originalStreamInfo = L.fromStrict $ BS.append streamInfoHeader streamInfoData

  -- Find where the audio data starts, collecting blocks we must keep
  -- verbatim; existing PICTURE blocks are kept when the art is unchanged
  let keepPictures = artUpdate == KeepExistingArt
  (audioDataOffset, preservedBlocks, originalVendor) <- scanMetadataBlocks handle 4 keepPictures

  -- Generate new metadata blocks with preserved STREAMINFO
  newMetadataBlocks <- generateMetadataBlocks metadata artUpdate originalStreamInfo preservedBlocks originalVendor
  let newMetadataSize = fromIntegral $ L.length newMetadataBlocks
  
  -- Get file size
  _ <- liftIO $ hFileSize handle
  
  -- Calculate size difference (new metadata vs old metadata)
  let oldMetadataSize = audioDataOffset - 4  -- Subtract "fLaC" signature
  let sizeDiff = newMetadataSize - oldMetadataSize
  
  if sizeDiff == 0 then do
    -- Same size, just overwrite metadata blocks
    liftIO $ do
      hSeek handle AbsoluteSeek 4  -- Position after "fLaC"
      L.hPut handle newMetadataBlocks
  else if sizeDiff > 0 then do
    -- Need to insert bytes
    insertBytesInFile handle sizeDiff audioDataOffset
    -- Write new metadata blocks
    liftIO $ do
      hSeek handle AbsoluteSeek 4
      L.hPut handle newMetadataBlocks
  else do
    -- Need to delete bytes
    let bytesToDelete = negate sizeDiff
    -- Write new metadata first
    liftIO $ do
      hSeek handle AbsoluteSeek 4
      L.hPut handle newMetadataBlocks
    -- Then delete extra space
    deleteBytesInFile handle bytesToDelete (4 + newMetadataSize)

-- | Walk the metadata blocks: return the audio data offset, the raw bytes
-- of every block we must carry over verbatim (SEEKTABLE, APPLICATION,
-- CUESHEET, unknown types - and PICTURE blocks when keepPictures is set),
-- and the original Vorbis vendor string. STREAMINFO and VORBIS_COMMENT
-- are regenerated; PADDING is dropped.
scanMetadataBlocks :: Handle -> Int -> Bool -> Writer (Int, [L.ByteString], Maybe BS.ByteString)
scanMetadataBlocks handle startOffset keepPictures = go startOffset [] Nothing
  where
    -- STREAMINFO, PADDING, VORBIS_COMMENT, and PICTURE unless kept
    regenerated = [0, 1, 4] ++ [6 | not keepPictures]

    go :: Int -> [L.ByteString] -> Maybe BS.ByteString -> Writer (Int, [L.ByteString], Maybe BS.ByteString)
    go currentOffset acc vendor = do
      liftIO $ hSeek handle AbsoluteSeek (fromIntegral currentOffset)
      headerBytes <- liftIO $ BS.hGet handle 4
      if BS.length headerBytes < 4
        then return (currentOffset, reverse acc, vendor)
        else do
          let header = runGet parseBlockHeader (L.fromStrict headerBytes)
              blockSize = fromIntegral (blockLength header)
              nextOffset = currentOffset + 4 + blockSize
          (acc', vendor') <- if blockType header `elem` regenerated
            then do
              -- Remember the original encoder's vendor string so a tag
              -- update does not claim the file was encoded by us
              newVendor <- if blockType header == 4 && vendor == Nothing
                then do
                  blockData <- liftIO $ BS.hGet handle blockSize
                  return $ extractVendor blockData
                else return vendor
              return (acc, newVendor)
            else do
              blockData <- liftIO $ BS.hGet handle blockSize
              if BS.length blockData < blockSize
                then throwError $ CorruptedWrite "Truncated metadata block"
                else pure ()
              -- Last-block flag is recomputed when the blocks are reassembled
              return (L.fromStrict (headerBytes <> blockData) : acc, vendor)
          if isLast header
            then return (nextOffset, reverse acc', vendor')
            else go nextOffset acc' vendor'

    extractVendor blockData
      | BS.length blockData < 4 = Nothing
      | otherwise =
          let vendorLen = fromIntegral $ runGet getWord32le (L.fromStrict (BS.take 4 blockData))
              vendorBytes = BS.take vendorLen (BS.drop 4 blockData)
          in if BS.length vendorBytes == vendorLen then Just vendorBytes else Nothing

-- | Insert bytes into file at given offset
insertBytesInFile :: Handle -> Int -> Int -> Writer ()
insertBytesInFile handle size offset = do
  -- Get current file size
  fileSize <- liftIO $ hFileSize handle
  let moveSize = fileSize - fromIntegral offset
  
  if moveSize < 0 then
    throwError $ WriteIOError "Invalid offset for insert"
  else do
    -- First, extend the file
    liftIO $ hSetFileSize handle (fileSize + fromIntegral size)
    
    -- Move data from offset to offset+size, working backwards
    moveDataBackwards handle (fromIntegral offset) (fromIntegral $ offset + size) moveSize

-- | Delete bytes from file at given offset  
deleteBytesInFile :: Handle -> Int -> Int -> Writer ()
deleteBytesInFile handle size offset = do
  -- Get current file size
  fileSize <- liftIO $ hFileSize handle
  let moveSize = fileSize - fromIntegral offset - fromIntegral size
  
  if moveSize < 0 then
    throwError $ WriteIOError "Invalid size/offset for delete"
  else do
    -- Move data from offset+size to offset
    moveDataForwards handle (fromIntegral $ offset + size) (fromIntegral offset) moveSize
    
    -- Truncate the file
    liftIO $ hSetFileSize handle (fileSize - fromIntegral size)

-- | Move data backwards in file (for insertions)
moveDataBackwards :: Handle -> Integer -> Integer -> Integer -> Writer ()
moveDataBackwards handle src dest count = do
  let go remaining' = do
        if remaining' <= 0 then
          return ()
        else do
          let chunkSize = min (fromIntegral bufferSize) remaining'
          -- Read from end of source region
          hSeek handle AbsoluteSeek (src + remaining' - chunkSize)
          chunk <- BS.hGet handle (fromIntegral chunkSize)
          -- Write to end of dest region
          hSeek handle AbsoluteSeek (dest + remaining' - chunkSize)
          BS.hPut handle chunk
          go (remaining' - chunkSize)
  
  liftIO $ go count

-- | Move data forwards in file (for deletions)
moveDataForwards :: Handle -> Integer -> Integer -> Integer -> Writer ()
moveDataForwards handle src dest count = do
  let go moved = do
        if moved >= count then
          return ()
        else do
          let chunkSize = min (fromIntegral bufferSize) (count - moved)
          -- Read from source
          hSeek handle AbsoluteSeek (src + moved)
          chunk <- BS.hGet handle (fromIntegral chunkSize)
          -- Write to dest
          hSeek handle AbsoluteSeek (dest + moved)
          BS.hPut handle chunk
          go (moved + chunkSize)
  
  liftIO $ go 0

-- | FLAC metadata block header
data BlockHeader = BlockHeader
  { isLast :: Bool
  , blockType :: Word8
  , blockLength :: Word32
  } deriving (Show)

-- | Parse FLAC metadata block header
parseBlockHeader :: Get BlockHeader
parseBlockHeader = do
  firstByte <- getWord8
  let lastFlag = (firstByte .&. 0x80) /= 0
  let bType = firstByte .&. 0x7F
  
  -- Block length is 24 bits
  b1 <- getWord8
  b2 <- getWord8
  b3 <- getWord8
  let len = (fromIntegral b1 `shiftL` 16) .|. 
            (fromIntegral b2 `shiftL` 8) .|. 
            fromIntegral b3
  
  return $ BlockHeader lastFlag bType len

-- | Generate new metadata blocks: STREAMINFO, preserved blocks, Vorbis
-- comment, and optionally a Picture block, with the last-block flag set on
-- the final block only. Kept pictures travel in preservedBlocks.
generateMetadataBlocks :: Metadata -> AlbumArtUpdate -> L.ByteString -> [L.ByteString] -> Maybe BS.ByteString -> Writer L.ByteString
generateMetadataBlocks metadata artUpdate originalStreamInfo preservedBlocks originalVendor = do
  vorbisBlock <- generateVorbisCommentBlock metadata originalVendor False
  pictureBlocks <- case artUpdate of
    SetAlbumArt albumArt -> (: []) <$> generatePictureBlock albumArt False
    _ -> return []

  let blocks = [originalStreamInfo] ++ preservedBlocks ++ [vorbisBlock] ++ pictureBlocks
  return $ L.concat $ markLastBlock blocks
  where
    markLastBlock [] = []
    markLastBlock [block] = [setLastFlag True block]
    markLastBlock (block:rest) = setLastFlag False block : markLastBlock rest

    setLastFlag set block = case L.uncons block of
      Just (firstByte, rest) ->
        L.cons (if set then firstByte .|. 0x80 else firstByte .&. 0x7F) rest
      Nothing -> block

-- | Generate Vorbis comment block, keeping the original vendor string
-- when there is one
generateVorbisCommentBlock :: Metadata -> Maybe BS.ByteString -> Bool -> Writer L.ByteString
generateVorbisCommentBlock metadata originalVendor isLastBlock = do
  let defaultVendor = TE.encodeUtf8 $ T.pack $ "Monatone " ++ showVersion version
  let vendorBytes = fromMaybe defaultVendor originalVendor
  let vendorLenBytes = runPut $ putWord32le $ fromIntegral $ BS.length vendorBytes
  
  -- Create comment list
  let comments = generateVorbisComments metadata
  let commentCount = length comments
  let commentCountBytes = runPut $ putWord32le $ fromIntegral commentCount
  
  -- Encode each comment
  let encodeComment (key, value) = 
        let text = key <> "=" <> value
            textBytes = TE.encodeUtf8 text
            lenBytes = runPut $ putWord32le $ fromIntegral $ BS.length textBytes
        in lenBytes <> L.fromStrict textBytes
  
  let encodedComments = L.concat $ map encodeComment comments
  
  -- Build complete Vorbis comment data
  let vorbisData = vendorLenBytes <> L.fromStrict vendorBytes <>
                  commentCountBytes <> encodedComments

  -- Create block header
  let headerByte = if isLastBlock then 0x84 else 0x04  -- Block type 4 = Vorbis comment
  header <- blockHeader headerByte vorbisData
  return $ header <> vorbisData

-- | Generate Vorbis comments from metadata. Every field the parser maps
-- is written back, and unmapped comments in rawTags are carried over so
-- an update never drops tags it does not understand. (Album art is not
-- included: FLAC stores it in a PICTURE block, OGG appends its own
-- METADATA_BLOCK_PICTURE comment.)
generateVorbisComments :: Metadata -> [(Text, Text)]
generateVorbisComments metadata = mappedComments ++ preservedComments
  where
    mbIds = musicBrainzIds metadata
    showT = T.pack . show

    -- A scalar field is the first value of a possibly multi-valued key.
    -- If the field still matches what was parsed, write every stored
    -- value back (several ARTIST comments survive an unrelated update);
    -- if it was edited, the edit replaces the whole set.
    multi key field = case field of
      Nothing -> []
      Just v ->
        let stored = HM.lookupDefault [] key (rawTags metadata)
        in if listToMaybe stored == Just v
             then [(key, s) | s <- stored]
             else [(key, v)]

    single key = maybe [] (\v -> [(key, v)])

    mappedComments = concat
      [ multi "TITLE" (title metadata)
      , multi "ARTIST" (artist metadata)
      , multi "ALBUM" (album metadata)
      , multi "ALBUMARTIST" (albumArtist metadata)
      , single "TRACKNUMBER" (showT <$> trackNumber metadata)
      , single "TRACKTOTAL" (showT <$> totalTracks metadata)
      , single "DISCNUMBER" (showT <$> discNumber metadata)
      , single "DISCTOTAL" (showT <$> totalDiscs metadata)
      , single "DATE" (date metadata <|> (showT <$> year metadata))
      , multi "GENRE" (genre metadata)
      , multi "COMMENT" (comment metadata)
      , multi "PUBLISHER" (publisher metadata)
      , single "BARCODE" (barcode metadata)
      , single "CATALOGNUMBER" (catalogNumber metadata)
      , multi "LABEL" (recordLabel metadata)
      , single "RELEASECOUNTRY" (releaseCountry metadata)
      , single "RELEASESTATUS" (releaseStatus metadata)
      , single "RELEASETYPE" (releaseType metadata)
      , single "MUSICBRAINZ_RELEASETRACKID" (mbTrackId mbIds)
      , single "MUSICBRAINZ_TRACKID" (mbRecordingId mbIds)
      , single "MUSICBRAINZ_ALBUMID" (mbReleaseId mbIds)
      , single "MUSICBRAINZ_RELEASEGROUPID" (mbReleaseGroupId mbIds)
      , multi "MUSICBRAINZ_ARTISTID" (mbArtistId mbIds)
      , multi "MUSICBRAINZ_ALBUMARTISTID" (mbAlbumArtistId mbIds)
      , single "MUSICBRAINZ_WORKID" (mbWorkId mbIds)
      , single "MUSICBRAINZ_DISCID" (mbDiscId mbIds)
      , single "ACOUSTID_FINGERPRINT" (acoustidFingerprint metadata)
      , single "ACOUSTID_ID" (acoustidId metadata)
      ]

    -- Keys the mapped fields own (whether or not they are set right now):
    -- stale rawTags copies of these must not be written back
    handledKeys =
      [ "TITLE", "ARTIST", "ALBUM", "ALBUMARTIST"
      , "TRACKNUMBER", "TRACKTOTAL", "DISCNUMBER", "DISCTOTAL"
      , "DATE", "YEAR", "GENRE", "COMMENT", "PUBLISHER"
      , "BARCODE", "CATALOGNUMBER", "LABEL", "RELEASECOUNTRY"
      , "RELEASESTATUS", "RELEASETYPE"
      , "MUSICBRAINZ_RELEASETRACKID", "MUSICBRAINZ_TRACKID"
      , "MUSICBRAINZ_ALBUMID", "MUSICBRAINZ_RELEASEGROUPID"
      , "MUSICBRAINZ_ARTISTID", "MUSICBRAINZ_ALBUMARTISTID"
      , "MUSICBRAINZ_WORKID", "MUSICBRAINZ_DISCID"
      , "ACOUSTID_FINGERPRINT", "ACOUSTID_ID"
      , "METADATA_BLOCK_PICTURE"  -- art is written as a PICTURE block instead
      ]

    preservedComments =
      [ (key, value)
      | (key, values) <- HM.toList (rawTags metadata)
      , T.toUpper key `notElem` handledKeys
      , value <- values
      ]

-- | Render the body of a FLAC PICTURE block (also the payload OGG encodes
-- as a base64 METADATA_BLOCK_PICTURE comment)
renderPictureData :: AlbumArt -> L.ByteString
renderPictureData art =
  let mimeBytes = TE.encodeUtf8 $ albumArtMimeType art
      descBytes = TE.encodeUtf8 $ albumArtDescription art
      imageData = albumArtData art
  in runPut $ do
       putWord32be $ fromIntegral $ albumArtPictureType art  -- Picture type
       putWord32be $ fromIntegral $ BS.length mimeBytes      -- MIME type length
       putByteString mimeBytes                               -- MIME type
       putWord32be $ fromIntegral $ BS.length descBytes      -- Description length
       putByteString descBytes                               -- Description
       putWord32be 0                                         -- Width (0 = unknown)
       putWord32be 0                                         -- Height (0 = unknown)
       putWord32be 0                                         -- Color depth (0 = unknown)
       putWord32be 0                                         -- Number of colors (0 = unknown)
       putWord32be $ fromIntegral $ BS.length imageData      -- Picture data length
       putByteString imageData                               -- Picture data

-- | Generate Picture block for album art
generatePictureBlock :: AlbumArt -> Bool -> Writer L.ByteString
generatePictureBlock art isLastBlock = do
  let pictureData = renderPictureData art
      headerByte = if isLastBlock then 0x86 else 0x06  -- Block type 6 = Picture

  header <- blockHeader headerByte pictureData
  return $ header <> pictureData

-- | Build a metadata block header, rejecting content that does not fit the
-- 24-bit length field (silently truncating it would corrupt the stream)
blockHeader :: Word8 -> L.ByteString -> Writer L.ByteString
blockHeader headerByte content = do
  let len = L.length content
  if len > 0xFFFFFF
    then throwError $ InvalidMetadata "Metadata block exceeds FLAC's 16MiB block size limit"
    else return $ runPut $ do
      putWord8 headerByte
      -- Write 24-bit length
      putWord8 $ fromIntegral $ (len `shiftR` 16) .&. 0xFF
      putWord8 $ fromIntegral $ (len `shiftR` 8) .&. 0xFF
      putWord8 $ fromIntegral $ len .&. 0xFF