OGG.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}

module Monatone.OGG
  ( parseOGG
  , loadAlbumArtOGG
  ) where

import Control.Applicative ((<|>))
import Control.Monad.Except (throwError)
import Control.Monad.IO.Class (liftIO)
import Data.Binary.Get
import Data.Bits ((.&.), shiftL)
import Data.Int (Int32)
import Data.Maybe (listToMaybe)
import Data.Word (Word64)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base64 as B64
import qualified Data.ByteString.Lazy as L
import System.IO (Handle, IOMode(..), hFileSize, hSeek, SeekMode(..))
import System.OsPath
import System.File.OsPath (withBinaryFile)
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 qualified Data.Text.Encoding.Error as TEE

import Monatone.Metadata  
import Monatone.Types

-- | OGG page header is always 27 bytes (before segment table)
oggPageHeaderSize :: Int
oggPageHeaderSize = 27

-- | Refuse to reassemble packets beyond this size (a comment packet with
-- embedded art is a few MB; this only guards against malformed lacing)
maxPacketSize :: Int
maxPacketSize = 64 * 1024 * 1024

-- | Parse OGG file efficiently - only the first two logical packets are
-- read: the identification header and the comment header (this layout is
-- shared by Vorbis and Opus)
parseOGG :: OsPath -> Parser Metadata
parseOGG filePath = do
  result <- liftIO $ withBinaryFile filePath ReadMode $ \handle -> do
    -- Read first page to check OGG signature
    firstHeader <- BS.hGet handle 27
    if BS.length firstHeader < 27 || BS.take 4 firstHeader /= "OggS"
      then return $ Left $ UnsupportedFormat "Not an OGG file"
      else do
        hSeek handle AbsoluteSeek 0
        packets <- readFirstPackets handle 2
        let metadata = foldl (flip applyPacket) (emptyMetadata OGG) packets
        maybeGranule <- findLastGranule handle
        return $ Right $ applyDuration packets maybeGranule metadata

  case result of
    Left err -> throwError err
    Right m -> return m

-- | Duration comes from the last page's granule position: total samples
-- for Vorbis, 48kHz sample count (including pre-skip) for Opus
applyDuration :: [BS.ByteString] -> Maybe Word64 -> Metadata -> Metadata
applyDuration packets maybeGranule metadata =
  case (maybeGranule, codec (audioProperties metadata)) of
    (Just granule, Just CodecOpus) ->
      setDuration (fromIntegral granule - opusPreSkip) 48000
    (Just granule, Just CodecVorbis) ->
      case sampleRate (audioProperties metadata) of
        Just rate | rate > 0 -> setDuration (fromIntegral granule) rate
        _ -> metadata
    _ -> metadata
  where
    setDuration :: Int -> Int -> Metadata
    setDuration samples rate
      | samples <= 0 = metadata
      | otherwise =
          let ms = round (fromIntegral samples * 1000 / fromIntegral rate :: Double)
          in metadata { audioProperties = (audioProperties metadata) { duration = Just ms } }

    opusPreSkip :: Int
    opusPreSkip = case [p | p <- packets, "OpusHead" `BS.isPrefixOf` p] of
      (p:_) | BS.length p >= 12 ->
        fromIntegral (BS.index p 10) + 256 * fromIntegral (BS.index p 11)
      _ -> 0

-- | Granule position of the file's last page (read from the tail without
-- scanning the whole file). Pages with the "no packet ends here" sentinel
-- granule (-1) are skipped.
findLastGranule :: Handle -> IO (Maybe Word64)
findLastGranule handle = do
  fileSize <- hFileSize handle
  let readLen = fromIntegral (min fileSize 65536)
  hSeek handle SeekFromEnd (negate (fromIntegral readLen))
  tailBytes <- BS.hGet handle readLen
  return $ case [g | off <- pageOffsets tailBytes
                   , Just g <- [granuleAt tailBytes off]
                   , g /= maxBound] of
    [] -> Nothing
    granules -> Just (last granules)
  where
    pageOffsets bs = go 0
      where
        go from = case BS.breakSubstring "OggS" (BS.drop from bs) of
          (before, rest)
            | BS.null rest -> []
            | otherwise ->
                let off = from + BS.length before
                in off : go (off + 4)

    granuleAt bs off
      | off + 14 > BS.length bs = Nothing
      | otherwise = Just $ sum
          [ fromIntegral (BS.index bs (off + 6 + i)) `shiftL` (8 * i) | i <- [0 .. 7] ]

