Writer.hs

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

module Monatone.FLAC.Writer
  ( writeFLACMetadata
  , WriteError(..)
  , Writer
  ) 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 (catMaybes)
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
-- Takes optional AlbumArt separately since Metadata only stores AlbumArtInfo
writeFLACMetadata :: Metadata -> Maybe AlbumArt -> OsPath -> Writer ()
writeFLACMetadata metadata maybeAlbumArt filePath = do
  -- Open file in read/write mode
  result <- liftIO $ tryIO $ withBinaryFile filePath ReadWriteMode $ \handle -> do
    runExceptT $ writeFLACHandleIncremental metadata maybeAlbumArt 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 -> Maybe AlbumArt -> Handle -> Writer ()
writeFLACHandleIncremental metadata maybeAlbumArt 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
  let originalStreamInfo = L.fromStrict $ BS.append streamInfoHeader streamInfoData

  -- Find where the audio data starts, collecting blocks we must keep verbatim
  (audioDataOffset, preservedBlocks) <- scanMetadataBlocks handle 4  -- Start after "fLaC"

  -- Generate new metadata blocks with preserved STREAMINFO
  newMetadataBlocks <- generateMetadataBlocks metadata maybeAlbumArt originalStreamInfo preservedBlocks
  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 plus the raw
-- bytes of every block we must carry over verbatim (SEEKTABLE, APPLICATION,
-- CUESHEET, unknown types). STREAMINFO, VORBIS_COMMENT and PICTURE are
-- regenerated; PADDING is dropped.
scanMetadataBlocks :: Handle -> Int -> Writer (Int, [L.ByteString])
scanMetadataBlocks handle startOffset = go startOffset []
  where
    regenerated = [0, 1, 4, 6]  -- STREAMINFO, PADDING, VORBIS_COMMENT, PICTURE

    go currentOffset acc = do
      liftIO $ hSeek handle AbsoluteSeek (fromIntegral currentOffset)
      headerBytes <- liftIO $ BS.hGet handle 4
      if BS.length headerBytes < 4
        then return (currentOffset, reverse acc)
        else do
          let header = runGet parseBlockHeader (L.fromStrict headerBytes)
              blockSize = fromIntegral (blockLength header)
              nextOffset = currentOffset + 4 + blockSize
          acc' <- if blockType header `elem` regenerated
            then return acc
            else do
              blockData <- liftIO $ BS.hGet handle blockSize
              -- Last-block flag is recomputed when the blocks are reassembled
              return $ L.fromStrict (headerBytes <> blockData) : acc
          if isLast header
            then return (nextOffset, reverse acc')
            else go nextOffset acc'

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

-- | Extract the original STREAMINFO block (already read from handle)
_extractStreamInfoBlock :: L.ByteString -> Writer L.ByteString
_extractStreamInfoBlock blockData = do
  if L.length blockData < 38  -- 4 byte header + 34 byte STREAMINFO
    then throwError $ CorruptedWrite "File too small for STREAMINFO block"
    else return $ L.take 38 blockData  -- Include header + data

-- | 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
generateMetadataBlocks :: Metadata -> Maybe AlbumArt -> L.ByteString -> [L.ByteString] -> Writer L.ByteString
generateMetadataBlocks metadata maybeAlbumArt originalStreamInfo preservedBlocks = do
  vorbisBlock <- generateVorbisCommentBlock metadata False
  pictureBlocks <- case maybeAlbumArt of
    Nothing -> return []
    Just albumArt -> (: []) <$> generatePictureBlock albumArt False

  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
generateVorbisCommentBlock :: Metadata -> Bool -> Writer L.ByteString
generateVorbisCommentBlock metadata isLastBlock = do
  -- Create vendor string
  let vendor = T.pack $ "Monatone " ++ showVersion version
  let vendorBytes = TE.encodeUtf8 vendor
  let vendorLenBytes = runPut $ putWord32le $ fromIntegral $ BS.length vendorBytes
  
  -- Create comment list
  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 blockLen = fromIntegral $ L.length vorbisData :: Word32
  let headerByte = if isLastBlock then 0x84 else 0x04  -- Block type 4 = Vorbis comment
  let header = runPut $ do
        putWord8 headerByte
        -- Write 24-bit length
        putWord8 $ fromIntegral $ (blockLen `shiftR` 16) .&. 0xFF
        putWord8 $ fromIntegral $ (blockLen `shiftR` 8) .&. 0xFF
        putWord8 $ fromIntegral $ blockLen .&. 0xFF
  
  return $ header <> vorbisData

-- | Generate Vorbis comments from metadata. Every field the FLAC parser
-- maps is written back, and unmapped comments in rawTags are carried over
-- so an update never drops tags it does not understand.
generateVorbisComments :: Metadata -> Writer [(Text, Text)]
generateVorbisComments metadata = return $ mappedComments ++ preservedComments
  where
    mbIds = musicBrainzIds metadata
    showT = T.pack . show

    mappedComments = catMaybes
      [ ("TITLE",) <$> title metadata
      , ("ARTIST",) <$> artist metadata
      , ("ALBUM",) <$> album metadata
      , ("ALBUMARTIST",) <$> albumArtist metadata
      , ("TRACKNUMBER",) . showT <$> trackNumber metadata
      , ("TRACKTOTAL",) . showT <$> totalTracks metadata
      , ("DISCNUMBER",) . showT <$> discNumber metadata
      , ("DISCTOTAL",) . showT <$> totalDiscs metadata
      , ("DATE",) <$> (date metadata <|> (showT <$> year metadata))
      , ("GENRE",) <$> genre metadata
      , ("COMMENT",) <$> comment metadata
      , ("PUBLISHER",) <$> publisher metadata
      , ("BARCODE",) <$> barcode metadata
      , ("CATALOGNUMBER",) <$> catalogNumber metadata
      , ("LABEL",) <$> recordLabel metadata
      , ("RELEASECOUNTRY",) <$> releaseCountry metadata
      , ("RELEASESTATUS",) <$> releaseStatus metadata
      , ("RELEASETYPE",) <$> releaseType metadata
      , ("MUSICBRAINZ_RELEASETRACKID",) <$> mbTrackId mbIds
      , ("MUSICBRAINZ_TRACKID",) <$> mbRecordingId mbIds
      , ("MUSICBRAINZ_ALBUMID",) <$> mbReleaseId mbIds
      , ("MUSICBRAINZ_RELEASEGROUPID",) <$> mbReleaseGroupId mbIds
      , ("MUSICBRAINZ_ARTISTID",) <$> mbArtistId mbIds
      , ("MUSICBRAINZ_ALBUMARTISTID",) <$> mbAlbumArtistId mbIds
      , ("MUSICBRAINZ_WORKID",) <$> mbWorkId mbIds
      , ("MUSICBRAINZ_DISCID",) <$> mbDiscId mbIds
      , ("ACOUSTID_FINGERPRINT",) <$> acoustidFingerprint metadata
      , ("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, value) <- HM.toList (rawTags metadata)
      , T.toUpper key `notElem` handledKeys
      ]

-- | Generate Picture block for album art
generatePictureBlock :: AlbumArt -> Bool -> Writer L.ByteString
generatePictureBlock art isLastBlock = do
  let mimeBytes = TE.encodeUtf8 $ albumArtMimeType art
      descBytes = TE.encodeUtf8 $ albumArtDescription art
      imageData = albumArtData art

      -- Build picture data according to FLAC spec
      pictureData = 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

      blockLen = fromIntegral $ L.length pictureData :: Word32
      headerByte = if isLastBlock then 0x86 else 0x06  -- Block type 6 = Picture

      -- Build block header
      header = runPut $ do
        putWord8 headerByte
        -- Write 24-bit length
        putWord8 $ fromIntegral $ (blockLen `shiftR` 16) .&. 0xFF
        putWord8 $ fromIntegral $ (blockLen `shiftR` 8) .&. 0xFF
        putWord8 $ fromIntegral $ blockLen .&. 0xFF

  return $ header <> pictureData