module Shader
    ( makeShader
    , addTexture
    , addUniforms
    , pokeShaders
    --, pokeShader
    , pokeShader'
    --, bindArrayBuffers
    , bindArrayBuffers'
    , bindShaderBuffers
    , makeSourcedShader
    , setupVAO
    , floatSize
    , numDrawableElements
    , drawShader
    , drawShaders
    , setShaderUniforms
    , resetShaderUniforms
    , extractProgAndUnis 
    , freeShaderPointers
    , setPerpMatUniform
    , module Shader.Data
    )
  where
import Geometry
import Shader.Data
import MatrixHelper
import Picture.Data

import Foreign
import Codec.Picture
import qualified Data.ByteString as BS
import qualified Data.Vector.Storable as V
import Control.Monad (unless, forM, zipWithM_, forM_, foldM)
import Control.Lens
import qualified Control.Foldl as F
import Graphics.Rendering.OpenGL hiding (Point,translate,scale,imageHeight)
--import Text.RawString.QQ
import Linear.Matrix
import Linear.V4

extractProgAndUnis :: FullShader -> (Program,UniformLocation)
extractProgAndUnis s = (_shaderProgram s, _shaderMatrixUniform s)

pokeShaders :: [FullShader] -> F.FoldM IO RenderType [Int]
pokeShaders = traverse pokeShader'
--pokeShaders = traverse pokeShader

pokeShader :: FullShader -> F.FoldM IO RenderType Int
pokeShader fs = F.FoldM (pokeRender fls (zip ptrs nAtss))  (return 0) return
  where 
    (_,ptrs,nAtss) = unzip3 $ _vaoBufferTargets $ _shaderVAO fs
    fls = _shaderPokeStrategy fs

pokeShader' :: FullShader -> F.FoldM IO RenderType Int
pokeShader' fs = F.FoldM (pokeRender' fls ptr stride)  (return 0) return
  where 
    theVBO = _vaoVBO $ _shaderVAO fs
    ptr = _vboPointer theVBO
    stride = sum $ _vboAttribSizes theVBO
    fls = _shaderPokeStrategy fs

pokeRender :: (RenderType -> [[[Float]]])
                -> [(Ptr Float,Int)] -> Int -> RenderType -> IO Int
pokeRender toFs ptrs n rt = pokeList ptrs n (toFs rt)

pokeRender' 
    :: (RenderType 
    -> [[[Float]]])
    -> Ptr Float
    -> Int  -- ^ stride
    -> Int 
    -> RenderType 
    -> IO Int
pokeRender' toFs ptr stride n rt = pokeList' ptr stride n (toFs rt)

pokeList :: [(Ptr Float,Int)] -> Int -> [[[Float]]] -> IO Int
pokeList ptrs = foldM (pokePtrs ptrs)

pokeList' :: Ptr Float -> Int -> Int -> [[[Float]]] -> IO Int
pokeList' ptr stride = foldM (pokePtrs' ptr stride)

pokePtrs' :: Ptr Float -> Int -> Int -> [[Float]] -> IO Int
pokePtrs' ptr stride n fss = do
    pokeArrayOff ptr (stride * n) (concat fss)
    return $ n + 1

pokePtrs :: [(Ptr Float,Int)] -> Int -> [[Float]] -> IO Int
pokePtrs ptrIs n fss = do
    zipWithM_ f ptrIs fss
    return $ n + 1
  where 
    f (ptr,i) fs = pokeArrayOff ptr (i*n) fs

pokeArrayOff :: Storable a => Ptr a -> Int -> [a] -> IO ()
pokeArrayOff ptr i = zipWithM_ (pokeElemOff ptr) [i..]

