{-# LANGUAGE LambdaCase, RecordWildCards, KindSignatures, GADTs #-}

-- TODO: Formulate this module as a patch against lambdacube-gl.

module MaskableStream where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Writer
import Data.Foldable
import Data.Function
import Data.Int
import Data.IORef
import Data.Maybe
import Data.Word
import qualified Data.Map.Strict as Map
import qualified Data.Vector as V
         ;import Data.Vector as V ((!),(//))
import Foreign.C.Types (CPtrdiff)
import Foreign.Marshal
import Foreign.Ptr
import Foreign.Storable
import Data.Dependent.Sum
import Data.Some
import Data.GADT.Show
import GHC.TypeLits
import System.IO

import LambdaCube.GL      as LC
import LambdaCube.GL.Data -- (uploadCubeMapToGPU,uploadTextureBufferToGPU,updateTextureBuffer)
import LambdaCube.GL.Mesh as LC
import LambdaCube.GL.Type
import LambdaCube.IR      as LC
import LambdaCube.GL.Util
import LambdaCube.GL.Input.Type
import LambdaCube.GL.Input hiding (createObjectCommands)

import Graphics.GL.Core33


-- based on addMeshToObjectArray
addToObjectArray :: GLStorage
                        -> String   -- ^ Slot name for a PrimitiveStream.
                        -> [String] -- ^ Uniform names.  IORefs will be put in 'objUniSetup'.
                        -> GPUData
                        -> IO Object
addToObjectArray input slotName objUniNames (GPUData prim streams indices _) = do
    let ObjectArraySchema _ slotStreams = fromMaybe (error $ "addMeshToObjectArray - missing object array: " ++ slotName)
                                                    $ Map.lookup slotName $! objectArrays $! schema input
    addObject input slotName prim indices (Map.intersection streams slotStreams) objUniNames

setUniformCommand :: (String -> GLUniform)
                               -> (String,GLint)
                               -> GLObjectCommand
setUniformCommand ulookup (n,i) = GLSetUniform i (ulookup n)

bindTextureCommand :: (String -> IORef GLint)
                   -> (String -> GLUniform)
                   -> String
                   -> GLObjectCommand
bindTextureCommand tlookup ulookup n = GLBindTexture (inputTypeToTextureTarget $ uniInputType u) (tlookup n) u
 where
    u = ulookup n
    uniInputType (GLTypedUniform ty _) = unwitnessType ty
    uniInputType (GLUniform r)         = objectType r

-- One of these:
--    VertexAttribPointer         -- Array
--    VertexAttribIPointer        --  streams.
--    VertexAttribI{1,2,3,4}[u]iv -- constant integer attribute.
-- Or some number of these in sequence:
--    VertexAttrib{1,2,3,4}fv -- constant vector or matrix of floats.
setVertexAttribCommmand :: (t -> Stream Buffer) -> (GLuint, t) -> GLObjectCommand
setVertexAttribCommmand alookup (i,name) = case alookup name of
    Stream ty (Buffer arrs bo) arrIdx start len -> mkAttrCmd i bo n (arrayTypeToGLType arrType) (intPtrToPtr $! offset)
      where
        (n, mkAttrCmd) = classifyStreamType ty
        ArrayDesc arrType arrLen arrOffs arrSize = arrs ! arrIdx
        offset = fromIntegral (arrOffs + start * fromIntegral n * sizeOfArrayType arrType)

    -- constant generic attribute
    constAttr -> GLSetVertexAttrib i constAttr


lookupOrLookup :: (Ord k, Show k) => String -> Map.Map k a -> Map.Map k a -> k -> a
lookupOrLookup callsite objUnis topUnis n = Map.findWithDefault (topUni n) n objUnis
 where
    topUni n = Map.findWithDefault (error $ "internal error ("++callsite++"): " ++ show n) n topUnis

classifyStreamType :: StreamType -> ( GLint
                                    , GLuint -> GLuint -> GLint -> GLenum -> Ptr () -> GLObjectCommand)
classifyStreamType = \case
    Attribute_Word   -> ( 1 , GLSetVertexAttribIArray )
    Attribute_V2U    -> ( 2 , GLSetVertexAttribIArray )
    Attribute_V3U    -> ( 3 , GLSetVertexAttribIArray )
    Attribute_V4U    -> ( 4 , GLSetVertexAttribIArray )
    Attribute_Int    -> ( 1 , GLSetVertexAttribIArray )
    Attribute_V2I    -> ( 2 , GLSetVertexAttribIArray )
    Attribute_V3I    -> ( 3 , GLSetVertexAttribIArray )
    Attribute_V4I    -> ( 4 , GLSetVertexAttribIArray )
    Attribute_Float  -> ( 1  , GLSetVertexAttribArray )
    Attribute_V2F    -> ( 2  , GLSetVertexAttribArray )
    Attribute_V3F    -> ( 3  , GLSetVertexAttribArray )
    Attribute_V4F    -> ( 4  , GLSetVertexAttribArray )
    Attribute_M22F   -> ( 4  , GLSetVertexAttribArray )
    Attribute_M23F   -> ( 6  , GLSetVertexAttribArray )
    Attribute_M24F   -> ( 8  , GLSetVertexAttribArray )
    Attribute_M32F   -> ( 6  , GLSetVertexAttribArray )
    Attribute_M33F   -> ( 9  , GLSetVertexAttribArray )
    Attribute_M34F   -> ( 12 , GLSetVertexAttribArray )
    Attribute_M42F   -> ( 8  , GLSetVertexAttribArray )
    Attribute_M43F   -> ( 12 , GLSetVertexAttribArray )
    Attribute_M44F   -> ( 16 , GLSetVertexAttribArray )

drawElementsCommand :: GLenum -> IndexStream Buffer -> GLObjectCommand
drawElementsCommand prim (IndexStream (Buffer arrs bo) arrIdx start idxCount)
        = GLDrawElements prim (fromIntegral idxCount) idxType bo ptr
 where
    ArrayDesc arrType arrLen arrOffs arrSize = arrs ! arrIdx
    idxType = arrayTypeToGLType arrType
    ptr    = intPtrToPtr $! fromIntegral (arrOffs + start * sizeOfArrayType arrType)

-- Note: all enabled attribute arrays should have the same count.
-- So an object should select similarly sized streams from 'dStreams' of GPUData.
-- So the 'mAttributes' map of a Mesh should contain equally sized arrays.
-- uploadMeshToGPU uses a single Buffer for all of a Mesh's attributes.
objectDrawStyle :: (GLsizei -> [(GLint,GLsizei)]) -- ^ mask, normally \x -> [(0,x)]
                    -> Object
                    -> Either [(GLint,GLsizei)] (IndexStream Buffer)
objectDrawStyle streamMask obj = case objIndices obj of
    Nothing -> Left $ let cnt = head [c | Stream _ _ _ _ c <- Map.elems (objAttributes obj)]
                       in streamMask $ fromIntegral cnt
    Just idxStream -> Right idxStream

-- backward-compatible non-masking interface.
createObjectCommands :: Map.Map String (IORef GLint) -> Map.Map String GLUniform -> Object -> GLProgram -> [GLObjectCommand]
createObjectCommands texUnitMap topUnis obj prg =
    createObjectCommands_ (objectEnvironment texUnitMap topUnis obj (pure . ((,) 0))) prg

data ObjectEnvironment = ObjectEnvironment
    { envPrim      :: GLenum -- GL version of LambdaCube.GL.Primitive
    , envDrawStyle :: Either [(GLint,GLsizei)]    -- Mask of array attributes (one draw command for pair).
                             (IndexStream Buffer) -- Indirect P_TriangleStripI or P_TrianglesI indices.
    , tlookup      :: String -> IORef GLint   -- lookup texture unit
    , ulookup      :: String -> GLUniform     -- lookup uniform
    , alookup      :: String -> Stream Buffer -- lookup vertex attribute
    }

objectEnvironment :: Map.Map String (IORef GLint)
                     -> Map.Map String GLUniform
                     -> Object
                     -> (GLsizei -> [(GLint, GLsizei)])
                     -> ObjectEnvironment
objectEnvironment texUnitMap topUnis obj streamMask = ObjectEnvironment
    { envPrim      = primitiveToGLType $ objPrimitive obj
    , envDrawStyle = objectDrawStyle streamMask obj
    , tlookup      = \n -> Map.findWithDefault (error $ "missing texture unit: " ++ show n) n texUnitMap
    , ulookup      = lookupOrLookup "missing uniform: " (objUniSetup obj) topUnis
    , alookup      = \n -> Map.findWithDefault (error $ "missing attribute: " ++ n) n $ objAttributes obj
    }


createObjectCommands_ :: ObjectEnvironment -> GLProgram -> [GLObjectCommand]
createObjectCommands_ ObjectEnvironment{..} prg = concat
    [ map (setUniformCommand ulookup)          $ Map.toList (inputUniforms prg)
    , map (bindTextureCommand tlookup ulookup) $ toList (inputTextureUniforms prg)
    , map (setVertexAttribCommmand alookup)    $ Map.elems (inputStreams prg)
    , case envDrawStyle of
        Left ranges   -> map (uncurry $ GLDrawArrays envPrim) ranges
        Right indexed -> pure $ drawElementsCommand envPrim indexed
    ]

updateCommands :: GLStorage
                  -> Object
                  -> (GLsizei -> [(GLint,GLsizei)])
                  -> IO Object
updateCommands input obj mask = do
    let cmdsRef = objCommands obj :: IORef (V.Vector (V.Vector [GLObjectCommand]))
        slotIdx = objSlot obj :: Int
    ppls <- readIORef $ pipelines input
    cmds <- V.forM (ppls :: V.Vector (Maybe GLRenderer)) $ \mp -> case mp of
        Nothing -> return V.empty
        Just p  -> do
            let env = objectEnvironment (glTexUnitMapping p) (uniformSetup input) obj mask
            Just ic <- readIORef $ glInput p
            case icSlotMapInputToPipeline ic ! slotIdx of
                Nothing         -> return V.empty -- this slot is not used in that pipeline
                Just pSlotIdx   -> do
                    let emptyV = V.replicate (V.length $ glPrograms p) []
                    return $ emptyV // [(prgIdx,createObjectCommands_ env (glPrograms p ! prgIdx))
                                        | prgIdx <- glSlotPrograms p ! pSlotIdx]
    writeIORef cmdsRef cmds
    return obj


componentCount :: GLABI c -> Int
componentCount x@(IsGLVector _) = fromIntegral $ natVal $ vectorLength x
componentCount x@(IsGLMatrix _) = let (r,c) = matrixDimensions x
                                    in fromIntegral (natVal r) * fromIntegral (natVal c)

uploadDynamicBuffer :: Int -> [Parameter] -> IO GPUData
uploadDynamicBuffer sz params = do
    bo <- alloca $! \pbo -> glGenBuffers 1 pbo >> peek pbo
    glBindBuffer GL_ARRAY_BUFFER bo
    let (mkstreams, mkarrays) = unzip $ mapMaybe attrInfo params
        attrInfo (Parameter n typ) = do
            atyp <- toStreamType typ
            This tt <- witnessType typ
            let abi = glABI tt
                cnt = componentCount abi
                arrtyp = let go :: GLPointerType typ -> ArrayType
                             go = \case GLPrimUInt  -> ArrWord32
                                        GLPrimInt   -> ArrInt32
                                        GLPrimFloat -> ArrFloat
                         in case abi of { IsGLVector p -> go p ; IsGLMatrix p -> go p }
            return ( \b i -> (n, Stream atyp b i 0 sz)
                   , \offset -> ArrayDesc arrtyp (cnt * sz) offset (4 * sz * cnt))
        streams = zipWith ($ buffer) mkstreams [0..]
        arrays = foldr mk (const []) mkarrays 0
            where mk f fin offset = let a = f offset
                                    in a : fin (offset + arrSize a)
        buffer = Buffer (V.fromList arrays) bo
        byteCount = sum $ map arrSize arrays
    glBufferData GL_ARRAY_BUFFER (fromIntegral byteCount) nullPtr GL_DYNAMIC_DRAW
    glBindBuffer GL_ARRAY_BUFFER 0
    let gd = GPUData PointList (Map.fromList streams) Nothing [buffer]
    return gd

incrementalUpdateBuffer :: MonadIO m => Buffer -> GLintptr -> GLsizeiptr -> Ptr a -> m ()
incrementalUpdateBuffer b byteoffset bytecount ptr = do
    glBindBuffer GL_ARRAY_BUFFER (bufGLObj b)
    glBufferSubData GL_ARRAY_BUFFER byteoffset bytecount ptr
    glBindBuffer GL_ARRAY_BUFFER 0

data AttributeKey c = AttributeKey (TypeTag c) Buffer CPtrdiff

attributeKey :: TypeTagable c => GPUData -> String -> Maybe (AttributeKey c)
attributeKey dta name = do
    stream <- Map.lookup name (dStreams dta)
    fix $ \mp -> let typ = typeTag (fromJust mp) in case stream of
        Stream t b i _ _ -> do
            let a = bufArrays b ! i
                -- arrType :: ArrayType
                -- arrLength :: Int -- number of 32 bit values
                -- arrOffset :: Int -- byte offset into buffer
                -- arrSize   :: Int -- byte count
                off = arrOffset a
            guard (fromStreamType t == unwitnessType typ)
            Just $ AttributeKey typ b (fromIntegral off)
        _ -> Nothing

(@<-) :: GLData a c => AttributeKey c -> a -> Writer [DSum AttributeKey GLUniformValue] ()
k @<- v = tell [k :=> GLUniformValue v]

updateAttributes :: Int -> Writer [DSum AttributeKey GLUniformValue] a -> IO ()
updateAttributes i writer = forM_ (execWriter writer) $ \case
    AttributeKey typ b base :=> GLUniformValue a -> do
        glBindBuffer GL_ARRAY_BUFFER (bufGLObj b)
        let abi = glABI typ
            attribSize = 4 * componentCount abi
        case marshalUniform abi a of

            Just (MarshalGLVector with) -> with $ \sz ptr -> do
                let sz' = fromIntegral $ attribSize * (fromIntegral sz)
                putStrLn $ "vector sz = " ++ show sz
                glBufferSubData GL_ARRAY_BUFFER (base + fromIntegral i * sz') sz' ptr

            Just (MarshalGLMatrix with) -> with $ \sz isrowcol ptr -> case isrowcol of
                0 -> do
                    let sz' = fromIntegral $ attribSize * (fromIntegral sz)
                    glBufferSubData GL_ARRAY_BUFFER (base + fromIntegral i * sz') sz' ptr
                _ -> hPutStrLn stderr $ "WARNING: (TODO) row-major matrix attribute update unimplemented."

            Nothing -> hPutStrLn stderr $ "Warning: dimension mismatch updating " ++ show (unwitnessType typ) ++ " attribute."
        glBindBuffer GL_ARRAY_BUFFER 0