module Shader.Compile
    ( makeShader
    , makeByteStringShader
    , makeByteStringShaderUsingVAO
    , makeShaderSized
    , makeShaderUsingShaderVAO
    , makeShaderUsingVAO
    , makeSourcedShader
    , setupVAO
    , setupVertexAttribPointer 
    ) where
import Shader.Data
import Shader.Parameters

import Foreign
import qualified Data.ByteString as BS
import Control.Monad
--import Control.Lens
import Graphics.Rendering.OpenGL hiding (Point,translate,scale,imageHeight)
{- | 
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
    -> EPrimitiveMode 
    -> IO FullShader
makeShader s shaderlist sizes pm = do
    prog <- makeSourcedShader s shaderlist
    vaob <- setupVAO sizes
    return $ FullShader 
        { _shadProg = prog
        , _shadVAO  = vaob
        , _shadPrim = pm
        , _shadTex  = Nothing
        , _shadUnis = []
        }
makeByteStringShader
    :: String       -- ^ (Arbitrary) name of the shader
    -> [(ShaderType,BS.ByteString)] -- ^ Filetype extensions and shader data
    -> [Int]        -- ^ The input vertex sizes
    -> EPrimitiveMode 
    -> IO FullShader
makeByteStringShader s shaderlist sizes pm = do
    prog <- makeShaderProgram s shaderlist
    vaob <- setupVAO sizes
    return $ FullShader 
        { _shadProg = prog
        , _shadVAO  = vaob
        , _shadPrim = pm
        , _shadTex  = Nothing
        , _shadUnis = []
        }
makeByteStringShaderUsingVAO
    :: String       -- ^ (Arbitrary) name of the shader
    -> [(ShaderType,BS.ByteString)] -- ^ Filetype extensions and shader data
    -> EPrimitiveMode 
    -> FullShader
    -> IO FullShader
makeByteStringShaderUsingVAO s shaderlist pm fs = do
    prog <- makeShaderProgram s shaderlist
    return $ fs 
        { _shadProg = prog
        , _shadPrim = pm
        , _shadTex  = Nothing
        , _shadUnis = []
        }

-- | Takes the VAO from elsewhere
makeShaderUsingVAO
    :: String       -- ^ First part of the name of the shader
    -> [ShaderType] -- ^ Filetype extensions
    -> EPrimitiveMode
    -> VAO
    -> IO FullShader
makeShaderUsingVAO s shaderlist pm theVAO = do
    prog <- makeSourcedShader s shaderlist
    return $ FullShader 
        { _shadProg = prog
        , _shadVAO = theVAO
        , _shadPrim = pm
        , _shadTex  = Nothing
        , _shadUnis = []
        }

-- | Takes the VAO from another shader
makeShaderUsingShaderVAO
    :: String       -- ^ First part of the name of the shader
    -> [ShaderType] -- ^ Filetype extensions
    -> EPrimitiveMode
    -> FullShader
    -> IO FullShader
makeShaderUsingShaderVAO s shaderlist pm fs = do
    prog <- makeSourcedShader s shaderlist
    return $ fs 
        { _shadProg = prog
        , _shadPrim = pm
        , _shadTex  = Nothing
        , _shadUnis = []
        }
{- | 
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.
-}
makeShaderSized
    :: String       -- ^ First part of the name of the shader
    -> [ShaderType] -- ^ Filetype extensions
    -> [Int]   -- ^ The input vertex sizes
    -> Int          -- ^ Number of vertexes that can be poked
    -> EPrimitiveMode
    -> IO FullShader
makeShaderSized s shaderlist sizes ndraw pm = do
    prog <- makeSourcedShader s shaderlist
    vaob <- setupVAOSized sizes ndraw
    return $ FullShader 
        { _shadProg = prog
        , _shadVAO  = vaob
        , _shadPrim = pm
        , _shadTex  = Nothing
        , _shadUnis = []
        }

-- | 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
makeSourcedShader s sts = do
    sources <- forM sts $ \st -> BS.readFile ("shader/" ++ s ++ shaderTypeExt st)
    makeShaderProgram s $ zip sts sources

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

setupVAO :: [Int] -> IO VAO
setupVAO sizes = do
    theVAO <- genObjectName
    bindVertexArrayObject $= Just theVAO
    theVBO <- setupVBO sizes
    return $ VAO 
        { _vao = theVAO
        , _vaoVBO = theVBO
        }
setupVAOSized :: [Int] -> Int -> IO VAO
setupVAOSized sizes ndraw = do
    theVAO <- genObjectName
    bindVertexArrayObject $= Just theVAO
    theVBO <- setupVBOSized sizes ndraw
    return $ VAO 
        { _vao = theVAO
        , _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)
    -- Allocate space
    bufferData ArrayBuffer $=
        (fromIntegral $ floatSize * numDrawableElements * strd
        , nullPtr
        , StreamDraw
        )
    return $ VBO
        { _vbo = vboName
        , _vboPtr = thePtr
        , _vboAttribSizes = sizes
        , _vboStride = sum sizes
        }
  where
    strd = sum sizes
    offs = scanl (+) 0 sizes

setupVBOSized :: [Int] -> Int -> IO VBO
setupVBOSized sizes ndraw = do
    vboName <- genObjectName
    bindBuffer ArrayBuffer $= Just vboName
    forM_ (zip3 [0..] sizes offs) $ \(loc,siz,off) -> do
        setupVertexAttribPointer loc siz strd off
    thePtr <- mallocArray (strd * ndraw)
    -- Allocate space
    bufferData ArrayBuffer $=
        (fromIntegral $ floatSize * ndraw * strd
        , nullPtr
        , StreamDraw
        )
    return $ VBO
        { _vbo = vboName
        , _vboPtr = thePtr
        , _vboAttribSizes = sizes
        , _vboStride = sum 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

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 ++ " :\n" ++ infoLog
    return theShader

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

