{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

module Hsm.Cam
  ( Cam
  , capturePng
  , runCam
  ) where

import Codec.Picture (Image(Image), encodePng)
import Codec.Picture.Types (PixelRGB8)
import Control.Concurrent (MVar, newEmptyMVar, putMVar, takeMVar)
import Control.Exception (mask_)
import Data.ByteString.Lazy (ByteString)
import Data.Primitive.Ptr (readOffPtr)
import Data.Vector.Storable (generateM)
import Effectful (Dispatch(Static), DispatchOf, Eff, IOE, (:>), liftIO)
import Effectful.Dispatch.Static (SideEffects(WithSideEffects), StaticRep, evalStaticRep, getStaticRep, unsafeEff_)
import Effectful.Exception (bracket, bracket_)
import Foreign.C.String (peekCString)
import Foreign.C.Types (CSize(CSize))
import Foreign.Ptr (Ptr, castPtr, freeHaskellFunPtr, nullPtr)
import Hsm.Cam.FFI
  ( acquireCamera
  , allocateFrameBuffer
  , createRequest
  , frameBufferLength
  , frameHeight
  , frameLine
  , frameStride
  , frameWidth
  , freeFrameBuffer
  , getDmaBufferFd
  , makeCallback
  , makeLogger
  , registerCallback
  , registerLogger
  , releaseCamera
  , requestFrame
  , startCamera
  , startCameraManager
  , stopCamera
  , stopCameraManager
  )
import Hsm.Log (Log, Severity(Info, Trace), getLoggerIO, logMsg)
import MMAP (mapShared, mkMmapFlags, mmap, munmap, protRead)

data Cam (a :: * -> *) (b :: *)

type instance DispatchOf Cam = Static WithSideEffects

data Rep = Rep
  { callbackMVar :: MVar ()
  , dmaBuffer :: Ptr ()
  }

newtype instance StaticRep Cam =
  Cam Rep

capturePng :: (Log "cam" :> es, Cam :> es) => Eff es ByteString
capturePng = do
  Cam Rep {..} <- getStaticRep
  logMsg Trace "Requesting frame"
  unsafeEff_ . mask_ $ requestFrame >> takeMVar callbackMVar
  logMsg Trace "Processing frame data"
  pixelVector <- unsafeEff_ . generateM (frameLine * frameHeight) $ mapPixel dmaBuffer
  logMsg Trace "Encoding PNG"
  return . encodePng $ Image @PixelRGB8 frameWidth frameHeight pixelVector
  where
    mapPixel dmaBuffer index = readOffPtr (castPtr dmaBuffer) offset
      where
        yIndex = index `div` frameLine
        xIndex = index `mod` frameLine
        offset = yIndex * frameStride + xIndex

runCam :: (IOE :> es, Log "cam" :> es) => Eff (Cam : es) a -> Eff es a
runCam action = do
  callbackMVar <- liftIO newEmptyMVar
  loggerBracket . callbackBracket callbackMVar . cameraManagerBracket . cameraAcquireBracket . frameBufferBracket . cameraStartBracket . createRequestBracket . mmapBracket $ \dmaBuffer -> do
    evalStaticRep (Cam Rep {..}) action
  where
    loggerBracket = bracket loggerAlloc loggerDealloc . const
      where
        loggerAlloc = do
          logMsg Info "Registering FFI logger"
          loggerIO <- getLoggerIO
          loggerFFI <- liftIO . makeLogger $ \severity message -> peekCString message >>= loggerIO (toEnum severity)
          liftIO $ registerLogger loggerFFI
          return loggerFFI
        loggerDealloc loggerFFI = do
          logMsg Info "Unregistering FFI logger"
          liftIO $ freeHaskellFunPtr loggerFFI
    callbackBracket callbackMVar = bracket callbackAlloc callbackDealloc . const
      where
        callbackAlloc = do
          logMsg Info "Registering FFI callback"
          callbackFFI <- liftIO . makeCallback $ putMVar callbackMVar ()
          liftIO $ registerCallback callbackFFI
          return callbackFFI
        callbackDealloc callbackFFI = do
          logMsg Info "Unregistering FFI callback"
          liftIO $ freeHaskellFunPtr callbackFFI
    -- Convenience wrappers for specific libcamera resources
    bracketFFI alloc dealloc = bracket_ (liftIO alloc) (liftIO dealloc)
    cameraManagerBracket = bracketFFI startCameraManager stopCameraManager
    cameraAcquireBracket = bracketFFI acquireCamera releaseCamera
    frameBufferBracket = bracketFFI allocateFrameBuffer freeFrameBuffer
    cameraStartBracket = bracketFFI startCamera stopCamera
    createRequestBracket = bracketFFI createRequest $ return ()
    -- Memory maps the camera's DMA buffer for frame access
    mmapBracket = bracket mmapDmaBuffer munmapDmaBuffer
      where
        mmapSize = CSize $ toEnum frameBufferLength
        mmapFlags = mkMmapFlags mapShared mempty
        mmapDmaBuffer = do
          logMsg Info "Mapping DMA buffer"
          liftIO $ getDmaBufferFd >>= \dmaBufferFd -> mmap nullPtr mmapSize protRead mmapFlags dmaBufferFd 0
        munmapDmaBuffer dmaBuffer = do
          logMsg Info "Unmapping DMA buffer"
          liftIO $ munmap dmaBuffer mmapSize