module Shader.Poke (
    pokeVerxs,
    pokeLayVerxs,
    pokeArrayOff,
    pokeShape,
    pokeWallsWindowsFloor,
    memoTopPrismEdgeIndices,
) where

import Control.Monad.Primitive
import qualified Data.Vector as V
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as UMV
import Foreign
import Geometry.Data
import Geometry.Polygon
import Graphics.GL.Core45
import Linear.V3 (cross)
import Picture.Data
import Shader.Data
import Shader.Parameters
import Shape.Data

pokeVerxs ::
    MV.MVector (PrimState IO) (FullShader, VBO) ->
    UMV.MVector (PrimState IO) Int ->
    Picture ->
    IO ()
pokeVerxs vbos count = VFSM.mapM_ (pokeVerx vbos count) . VFSM.fromList

pokeVerx :: MV.MVector (PrimState IO) (FullShader, VBO) -> UMV.MVector (PrimState IO) Int -> Verx -> IO ()
pokeVerx vbos offsets Verx{_vxPos = thePos, _vxCol = theCol, _vxExt = ext, _vxShadNum = theShadNum} = do
    typeOff <- UMV.unsafeRead offsets sn
    basePtr <- _vboPtr . snd <$> MV.unsafeRead vbos sn
    let thePtr = plusPtr basePtr (typeOff * pokeStride sn * floatSize)
    poke34 thePtr thePos theCol
    pokeArrayOff thePtr 7 ext
    UMV.unsafeModify offsets (+ 1) sn
  where
    sn = _unShadNum theShadNum

pokeWallsWindowsFloor ::
    Ptr Float ->
    Ptr Float ->
    Ptr Float ->
    [((Point2, Point2), Point4)] ->
    [((Point2, Point2), Point4)] ->
    [(Point3, Point3)] ->
    IO (Int, Int, Int)
pokeWallsWindowsFloor wlptr wiptr flptr wls wis fls = do
    wlcounts1 <- VFSM.foldlM' (pokeW wlptr) 0 (VFSM.fromList wls)
    wlcounts2 <- VFSM.foldlM' (pokeW wiptr) 0 (VFSM.fromList wis)
    flcounts <- VFSM.foldlM' (pokeF flptr) 0 (VFSM.fromList fls)
    return (wlcounts1, wlcounts2, flcounts)

pokeF :: Ptr Float -> Int -> (Point3, Point3) -> IO Int
pokeF ptr i' (V3 a b c, V3 d e f) = do
    let i = i' * 6
    pokeElemOff ptr (i + 0) a
    pokeElemOff ptr (i + 1) b
    pokeElemOff ptr (i + 2) c
    pokeElemOff ptr (i + 3) d
    pokeElemOff ptr (i + 4) e
    pokeElemOff ptr (i + 5) f
    return $ i' + 1

pokeW :: Ptr Float -> Int -> ((Point2, Point2), Point4) -> IO Int
pokeW ptr i' ((V2 a b, V2 c d), V4 e f g h) = do
    let i = i' * 8
    pokeElemOff ptr (i + 0) a
    pokeElemOff ptr (i + 1) b
    pokeElemOff ptr (i + 2) c
    pokeElemOff ptr (i + 3) d
    pokeElemOff ptr (i + 4) e
    pokeElemOff ptr (i + 5) f
    pokeElemOff ptr (i + 6) g
    pokeElemOff ptr (i + 7) h
    return $ i' + 1

pokeShape ::
    Ptr Float ->
    Ptr GLushort ->
    Ptr GLushort ->
    (Int, Int, Int) ->
    [Surface] ->
    IO (Int, Int, Int)
pokeShape ptr iptr ieptr is = VFSM.foldlM' (pokeShapeObj ptr iptr ieptr) is . VFSM.fromList

