module IPS.Apply (apply, applyM, unsafeApplyM) where

import IPS.Types
import Data.Vector.Storable.Mutable (MVector)
import qualified Data.Vector.Storable.Mutable as MVector
import qualified Data.Vector.Storable as Vector
import Control.Monad.ST
import Control.Monad.Primitive (PrimMonad (PrimState))
import Data.Foldable


type BytesM m = MVector (PrimState m) Word8

-- explodes if the patch tries to write outside the existing space
unsafeApplyM :: PrimMonad m => Patch -> BytesM m -> m ()
unsafeApplyM p buf =
  for_ p \Chunk {offset, body} ->
    case body of
      Normal bs -> do
        let offset' = fromIntegral offset
        let sz = Vector.length bs
        let dst = MVector.slice offset' sz buf
        Vector.copy dst bs
      RLE {size, value} -> do
        let offset' = fromIntegral offset
        let sz = fromIntegral size
        let dst = MVector.slice offset' sz buf
        MVector.set dst value

-- size of patch output
-- (patches can write past the end of the input file)
outSize :: Patch -> Int
outSize = foldl' (\m c -> max m (end c)) 0 where
  end (Chunk {offset, body}) = case body of
    Normal bs  -> offset' + Vector.length bs
    RLE {size} -> offset' + fromIntegral size
    where offset' = fromIntegral offset

-- applies in place if it fits, otherwise allocates a new buffer for it
applyM :: PrimMonad m => Patch -> BytesM m -> m (Maybe (BytesM m))
applyM p buf = do
  let size = outSize p
  (buf', realloc) <-
    if MVector.length buf >= size then
      pure (buf, False)
    else do
      newBuf <- MVector.new size
      let dst = MVector.slice 0 (MVector.length buf) newBuf
      MVector.copy dst buf
      pure (newBuf, True)
  unsafeApplyM p buf'
  pure $ if realloc then Just buf' else Nothing

-- copies an immutable vector (resizing if needed) and patches it
apply :: Patch -> Bytes -> Bytes
apply p inp = runST do
  let size = outSize p `max` Vector.length inp
  buf <- MVector.new size
  let dst = MVector.slice 0 (Vector.length inp) buf
  Vector.copy dst inp
  unsafeApplyM p buf
  Vector.unsafeFreeze buf