-- | Apply a metadata packet to the metadata being built
applyPacket :: BS.ByteString -> Metadata -> Metadata
applyPacket packet metadata
  | "\x01vorbis" `BS.isPrefixOf` packet = parseVorbisInfo packet metadata
  | "OpusHead" `BS.isPrefixOf` packet = parseOpusInfo packet metadata
  | otherwise = case commentPayload packet of
      Just payload -> case runGetOrFail (parseVorbisCommentGet metadata) payload of
        Left _ -> metadata
        Right (_, _, result) -> result
      Nothing -> metadata

-- | Strip the header magic from a comment packet, if it is one. Vorbis and
-- Opus comment packets share the Vorbis comment structure after the magic
commentPayload :: BS.ByteString -> Maybe L.ByteString
commentPayload packet
  | "\x03vorbis" `BS.isPrefixOf` packet = Just $ L.fromStrict $ BS.drop 7 packet
  | "OpusTags" `BS.isPrefixOf` packet = Just $ L.fromStrict $ BS.drop 8 packet
  | otherwise = Nothing

-- | Read the first n complete logical packets of an OGG stream,
-- reassembling packets that span multiple segments and pages: a packet is
-- the concatenation of consecutive segments up to and including the first
-- segment shorter than 255 bytes, and may continue across pages (header
-- flag 0x01 on the continuing page)
readFirstPackets :: Handle -> Int -> IO [BS.ByteString]
readFirstPackets handle limit = goPage [] []
  where
    -- completed and partial both hold newest-first chunks
    goPage completed partial
      | length completed >= limit = return $ take limit $ reverse completed
      | otherwise = do
          maybePage <- readOggPage handle
          case maybePage of
            Nothing -> return $ take limit $ reverse completed  -- EOF or corrupt page
            Just (continued, segments) -> do
              -- A page that does not continue a packet implicitly abandons
              -- any leftover partial data (malformed stream)
              let partial0 = if continued then partial else []
                  (completed', partial') = foldl step (completed, partial0) segments
              if sum (map BS.length partial') > maxPacketSize
                then return $ take limit $ reverse completed'
                else goPage completed' partial'

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

-- | Read one OGG page: the continuation flag and the segments with a marker
-- for whether each segment ends a packet (lacing value < 255)
readOggPage :: Handle -> IO (Maybe (Bool, [(BS.ByteString, Bool)]))
readOggPage handle = do
  headerBytes <- BS.hGet handle oggPageHeaderSize
  if BS.length headerBytes < oggPageHeaderSize || BS.take 4 headerBytes /= "OggS"
    then return Nothing
    else do
      let continued = BS.index headerBytes 5 .&. 0x01 /= 0
          numSegments = fromIntegral $ BS.index headerBytes 26
      segmentTable <- BS.hGet handle numSegments
      if BS.length segmentTable < numSegments
        then return Nothing
        else do
          let lacings = map fromIntegral $ BS.unpack segmentTable :: [Int]
          pageData <- BS.hGet handle (sum lacings)
          if BS.length pageData < sum lacings
            then return Nothing
            else return $ Just (continued, splitSegments lacings pageData)
  where
    splitSegments [] _ = []
    splitSegments (lacing:rest) bs =
      let (segment, unread) = BS.splitAt lacing bs
      in (segment, lacing < 255) : splitSegments rest unread

-- | Parse Vorbis identification header (packet type 1)
parseVorbisInfo :: BS.ByteString -> Metadata -> Metadata
parseVorbisInfo bs metadata =
  if BS.length bs < 30  -- Minimum size for valid header
    then metadata
    else
      let lazyBs = L.fromStrict bs
      in case runGetOrFail (parseVorbisInfoGet metadata) (L.drop 7 lazyBs) of
        Left _ -> metadata
        Right (_, _, result) -> result

parseVorbisInfoGet :: Metadata -> Get Metadata
parseVorbisInfoGet metadata = do
  _ <- getWord32le  -- vorbisVersion
  audioChannels <- getWord8
  audioSampleRate <- getWord32le
  -- The bitrate fields are signed per the Vorbis spec; -1 means "unset"
  -- and must not be treated as a huge unsigned value
  bitrateMaximum <- fromIntegral <$> getWord32le :: Get Int32
  bitrateNominal <- fromIntegral <$> getWord32le :: Get Int32
  bitrateMinimum <- fromIntegral <$> getWord32le :: Get Int32

  let bitrate' = if bitrateNominal > 0
                then Just $ fromIntegral $ bitrateNominal `div` 1000
                else if bitrateMaximum > 0 && bitrateMinimum > 0
                then Just $ fromIntegral $ (bitrateMaximum + bitrateMinimum) `div` 2000
                else Nothing
  
  return $ metadata
    { audioProperties = AudioProperties
      { sampleRate = Just $ fromIntegral audioSampleRate
      , channels = Just $ fromIntegral audioChannels
      , bitrate = bitrate'
      , bitsPerSample = Nothing  -- Not in Vorbis info
      , duration = Nothing  -- Would need granule position from last page
      , codec = Just CodecVorbis
      }
    }

-- | Parse Opus identification header (the "OpusHead" packet)
parseOpusInfo :: BS.ByteString -> Metadata -> Metadata
parseOpusInfo bs metadata =
  if BS.length bs < 19  -- 8 (magic) + version + channels + pre-skip + input rate + gain + mapping
    then metadata
    else
      let lazyBs = L.fromStrict bs
      in case runGetOrFail (parseOpusInfoGet metadata) (L.drop 8 lazyBs) of
        Left _ -> metadata
        Right (_, _, result) -> result

parseOpusInfoGet :: Metadata -> Get Metadata
parseOpusInfoGet metadata = do
  _ <- getWord8            -- version
  opusChannels <- getWord8
  _ <- getWord16le         -- pre-skip
  _ <- getWord32le         -- input sample rate (informational only, RFC 7845)
  return $ metadata
    { audioProperties = emptyAudioProperties
      -- Opus always decodes at 48kHz; the header's input rate is just a
      -- record of what the encoder was fed
      { sampleRate = Just 48000
      , channels = Just $ fromIntegral opusChannels
      , codec = Just CodecOpus
      }
    }

-- | Parse the Vorbis comment structure (shared by Vorbis and Opus, after
-- the packet magic has been stripped)
parseVorbisCommentGet :: Metadata -> Get Metadata
parseVorbisCommentGet metadata = do
  -- Read vendor string length (little-endian 32-bit)
  vendorLength <- getWord32le
  -- Skip vendor string
  skip (fromIntegral vendorLength)
  
  -- Read number of comments
  numComments <- getWord32le
  
  -- Read each comment
  comments <- parseCommentList (fromIntegral numComments)

  -- Vorbis comments may repeat a key (e.g. several ARTIST entries); keep
  -- every value. The scalar metadata fields take the first one.
  let tagMap = HM.fromListWith (flip (<>)) [(k, [v]) | (k, v) <- comments]
      firstOf key = HM.lookup key tagMap >>= listToMaybe

  -- Extract standard fields
  return $ metadata
    { title = firstOf "TITLE"
    , artist = firstOf "ARTIST"
    , album = firstOf "ALBUM"
    , albumArtist = firstOf "ALBUMARTIST"
    , year = (firstOf "YEAR" >>= readInt)
             <|> (firstOf "DATE" >>= extractYearFromDate)
    , date = firstOf "DATE"
    , comment = firstOf "COMMENT"
    , genre = firstOf "GENRE"
    , trackNumber = firstOf "TRACKNUMBER" >>= readInt
    , totalTracks = firstOf "TRACKTOTAL" >>= readInt
    , discNumber = firstOf "DISCNUMBER" >>= readInt
    , totalDiscs = firstOf "DISCTOTAL" >>= readInt
    , releaseCountry = firstOf "RELEASECOUNTRY"
    , recordLabel = firstOf "LABEL"
    , catalogNumber = firstOf "CATALOGNUMBER"
    , barcode = firstOf "BARCODE"
    , releaseStatus = firstOf "RELEASESTATUS"
    , releaseType = firstOf "RELEASETYPE"
    , albumArtInfo = firstOf "METADATA_BLOCK_PICTURE" >>= parseVorbisPictureInfo
    , musicBrainzIds = MusicBrainzIds
      { mbTrackId = firstOf "MUSICBRAINZ_RELEASETRACKID"
      , mbRecordingId = firstOf "MUSICBRAINZ_TRACKID"
      , mbReleaseId = firstOf "MUSICBRAINZ_ALBUMID"
      , mbReleaseGroupId = firstOf "MUSICBRAINZ_RELEASEGROUPID"
      , mbArtistId = firstOf "MUSICBRAINZ_ARTISTID"
      , mbAlbumArtistId = firstOf "MUSICBRAINZ_ALBUMARTISTID"
      , mbWorkId = firstOf "MUSICBRAINZ_WORKID"
      , mbDiscId = firstOf "MUSICBRAINZ_DISCID"
      }
    , rawTags = tagMap
    }
  where
    parseCommentList :: Int -> Get [(Text, Text)]
    parseCommentList 0 = return []
    parseCommentList n = do
      -- Read comment length
      commentLength <- getWord32le
      -- Read comment data
      commentBytes <- getByteString (fromIntegral commentLength)
      -- Parse the comment (format: "KEY=value")
      let comment' = case BS.split 0x3D commentBytes of -- Split on '='
            (key:value:rest') -> 
              let keyText = T.toUpper $ TE.decodeUtf8With TEE.lenientDecode key
                  valueText = TE.decodeUtf8With TEE.lenientDecode (BS.intercalate "=" (value:rest'))
              in Just (keyText, valueText)
            _ -> Nothing
      rest <- parseCommentList (n - 1)
      return $ case comment' of
        Just c -> c : rest
        Nothing -> rest

-- | Parse Vorbis picture info (base64-encoded FLAC picture block, metadata only)
parseVorbisPictureInfo :: Text -> Maybe AlbumArtInfo
parseVorbisPictureInfo encodedData =
  case B64.decode (TE.encodeUtf8 encodedData) of
    Left _ -> Nothing
    Right pictureData -> parseFLACPictureBlockInfo pictureData
  where
    parseFLACPictureBlockInfo :: BS.ByteString -> Maybe AlbumArtInfo
    parseFLACPictureBlockInfo bs =
      let lazyBs = L.fromStrict bs
      in case runGetOrFail parsePictureInfo lazyBs of
        Left _ -> Nothing
        Right (_, _, artInfo) -> Just artInfo

    parsePictureInfo :: Get AlbumArtInfo
    parsePictureInfo = do
      pictureType <- getWord32be
      mimeLength <- getWord32be
      mimeType <- getByteString (fromIntegral mimeLength)
      descLength <- getWord32be
      description <- getByteString (fromIntegral descLength)
      _width <- getWord32be
      _height <- getWord32be
      _colorDepth <- getWord32be
      _numColors <- getWord32be
      pictureDataLength <- getWord32be
      -- Skip reading the actual picture data for performance
      -- skip (fromIntegral pictureDataLength)

      return $ AlbumArtInfo
        { albumArtInfoMimeType = TE.decodeUtf8With TEE.lenientDecode mimeType
        , albumArtInfoPictureType = fromIntegral pictureType
        , albumArtInfoDescription = TE.decodeUtf8With TEE.lenientDecode description
        , albumArtInfoSizeBytes = fromIntegral pictureDataLength
        }

-- | Load album art from OGG file (full binary data for writing)
loadAlbumArtOGG :: OsPath -> Parser (Maybe AlbumArt)
loadAlbumArtOGG filePath = do
  liftIO $ withBinaryFile filePath ReadMode $ \handle -> do
    -- Read first page to check OGG signature
    firstHeader <- BS.hGet handle 27
    if BS.length firstHeader < 27 || BS.take 4 firstHeader /= "OggS"
      then return Nothing
      else do
        -- The comment packet is the second logical packet in the stream
        hSeek handle AbsoluteSeek 0
        packets <- readFirstPackets handle 2
        return $ case [p | packet <- packets, Just p <- [commentPayload packet]] of
          (payload:_) -> extractPictureFromComment payload
          [] -> Nothing
  where
    extractPictureFromComment :: L.ByteString -> Maybe AlbumArt
    extractPictureFromComment payload =
      case runGetOrFail parseVorbisCommentForPicture payload of
        Left _ -> Nothing
        Right (_, _, result) -> result

    parseVorbisCommentForPicture :: Get (Maybe AlbumArt)
    parseVorbisCommentForPicture = do
      vendorLength <- getWord32le
      skip (fromIntegral vendorLength)
      numComments <- getWord32le
      findPictureComment (fromIntegral numComments)

    findPictureComment :: Int -> Get (Maybe AlbumArt)
    findPictureComment 0 = return Nothing
    findPictureComment n = do
      commentLength <- getWord32le
      commentBytes <- getByteString (fromIntegral commentLength)
      case BS.split 0x3D commentBytes of
        (key:value:rest) ->
          -- Rejoin on '=' so base64 padding at the end of the value survives
          let keyText = T.toUpper $ TE.decodeUtf8With TEE.lenientDecode key
              valueText = TE.decodeUtf8With TEE.lenientDecode (BS.intercalate "=" (value:rest))
          in if keyText == "METADATA_BLOCK_PICTURE"
             then return $ parseVorbisPictureFull valueText
             else findPictureComment (n - 1)
        _ -> findPictureComment (n - 1)

    parseVorbisPictureFull :: Text -> Maybe AlbumArt
    parseVorbisPictureFull encodedData =
      case B64.decode (TE.encodeUtf8 encodedData) of
        Left _ -> Nothing
        Right pictureData -> parseFLACPictureBlockFull pictureData

    parseFLACPictureBlockFull :: BS.ByteString -> Maybe AlbumArt
    parseFLACPictureBlockFull bs =
      let lazyBs = L.fromStrict bs
      in case runGetOrFail parsePictureData lazyBs of
        Left _ -> Nothing
        Right (_, _, art) -> Just art

    parsePictureData :: Get AlbumArt
    parsePictureData = do
      pictureType <- getWord32be
      mimeLength <- getWord32be
      mimeType <- getByteString (fromIntegral mimeLength)
      descLength <- getWord32be
      description <- getByteString (fromIntegral descLength)
      _width <- getWord32be
      _height <- getWord32be
      _colorDepth <- getWord32be
      _numColors <- getWord32be
      pictureDataLength <- getWord32be
      pictureData <- getByteString (fromIntegral pictureDataLength)

      return $ AlbumArt
        { albumArtMimeType = TE.decodeUtf8With TEE.lenientDecode mimeType
        , albumArtPictureType = fromIntegral pictureType
        , albumArtDescription = TE.decodeUtf8With TEE.lenientDecode description
        , albumArtData = pictureData
        }

-- | Extract year from DATE field (YYYY-MM-DD or just YYYY)
extractYearFromDate :: T.Text -> Maybe Int
extractYearFromDate dateText =
  let yearStr = T.takeWhile (/= '-') dateText
  in readInt yearStr