{-# LANGUAGE BangPatterns, FlexibleContexts, TypeSynonymInstances, FlexibleInstances #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE StandaloneDeriving         #-}
module LambdaCube.GL.Input where

import Control.Applicative
import Control.Exception
import Control.Monad
import Control.Monad.Reader
import Control.Monad.Writer
import Data.Maybe
import Data.IORef
import Data.Map (Map)
import Data.IntMap (IntMap)
import Data.Vector (Vector,(//),(!))
import Data.Word
import Data.String
import Data.Typeable
import Foreign
import qualified Data.Dependent.Map as DMap
import qualified Data.IntMap as IM
import qualified Data.Set as S
import qualified Data.Map as Map
import qualified Data.Vector as V
import qualified Data.Vector.Algorithms.Intro as I
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as SB

import Graphics.GL.Core33

import LambdaCube.GL.Input.Type
import LambdaCube.GL.Type as T
import LambdaCube.GL.Util
import LambdaCube.IR as IR
import LambdaCube.Linear as IR
import LambdaCube.PipelineSchema

import qualified LambdaCube.IR as IR

schemaFromPipeline :: IR.Pipeline -> PipelineSchema
schemaFromPipeline a = PipelineSchema (Map.fromList sl) (foldl Map.union Map.empty ul)
  where
    (sl,ul) = unzip [( (sName,ObjectArraySchema sPrimitive (fmap cvt sStreams))
                     , sUniforms
                     )
                    | IR.Slot sName sStreams sUniforms sPrimitive _ <- V.toList $ IR.slots a
                    ]
    cvt a = case toStreamType a of
        Just v  -> v
        Nothing -> error "internal error (schemaFromPipeline)"

mkUniform :: [(String,InputType)] -> IO (Map String GLUniform)
mkUniform l = do
    unis <- forM l $ \(n,t) -> do
        uni <- initializeUniform t
        return (n,uni)
    return (Map.fromList unis)


allocStorage :: PipelineSchema -> IO GLStorage
allocStorage sch = do
    let sm  = Map.fromList $ zip (Map.keys $ objectArrays sch) [0..]
        len = Map.size sm
    unis <- mkUniform $ Map.toList $ uniforms sch
    seed <- newIORef 0
    slotV <- V.replicateM len $ newIORef (GLSlot IM.empty V.empty Ordered)
    size <- newIORef (0,0)
    ppls <- newIORef $ V.singleton Nothing
    return $ GLStorage
        { schema        = sch
        , slotMap       = sm
        , slotVector    = slotV
        , objSeed       = seed
        , uniformSetup  = unis
        , screenSize    = size
        , pipelines     = ppls
        }

disposeStorage :: GLStorage -> IO ()
disposeStorage _ = putStrLn "not implemented: disposeStorage"

-- object
addObject :: GLStorage -> String -> Primitive -> Maybe (IndexStream Buffer) -> Map String (Stream Buffer) -> [String] -> IO Object
addObject input slotName prim indices attribs uniformNames = do
    let sch = schema input
    forM_ uniformNames $ \n -> case Map.lookup n (uniforms sch) of
        Nothing -> fail $ "Unknown uniform: " ++ show n
        _ -> return ()
    case Map.lookup slotName (objectArrays sch) of
        Nothing -> fail $ "Unknown slot: " ++ show slotName
        Just (ObjectArraySchema sPrim sAttrs) -> do
            when (sPrim /= (primitiveToFetchPrimitive prim)) $ fail $
                "Primitive mismatch for slot (" ++ show slotName ++ ") expected " ++ show sPrim  ++ " but got " ++ show prim
            let sType = fmap streamToStreamType attribs
            when (sType /= sAttrs) $ fail $ unlines $ 
                [ "Attribute stream mismatch for slot (" ++ show slotName ++ ") expected "
                , show sAttrs
                , " but got "
                , show sType
                ]
                
    let slotIdx = case slotName `Map.lookup` slotMap input of
            Nothing -> error $ "internal error (slot index): " ++ show slotName
            Just i  -> i
        seed = objSeed input
    order <- newIORef 0
    enabled <- newIORef True
    index <- readIORef seed
    modifyIORef seed (1+)
    unis <- mkUniform [(n,t) | n <- uniformNames,
                                         let t = fromMaybe (error $ "missing uniform: " ++ n)
                                                   $ Map.lookup n (uniforms sch)]
    cmdsRef <- newIORef (V.singleton V.empty)
    let obj = Object
            { objSlot       = slotIdx
            , objPrimitive  = prim
            , objIndices    = indices
            , objAttributes = attribs
            , objUniSetup   = unis
            , objOrder      = order
            , objEnabled    = enabled
            , objId         = index
            , objCommands   = cmdsRef
            }

    modifyIORef' (slotVector input ! slotIdx) $ \(GLSlot objs _ _) -> GLSlot (IM.insert index obj objs) V.empty Generate

    -- generate GLObjectCommands for the new object
    {-
        foreach pipeline:
            foreach realted program:
                generate commands
    -}
    ppls <- readIORef $ pipelines input
    let topUnis = uniformSetup input
    cmds <- V.forM ppls $ \mp -> case mp of
        Nothing -> return V.empty
        Just p  -> do
            Just ic <- readIORef $ glInput p
            case icSlotMapInputToPipeline ic ! slotIdx of
                Nothing         -> do
                    --putStrLn $ " ** slot is not used!"
                    return V.empty   -- this slot is not used in that pipeline
                Just pSlotIdx   -> do
                    --putStrLn "slot is used!" 
                    --where
                    let emptyV = V.replicate (V.length $ glPrograms p) []
                    return $ emptyV // [(prgIdx,createObjectCommands (glTexUnitMapping p) topUnis obj (glPrograms p ! prgIdx))| prgIdx <- glSlotPrograms p ! pSlotIdx]
    writeIORef cmdsRef cmds
    return obj

removeObject :: GLStorage -> Object -> IO ()
removeObject p obj = modifyIORef (slotVector p ! objSlot obj) $ \(GLSlot !objs _ _) -> GLSlot (IM.delete (objId obj) objs) V.empty Generate

enableObject :: Object -> Bool -> IO ()
enableObject obj b = writeIORef (objEnabled obj) b

setObjectOrder :: GLStorage -> Object -> Int -> IO ()
setObjectOrder p obj i = do
    writeIORef (objOrder obj) i
    modifyIORef (slotVector p ! objSlot obj) $ \(GLSlot objs sorted _) -> GLSlot objs sorted Reorder

uniformSetter :: GLStorage -> Map String InputSetter
uniformSetter = uniformSetup

objectUniformSetter :: Object -> Map GLUniformName InputSetter
objectUniformSetter = objUniSetup

setScreenSize :: GLStorage -> Word -> Word -> IO ()
setScreenSize p w h = writeIORef (screenSize p) (w,h)

sortSlotObjects :: GLStorage -> IO ()
sortSlotObjects p = V.forM_ (slotVector p) $ \slotRef -> do
    GLSlot objMap sortedV ord <- readIORef slotRef
    let cmpFun (a,_) (b,_) = a `compare` b
        doSort objs = do
            ordObjsM <- V.thaw objs
            I.sortBy cmpFun ordObjsM
            ordObjs <- V.freeze ordObjsM
            writeIORef slotRef (GLSlot objMap ordObjs Ordered)
    case ord of
        Ordered -> return ()
        Generate -> do
            objs <- V.forM (V.fromList $ IM.elems objMap) $ \obj -> do
                ord <- readIORef $ objOrder obj
                return (ord,obj)
            doSort objs
        Reorder -> do
            objs <- V.forM sortedV $ \(_,obj) -> do
                ord <- readIORef $ objOrder obj
                return (ord,obj)
            doSort objs

createObjectCommands :: Map String (IORef GLint) -> Map String GLUniform -> Object -> GLProgram -> [GLObjectCommand]
createObjectCommands texUnitMap topUnis obj prg = objUniCmds ++ objStreamCmds ++ [objDrawCmd]
  where
    -- object draw command
    objDrawCmd = case objIndices obj of
        Nothing -> GLDrawArrays prim 0 (fromIntegral count)
        Just (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)
      where
        objAttrs = objAttributes obj
        prim = primitiveToGLType $ objPrimitive obj
        count = head [c | Stream _ _ _ _ c <- Map.elems objAttrs]

    -- object uniform commands
    -- texture slot setup commands
    objUniCmds = uniCmds ++ texCmds
      where
        uniCmds = [GLSetUniform i u | (n,i) <- uniMap, let u = Map.findWithDefault (topUni n) n objUnis]
        uniMap  = Map.toList $ inputUniforms prg
        topUni n = Map.findWithDefault (error $ "internal error (createObjectCommands): " ++ show n) n topUnis
        objUnis = objUniSetup obj
        texUnis = S.toList $ inputTextureUniforms prg
        texCmds = [ GLBindTexture (inputTypeToTextureTarget $ uniInputType u) texUnit u
                  | n <- texUnis
                  , let u = Map.findWithDefault (topUni n) n objUnis
                  , let texUnit = Map.findWithDefault (error $ "internal error (createObjectCommands - Texture Unit): " ++ show n) n texUnitMap
                  ]
        uniInputType (GLTypedUniform ty _) = unwitnessType ty
        uniInputType (GLUniform r)         = objectType r

    -- object attribute stream commands
    objStreamCmds = [attrCmd i s | (i,name) <- Map.elems attrMap, let s = fromMaybe (error $ "missing attribute: " ++ name) $ Map.lookup name objAttrs]
      where 
        attrMap = inputStreams prg
        objAttrs = objAttributes obj
        attrCmd i s = case s of
            Stream ty (Buffer arrs bo) arrIdx start len -> case ty of
                Attribute_Word   -> setIntAttrib 1
                Attribute_V2U    -> setIntAttrib 2
                Attribute_V3U    -> setIntAttrib 3
                Attribute_V4U    -> setIntAttrib 4
                Attribute_Int    -> setIntAttrib 1
                Attribute_V2I    -> setIntAttrib 2
                Attribute_V3I    -> setIntAttrib 3
                Attribute_V4I    -> setIntAttrib 4
                Attribute_Float  -> setFloatAttrib 1
                Attribute_V2F    -> setFloatAttrib 2
                Attribute_V3F    -> setFloatAttrib 3
                Attribute_V4F    -> setFloatAttrib 4
                Attribute_M22F   -> setFloatAttrib 4
                Attribute_M23F   -> setFloatAttrib 6
                Attribute_M24F   -> setFloatAttrib 8
                Attribute_M32F   -> setFloatAttrib 6
                Attribute_M33F   -> setFloatAttrib 9
                Attribute_M34F   -> setFloatAttrib 12
                Attribute_M42F   -> setFloatAttrib 8
                Attribute_M43F   -> setFloatAttrib 12
                Attribute_M44F   -> setFloatAttrib 16
              where
                setFloatAttrib n = GLSetVertexAttribArray i bo n glType (ptr n)
                setIntAttrib n = GLSetVertexAttribIArray i bo n glType (ptr n)
                ArrayDesc arrType arrLen arrOffs arrSize = arrs ! arrIdx
                glType = arrayTypeToGLType arrType
                ptr compCnt   = intPtrToPtr $! fromIntegral (arrOffs + start * fromIntegral compCnt * sizeOfArrayType arrType)

            -- constant generic attribute
            constAttr -> GLSetVertexAttrib i constAttr

newtype UniM a = UniM (ReaderT (Map GLUniformName GLUniform) (Writer [IO ()]) a)
deriving instance Functor UniM
deriving instance Applicative UniM
deriving instance Monad UniM
deriving instance MonadReader (Map String GLUniform) UniM
deriving instance MonadWriter [IO ()] UniM


(@=) :: (Typeable a, Uniformable a) => String -> IO a -> UniM ()
name @= val = do
    u <- do
        us <- ask
        return $ us Map.! name
    case u of

        GLTypedUniform ty ref -> do
            case DMap.lookup ty (uniformContexts val) of
                Just UniformContext -> do
                    tell [val >>= writeIORef ref . GLUniformValue]
                Nothing -> do
                    tell [throwIO $ typeMismatch ty ref]

        GLUniform ref -> case withTypes val ref <$> eqT of
            Just Refl -> tell [val >>= writeIORef ref]
            Nothing -> tell [ Prelude.putStrLn $ "WARNING: "++show (objectType ref)++" variable "
                                                 ++ show name
                                                 ++ " cannot recieve value " ++ show (typeRep val)
                            , throwIO $ typeMismatch ref val
                            ]


updateUniforms :: GLStorage -> UniM a -> IO ()
updateUniforms storage (UniM m) = sequence_ l where
  setters = uniformSetup storage
  l = execWriter $ runReaderT m setters

updateObjectUniforms :: Object -> UniM a -> IO ()
updateObjectUniforms object (UniM m) = sequence_ l where
  setters = objectUniformSetter object
  l = execWriter $ runReaderT m setters

-- | Set a uniform ref.
setGLUniform :: Typeable a =>
    (forall v. Typeable v => TypeTag v -> Maybe (UniformContext a v))
    -> String       -- ^ For warning messages, name of uniform.
    -> GLUniform -- ^ Uniform ref to set.
    -> a         -- ^ Value to store.
    -> IO ()
setGLUniform resolv name u val = case u of
        GLTypedUniform ty ref -> do
            case resolv ty of
                Just UniformContext -> writeIORef ref $ GLUniformValue val
                Nothing -> warn $ unwords [ "Cannot set", show $ unwitnessType ty
                                          , "uniform", name
                                          , "to", show (typeOf val)
                                          , "value." ]

        GLUniform ref -> case withTypes (Just val) ref <$> eqT of
            Just Refl -> writeIORef ref val
            Nothing -> warn $ unwords [ "uniform", name
                                      , "only accepts values of type"
                                      , show $ typeRep ref ]
    where warn s = putStrLn $ "WARNING: " ++ s

-- | Lookup and set a Uniform ref.
setUniformRef :: ( Typeable a
                 , Show name, Ord name
                 ) => (forall v. Typeable v => TypeTag v -> Maybe (UniformContext a v))
                   -> name
                   -> Map name GLUniform
                   -> a
                   -> IO ()
setUniformRef resolv name us val = case Map.lookup name us of
    Nothing -> warn $ "unknown uniform: " ++ show name
    Just u -> setGLUniform resolv (show name) u val
    where warn s = putStrLn $ "WARNING: " ++ s

uniformOf :: (Show name, Ord name, GLData a have) =>
    TypeTag have
    -> name
    -> Map name GLUniform
    -> a
    -> IO ()
uniformOf have = setUniformRef $ knownContext have

uniform :: (Typeable a, Show name, Ord name, Uniformable a) =>
                 name -> Map name GLUniform -> a -> IO ()
uniform n o a = setUniformRef (resolveContext a) n o a


uniformFTexture2D :: SB.ByteString -> Map GLUniformName GLUniform -> TextureData -> IO ()
uniformFTexture2D =
    -- TODO: Check that the uniform is of the expected FTexture2D type.
    uniform . SB.unpack

uniformBool :: (Show name, Ord name) => name -> Map name GLUniform -> Bool -> IO ()
uniformV2B :: (Show name, Ord name, GLData a (GLVector 2 Word32)) => name -> Map name GLUniform -> a -> IO ()
uniformV3B :: (Show name, Ord name, GLData a (GLVector 3 Word32)) => name -> Map name GLUniform -> a -> IO ()
uniformV4B :: (Show name, Ord name, GLData a (GLVector 4 Word32)) => name -> Map name GLUniform -> a -> IO ()
uniformBool = uniformOf TypeBool
uniformV2B = uniformOf TypeV2B
uniformV3B = uniformOf TypeV3B
uniformV4B = uniformOf TypeV4B

uniformWord :: (Show name, Ord name) => name -> Map name GLUniform -> Word32 -> IO ()
uniformV2U :: (Typeable f, GLData (f Word32) (GLVector 2 Word32)) =>
              String -> Map GLUniformName GLUniform -> f Word32 -> IO ()
uniformV3U :: (Typeable f, GLData (f Word32) (GLVector 3 Word32)) =>
              String -> Map GLUniformName GLUniform -> f Word32 -> IO ()
uniformV4U :: (Typeable f, GLData (f Word32) (GLVector 4 Word32)) =>
              String -> Map GLUniformName GLUniform -> f Word32 -> IO ()
uniformWord = uniformOf TypeWord
uniformV2U = uniformOf TypeV2U
uniformV3U = uniformOf TypeV3U
uniformV4U = uniformOf TypeV4U

uniformFloat :: (Show name, Ord name) => name -> Map name GLUniform -> Float -> IO ()
uniformV2F :: (Typeable f, GLData (f Float) (GLVector 2 Float)) =>
              String -> Map GLUniformName GLUniform -> f Float -> IO ()
uniformV3F :: (Typeable f, GLData (f Float) (GLVector 3 Float)) =>
              String -> Map GLUniformName GLUniform -> f Float -> IO ()
uniformV4F :: (Typeable f, GLData (f Float) (GLVector 4 Float)) =>
              String -> Map GLUniformName GLUniform -> f Float -> IO ()
uniformFloat = uniformOf TypeFloat
uniformV2F = setUniformRef (knownContext TypeV2F)
uniformV3F = setUniformRef (knownContext TypeV3F)
uniformV4F = setUniformRef (knownContext TypeV4F)

uniformInt :: (Show name, Ord name) => name -> Map name GLUniform -> Int32 -> IO ()
uniformV2I :: (Typeable f, GLData (f Int32) (GLVector 2 Int32)) =>
              String -> Map GLUniformName GLUniform -> f Int32 -> IO ()
uniformV3I :: (Typeable f, GLData (f Int32) (GLVector 3 Int32)) =>
              String -> Map GLUniformName GLUniform -> f Int32 -> IO ()
uniformV4I :: (Typeable f, GLData (f Int32) (GLVector 4 Int32)) =>
              String -> Map GLUniformName GLUniform -> f Int32 -> IO ()
uniformInt = uniformOf TypeInt
uniformV2I = uniformOf TypeV2I
uniformV3I = uniformOf TypeV3I
uniformV4I = uniformOf TypeV4I

{-
Note: This works to infer the type Float for literals without fixing the matrix type:

    type family MatrixComponent m where
        MatrixComponent (f (g c)) = c
        MatrixComponent (f c) = c

    uniformM44F :: ( MatrixComponent a ~ Float , GLData a (GLMatrix 4 4 Float))
                    => String -> Map String GLUniform -> a -> IO ()

However, it breaks the ability to partially apply without a type signature.
Therefore, I'm forcing LambdaCube's internal matrix types for uniformM*
functions.
-}
uniformM22F :: (Show name, Ord name) => name -> Map name GLUniform -> M22F -> IO ()
uniformM23F :: (Show name, Ord name) => name -> Map name GLUniform -> M23F -> IO ()
uniformM24F :: (Show name, Ord name) => name -> Map name GLUniform -> M24F -> IO ()
uniformM32F :: (Show name, Ord name) => name -> Map name GLUniform -> M32F -> IO ()
uniformM33F :: (Show name, Ord name) => name -> Map name GLUniform -> M33F -> IO ()
uniformM34F :: (Show name, Ord name) => name -> Map name GLUniform -> M34F -> IO ()
uniformM42F :: (Show name, Ord name) => name -> Map name GLUniform -> M42F -> IO ()
uniformM43F :: (Show name, Ord name) => name -> Map name GLUniform -> M43F -> IO ()
uniformM44F :: (Show name, Ord name) => name -> Map name GLUniform -> M44F -> IO ()
uniformM22F = uniformOf TypeM22F
uniformM23F = uniformOf TypeM23F
uniformM24F = uniformOf TypeM24F
uniformM32F = uniformOf TypeM32F
uniformM33F = uniformOf TypeM33F
uniformM34F = uniformOf TypeM34F
uniformM42F = uniformOf TypeM42F
uniformM43F = uniformOf TypeM43F
uniformM44F = uniformOf TypeM44F