pokeShapeObj ::
    Ptr Float ->
    Ptr GLushort ->
    Ptr GLushort ->
    (Int, Int, Int) ->
    Surface ->
    IO (Int, Int, Int)
{-# INLINE pokeShapeObj #-}
pokeShapeObj ptr iptr ieptr counts (Surface shtype shVerts col) = case shtype of
    TopPrism size -> pokeTopPrism midp col size ptr iptr ieptr counts (VFSM.fromList shVerts)
    FlatFaces size -> pokeBox col size ptr iptr ieptr counts shVerts
    RoundedFaces size -> pokeRoundedFaces col size ptr iptr ieptr counts shVerts
    Cylinder size -> pokeCylinder col size ptr iptr ieptr counts shVerts
  where
    midp = centroidNum shVerts

pokeRoundedFaces ::
    Point4 ->
    Int ->
    Ptr Float ->
    Ptr GLushort ->
    Ptr GLushort ->
    (Int, Int, Int) ->
    [Point3] ->
    IO (Int, Int, Int)
pokeRoundedFaces col size ptr iptr ieptr (nv, nsi, nei) (tc : bc : svs) = do
    nv' <- pokeRoundedCurve col ptr tc bc svs nv
    nsi' <-
        UV.foldM'
            (pokeIndex nv iptr)
            nsi
            (memoTopPrismIndices V.! (size - 3))
    nei' <-
        UV.foldM' (pokeIndex nv ieptr) nei $
            memoTopPrismEdgeIndices V.! (size - 3)
    return (nv', nsi', nei')
pokeRoundedFaces _ _ _ _ _ _ _ = undefined

pokeCylinder ::
    Point4 ->
    Int ->
    Ptr Float ->
    Ptr GLushort ->
    Ptr GLushort ->
    (Int, Int, Int) ->
    [Point3] ->
    IO (Int, Int, Int)
pokeCylinder col size ptr iptr ieptr (nv, nsi, nei) (tc : bc : svs) = do
    nv' <- pokeRoundedCurve col ptr tc bc svs nv >>= pokeCylinderCaps col ptr tc bc svs
    nsi' <-
        UV.foldM'
            (pokeIndex nv iptr)
            nsi
            (memoCylinderIndices V.! (size - 3))
    nei' <-
        UV.foldM' (pokeIndex nv ieptr) nei $
            memoTopPrismEdgeIndices V.! (size - 3)
    return (nv', nsi', nei')
pokeCylinder _ _ _ _ _ _ _ = undefined

pokeRoundedCurve :: Point4 -> Ptr Float -> Point3 -> Point3 -> [Point3] -> Int -> IO Int
pokeRoundedCurve col ptr tc bc = go True
  where
    go True (x : xs) n = pokeJustV tc col ptr n x >>= go False xs
    go False (x : xs) n = pokeJustV bc col ptr n x >>= go True xs
    go _ [] n = return n

pokeCylinderCaps :: Point4 -> Ptr Float -> Point3 -> Point3 -> [Point3] -> Int -> IO Int
pokeCylinderCaps col ptr tc bc = go True
  where
    go True (x : xs) n = pokeJustVInvNormal tc col ptr n x >>= go False xs
    go False (x : xs) n = pokeJustVInvNormal bc col ptr n x >>= go True xs
    go _ [] n = return n

pokeBox ::
    Point4 ->
    Int ->
    Ptr Float ->
    Ptr GLushort ->
    Ptr GLushort ->
    (Int, Int, Int) ->
    [Point3] ->
    IO (Int, Int, Int)
{-# INLINE pokeBox #-}
pokeBox col size ptr iptr ieptr (nv, nsi, nei) svs = do
    nv' <- VFSM.foldM' (pokeBoxSurface col ptr svsv) nv $ VFSM.fromList $ boxSurfaces size
    nsi' <-
        UV.foldM'
            (pokeIndex nv iptr)
            nsi
            (memoFlatIndices V.! (size -3))
    nei' <-
        UV.foldM' (pokeIndex nv ieptr) nei $
            memoBoxEdgeIndices V.! (size - 3)
    return (nv', nsi', nei')
  where
    svsv = UV.fromList svs

-- should probably use a vector of Point3
pokeBoxSurface :: Point4 -> Ptr Float -> UV.Vector Point3 -> Int -> [Int] -> IO Int
pokeBoxSurface col ptr vs n is =
    UV.foldM'
        (pokeFlatV norm col ptr)
        n
        v
  where
    v = UV.backpermute vs (UV.fromList is)
    x = v UV.! 0
    y = v UV.! 1
    z = v UV.! 2
    norm = negate $ cross (x - y) (z - y)

-- should probably memoize this
boxSurfaces :: Int -> [[Int]]
boxSurfaces size =
    [0, 2 .. size * 2 -1] :
    [1, 3 .. size * 2] :
    map f [0, 2 .. size * 2 - 1]
  where
    f x = map (`mod` (size * 2)) [x, x + 1, x + 3, x + 2]

boxSurfaces' :: Int -> [[Int]]
boxSurfaces' n =
    [0 .. n -1] :
    reverse [n .. n * 2 -1] :
    [ map ((2 * n) +) [4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3] | i <- [0 .. n -1]]

pokeTopPrism ::
    Point3 ->
    Point4 ->
    Int ->
    Ptr Float ->
    Ptr GLushort ->
    Ptr GLushort ->
    (Int, Int, Int) ->
    VFSM.Stream IO Point3 ->
    IO (Int, Int, Int)
{-# INLINE pokeTopPrism #-}
pokeTopPrism cp col size ptr iptr ieptr (nv, nsi, nei) svs = do
    nv' <- VFSM.foldlM' (pokeJustV cp col ptr) nv svs
    nsi' <-
        UV.foldM'
            (pokeIndex nv iptr)
            nsi
            (memoTopPrismIndices V.! (size - 3))
    nei' <-
        UV.foldM'
            (pokeIndex nv ieptr)
            nei
            (memoTopPrismEdgeIndices V.! (size - 3))
    return (nv', nsi', nei')

pokeIndex ::
    -- | base index
    Int ->
    Ptr GLushort ->
    -- | number poked
    Int ->
    -- | index offset
    Int ->
    IO Int
{-# INLINE pokeIndex #-}
pokeIndex nv eiptr ni ioff = do
    pokeElemOff eiptr ni (fromIntegral $ nv + ioff)
    return $ ni + 1

--pokeTopPrismIndex ::
--    Int ->
--    Ptr GLushort ->
--    Int ->
--    Int ->
--    IO Int
--{-# INLINE pokeTopPrismIndex #-}
--pokeTopPrismIndex nv iptr nshapeindices ioff = do
--    pokeElemOff iptr nshapeindices (fromIntegral $ nv + ioff)
--    return $ nshapeindices + 1

triangulate :: [Int] -> [Int]
triangulate is = V.toList . V.backpermute (V.fromList is) . V.fromList $ triangulateIndices (length is)

triangulateIndices :: Int -> [Int]
triangulateIndices i = concatMap f [0 .. i -3]
  where
    f x
        | even x = [0, x + 1, x + 2]
        | otherwise = [0, x + 2, x + 1]

memoFlatIndices :: V.Vector (UV.Vector Int)
memoFlatIndices =
    V.generate 10 $
        UV.fromList . concatMap triangulate . boxSurfaces' . (+ 3)

memoTopPrismIndices :: V.Vector (UV.Vector Int)
memoTopPrismIndices =
    V.generate 10 $
        UV.fromList . topPrismIndices . (+ 3)

memoCylinderIndices :: V.Vector (UV.Vector Int)
memoCylinderIndices =
    V.generate 10 $
        UV.fromList . cylinderIndices . (+ 3)

memoTopPrismEdgeIndices :: V.Vector (UV.Vector Int)
memoTopPrismEdgeIndices =
    V.generate 10 $
        UV.fromList . topPrismEdgeIndices . (+ 3)

topPrismEdgeIndices ::
    -- | the number of vertices on the top surface, ie. half the total
    Int ->
    [Int]
topPrismEdgeIndices n = concatMap f [0 .. n -1]
  where
    f i = map g
            [ 0 , 2 ,  1 , 4 
            , 0 , 1 , -2 , 3 
            , 1 , 3 , -1 , 2 
            ]
      where
        g j = (2 * i + j) `mod` (2 * n)

memoBoxEdgeIndices :: V.Vector (UV.Vector Int)
memoBoxEdgeIndices =
    V.generate 10 $
        UV.fromList . boxEdgeIndices . (+ 3)

boxEdgeIndices ::
    -- | the number of vertices on the top surface, ie. half the total
    Int ->
    [Int]
boxEdgeIndices n = concatMap f [0 .. n -1]
  where
    f i =
        map
            g
            [ 0 , 1 , n , 2 
            , 0 , n , -1 , n + 1
            , n , n + 1 , n -1 , 1
            ]
      where
        g j = (i + j) `mod` (2 * n)

cylinderIndices :: Int -> [Int]
cylinderIndices n =
    cylinderRoundIndices n
        ++ triangulate [2 * n, 2 * n + 2 .. 4 * n - 1]
        ++ triangulate [2 * n + 1, 2 * n + 3 .. 4 * n - 1]

cylinderRoundIndices :: Int -> [Int]
cylinderRoundIndices n =
    [ 2 * n -2 , 2 * n -1 , 1
    , 2 * n -2 , 1 , 0 -- last side triangle (applies mod 2n)
    ]
        ++ concatMap g [0 .. n -2] -- other triangles on sides
  where
    g x =
        [ 2 * x , 2 * x + 1 , 2 * x + 3
        , 2 * x , 2 * x + 3 , 2 * x + 2
        ]

topPrismIndices :: Int -> [Int]
topPrismIndices n =
    concatMap f [1 .. n -2] -- triangles on top face
        ++ concatMap f' [1 .. n -2] -- triangles on bottom face
        -- these should be checked
        ++ [ 2 * n -2 , 2 * n -1 , 1
           , 2 * n -2 , 1 , 0 -- last side triangle (applies mod 2n)
           ]
        ++ concatMap g [0 .. n -2] -- other triangles on sides
  where
    f x = [0, 2 * x, 2 * x + 2]
    f' x = [1, 2 * x + 3, 2 * x + 1]
    g x =
        [ 2 * x , 2 * x + 1 , 2 * x + 3
        , 2 * x , 2 * x + 3 , 2 * x + 2
        ]

-- consider changing the position to a vec4
-- and just doing two pokes rather than seven
-- (especially if adding normal data)
pokeJustV ::
    Point3 ->
    Point4 ->
    Ptr Float ->
    Int ->
    Point3 ->
    IO Int
{-# INLINE pokeJustV #-}
pokeJustV cp col ptr nv sh = do
    UV.imapM_ f $ UV.fromList [x, y, z, 1, r, g, b, a, nx, ny, nz, 1]
    return (nv + 1)
  where
    f i = pokeElemOff ptr (nv * nShapeVerxComp + i)
    V3 x y z = sh
    V4 r g b a = col
    V3 nx ny nz = cp

pokeJustVInvNormal ::
    Point3 ->
    Point4 ->
    Ptr Float ->
    Int ->
    Point3 ->
    IO Int
{-# INLINE pokeJustVInvNormal #-}
pokeJustVInvNormal cp col ptr nv sh = do
    UV.imapM_ f $ UV.fromList [x, y, z, 1, r, g, b, a, nx, ny, nz, 1]
    return (nv + 1)
  where
    f i = pokeElemOff ptr (nv * nShapeVerxComp + i)
    V3 x y z = sh
    V4 r g b a = col
    V3 nx ny nz = (2 * sh) - cp

pokeFlatV ::
    Point3 ->
    Point4 ->
    Ptr Float ->
    Int ->
    Point3 ->
    IO Int
{-# INLINE pokeFlatV #-}
pokeFlatV norm col ptr nv sh = do
    UV.imapM_ f $ UV.fromList [x, y, z, 1, r, g, b, a, nx, ny, nz, 1]
    return (nv + 1)
  where
    f i = pokeElemOff ptr (nv * nShapeVerxComp + i)
    V3 x y z = sh
    V4 r g b a = col
    V3 nx ny nz = sh - norm

pokeLayVerxs ::
    MV.MVector (PrimState IO) (FullShader, VBO) ->
    UMV.MVector (PrimState IO) Int ->
    Picture ->
    IO ()
pokeLayVerxs vbos counts = VFSM.mapM_ (pokeLayVerx vbos counts) . VFSM.fromList

pokeLayVerx :: MV.MVector (PrimState IO) (FullShader, VBO) -> UMV.MVector (PrimState IO) Int -> Verx -> IO ()
{-# INLINE pokeLayVerx #-}
pokeLayVerx vbos counts vx = do
    theOff <- UMV.unsafeRead counts vecPos
    basePtr <- _vboPtr . snd <$> MV.unsafeRead vbos sn
    let thePtr = plusPtr basePtr ((theOff + layOff) * theStride * floatSize)
    poke34 thePtr thePos theCol
    pokeArrayOff thePtr 7 (_vxExt vx)
    UMV.unsafeModify counts (+ 1) vecPos
  where
    sn = _unShadNum (_vxShadNum vx)
    vecPos = theLayer * numLayers + sn
    theLayer = layerNum $ _vxLayer vx
    thePos = _vxPos vx
    theCol = _vxCol vx
    layOff = theLayer * numSubElements
    theStride = pokeStride sn

pokeStride :: Int -> Int
{-# INLINE pokeStride #-}
pokeStride 0 = 7
pokeStride 1 = 8
pokeStride 2 = 11
pokeStride 3 = 9
pokeStride 4 = 10
pokeStride 5 = 7
pokeStride _ = undefined

poke34 :: Ptr Float -> Point3 -> Point4 -> IO ()
{-# INLINE poke34 #-}
poke34 ptr (V3 a b c) (V4 d e f g) = do
    pokeElemOff ptr 0 a
    pokeElemOff ptr 1 b
    pokeElemOff ptr 2 c
    pokeElemOff ptr 3 d
    pokeElemOff ptr 4 e
    pokeElemOff ptr 5 f
    pokeElemOff ptr 6 g

pokeArrayOff :: Ptr Float -> Int -> [Float] -> IO ()
{-# INLINE pokeArrayOff #-}
pokeArrayOff ptr i = pokeArray (plusPtr ptr (floatSize * i))