bindArrayBuffers :: Int -> [(BufferObject,Ptr Float,Int)] -> IO ()
{-# INLINE bindArrayBuffers #-}
bindArrayBuffers numVs ps = do
    forM_ ps $ \(bo,ptr,i) -> do
        bindBuffer ArrayBuffer $= Just bo
        bufferData ArrayBuffer $= (fromIntegral $ floatSize * numVs * i, ptr, DynamicDraw)

bindArrayBuffers' :: Int -> VBO -> IO ()
bindArrayBuffers' numVs theVBO = do
    bindBuffer ArrayBuffer $= Just (_vbo theVBO)
    bufferData ArrayBuffer $= 
        (fromIntegral $ floatSize * numVs * (sum $ _vboAttribSizes theVBO), _vboPointer theVBO, DynamicDraw)
    

bindShaderBuffers :: [FullShader] -> [Int] -> IO ()
bindShaderBuffers = zipWithM_ f
  where 
    --f fs i = bindArrayBuffers i $ _vaoBufferTargets $ _shaderVAO fs
    f fs i = bindArrayBuffers' i $ _vaoVBO $ _shaderVAO fs

drawShaders :: [FullShader] -> [Int] -> IO ()
drawShaders = zipWithM_ drawShader

drawShader :: FullShader -> Int -> IO ()
{-# INLINE drawShader #-}
drawShader fs i = do 
    currentProgram $= Just (_shaderProgram fs)
    bindVertexArrayObject $= Just (_vao $ _shaderVAO fs)
    case _shaderTexture fs of 
        Just ShaderTexture{_textureObject = txo}
            -> textureBinding Texture2D $= Just txo
        _   -> return ()
    drawArrays (_shaderDrawPrimitive fs) 0 (fromIntegral i)

-- I am not sure if this assumes that the shader is constructed directly before
-- the texture is added...
addTexture :: String -> FullShader -> IO (FullShader)
addTexture texturePath shad = do
    Right cmap <- readImage texturePath
    let tex = convertRGBA8 cmap
    textureOb <- genObjectName
    textureBinding Texture2D $= Just textureOb
    let texData = V.toList $ imageData tex
        wtex = fromIntegral $ imageWidth  tex
        htex = fromIntegral $ imageHeight tex
    withArray texData $ \ptr -> do
        texImage2D Texture2D NoProxy 0 RGBA8 (TextureSize2D wtex htex) 0 
                        (PixelData RGBA UnsignedByte ptr)
    generateMipmap' Texture2D
    textureFilter Texture2D $= ((Linear',Just Linear') , Nearest)
    return $ shad & shaderTexture ?~ ShaderTexture {_textureObject = textureOb}

{- | 
Compiles a full shader found within the shader directory.
The shader is made up of files begining with the inputted string with extensions .vert, .geom etc.
-}
makeShader 
    :: String       -- ^ First part of the name of the shader
    -> [ShaderType] -- ^ Filetype extensions
    -> [Int]   -- ^ The input vertex sizes
    -> PrimitiveMode
    -> (RenderType -> [[[Float]]])   -- ^ Poke strategy: method for creating a list of vertex data to be bound
    -> IO (FullShader)
makeShader s shaderlist sizes pm renStrat = do
    (prog,unis)  <- makeSourcedShader s shaderlist
    vaob        <- setupVAO sizes
    return $ FullShader { _shaderProgram = prog
                        , _shaderMatrixUniform = unis
                        , _shaderVAO = vaob
                        , _shaderPokeStrategy = renStrat
                        , _shaderDrawPrimitive = pm
                        , _shaderTexture = Nothing
                        , _shaderCustomUnis = []
                        }
floatSize :: Int
floatSize = sizeOf (0.5 :: GLfloat)
{-# INLINE floatSize #-}

bufferOffset :: Integral a => a -> Ptr b
bufferOffset = plusPtr nullPtr . (* floatSize) . fromIntegral

numDrawableElements :: Int
{-# INLINE numDrawableElements #-}
numDrawableElements = 50000

setupVAO :: [Int] -> IO VAO
setupVAO sizes = do
    theVAO <- genObjectName
    bindVertexArrayObject $= Just theVAO
    vbos <- forM (zip (map fromIntegral [(0::Int)..]) sizes) setupArrayBuffer
    ptrs <- forM (zip vbos sizes) setupVBOPointers
    theVBO <- setupVBO $ sizes
    return $ VAO 
        { _vao = theVAO
        , _vaoBufferTargets = ptrs
        , _vaoVBO = theVBO
        }

setupVBO :: [Int] -> IO VBO
setupVBO sizes = do
    vboName <- genObjectName
    bindBuffer ArrayBuffer $= Just vboName
    forM_ (zip3 [0..] sizes offs) $ \(loc,siz,off) -> do
        setupVertexAttribPointer loc siz strd off
    thePtr <- mallocArray (strd * numDrawableElements)
    return $ VBO
        { _vbo = vboName
        , _vboPointer = thePtr
        , _vboAttribSizes = sizes
        }
  where
    strd = sum sizes
    offs = scanl (+) 0 sizes
    
{- | Assumes the correct VBO is bound -}
setupVertexAttribPointer 
    :: Int  -- ^ Atrib location
    -> Int  -- ^ Size
    -> Int  -- ^ Stride 
    -> Int  -- ^ Offset
    -> IO ()
setupVertexAttribPointer loc siz strd off = do
    vertexAttribPointer (AttribLocation (fi loc)) $=
        (ToFloat, VertexArrayDescriptor (fi' siz) Float (fi'' $ floatSize * strd) (bufferOffset off))
    vertexAttribArray (AttribLocation (fi loc)) $= Enabled
  where
    fi = fromIntegral
    fi' = fromIntegral
    fi'' = fromIntegral

setupVBOPointers :: (BufferObject,Int) -> IO (BufferObject,Ptr Float,Int)
setupVBOPointers (vbo,vsize) = do
    thePtr <- mallocArray (vsize * numDrawableElements)
    return (vbo,thePtr,vsize)

setupArrayBuffer :: (GLuint,Int) -> IO BufferObject
setupArrayBuffer (aloc,i) = do
    vbo <- genObjectName
    bindBuffer ArrayBuffer $= Just vbo
    vertexAttribPointer (AttribLocation aloc) $=
        ( ToFloat
        , VertexArrayDescriptor (fromIntegral i) 
                                Float 
                                (fromIntegral $ floatSize * i) 
                                (bufferOffset (0::Int))
        )
    vertexAttribArray (AttribLocation aloc) $= Enabled
    return vbo

-- | Compile shader and get its uniform locations.
-- supposes the shader code is in the shader folder, with the string names
-- followed by .vert/.geom/.frag.
makeSourcedShader :: String -> [ShaderType] -> IO (Program, UniformLocation)
makeSourcedShader s sts = do
    sources <- forM sts $ \st -> BS.readFile ("shader/" ++ s ++ shaderTypeExt st)
    prog <- makeShaderProgram s $ zip sts sources
    uniformLocations <- uniformLocation prog "worldMat"
    return (prog,uniformLocations)

addUniforms :: [String] -> FullShader -> IO (FullShader)
addUniforms uniStrings shad = do
    uniLocs <- mapM (uniformLocation $ _shaderProgram shad) uniStrings
    return $ shad & shaderCustomUnis %~ (++ uniLocs)

shaderTypeExt :: ShaderType -> String
shaderTypeExt VertexShader   = ".vert"
shaderTypeExt GeometryShader = ".geom"
shaderTypeExt FragmentShader = ".frag"
shaderTypeExt _ = undefined

makeShaderProgram :: String -> [(ShaderType,BS.ByteString)] -> IO Program
makeShaderProgram str sources = do
    theShaderProgram <- createProgram
    shaders <- mapM (compileAndCheckShader str) sources
    mapM_ (attachShader theShaderProgram) shaders

    linkProgram theShaderProgram
    linkingSuccess <- linkStatus theShaderProgram
    unless linkingSuccess $ do
        infoLog <- get (programInfoLog theShaderProgram)
        putStrLn $ str ++ ": Program Linking" ++ infoLog
    return theShaderProgram

compileAndCheckShader :: String -> (ShaderType,BS.ByteString) -> IO Shader
compileAndCheckShader str (theShaderType,sourceCode) = do
    theShader <- createShader theShaderType
    shaderSourceBS theShader $= sourceCode
    compileShader theShader
    success <- compileStatus theShader
    unless success $ do
        infoLog <- get (shaderInfoLog theShader)
        putStrLn $ str ++ ": Shader compile: " ++ show theShaderType ++ " : " ++ show infoLog
    return theShader

resetShaderUniforms :: [(Program, UniformLocation)] -> IO ()
{-# INLINE resetShaderUniforms #-}
resetShaderUniforms = setShaderUniforms 0 1 (0,0) (2,2)

setShaderUniforms :: Float -> Float -> Point2 -> Point2 -> [(Program,UniformLocation)] -> IO ()
{-# INLINE setShaderUniforms #-}
setShaderUniforms rot czoom (tranx,trany) (winx,winy) fss = do
    let scalMat = Linear.Matrix.transpose $
            V4 (V4 (2*czoom/winx) 0             0 (0::GLfloat))
                (V4 0             (2*czoom/winy) 0 0)
                (V4 0             0             1 0)
                (V4 0             0             0 1)
    let rotMat = Linear.Matrix.transpose $
            V4 (V4 (cos rot) (sin (-rot)) 0 0)
                (V4 (sin rot) (cos rot)    0 0)
                (V4 0             0        1 0)
                (V4 0             0        0 1)
    let tranMat = Linear.Matrix.transpose $
            V4 (V4 1        0        0 0)
                (V4 0        1        0 0)
                (V4 0        0        1 0)
                (V4 (-tranx) (-trany) 0 1)
    let wmat = scalMat !*! rotMat !*! tranMat
        vToL (V4 a b c d) = [a,b,c,d]
    wmata <- (newMatrix RowMajor $ concatMap vToL $ vToL wmat) :: IO (GLmatrix GLfloat)
-- set common uniforms
    forM_ fss $ \shad -> do
        currentProgram $= Just (fst shad)
        uniform (snd shad) $= wmata

setPerpMatUniform 
    :: Float -- ^ rotation
    -> Float -- ^ zoom
    -> Point2 -- ^ translation
    -> Point2 -- ^ window size
    -> Point2 -- ^ viewfrom point
    -> FullShader
    -> IO ()
setPerpMatUniform rot czoom trans wins vFrom shad = do
    pmat <- (newMatrix RowMajor $ perspectiveMatrix rot czoom trans wins vFrom) :: IO (GLmatrix GLfloat)
    currentProgram $= Just (_shaderProgram shad)
    uniform (_shaderMatrixUniform shad) $= pmat
    return ()

freeShaderPointers :: FullShader -> IO ()
freeShaderPointers fs = do
    forM_ (_vaoBufferTargets $ _shaderVAO fs) $ \(_,ptr,_) ->
        free ptr