{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
module MeshSketch where

import Codec.Picture             as Juicy
import Control.Concurrent
import Control.Monad
import Data.Word
import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Int
import Data.IORef
import Data.Text (Text)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Data.Vector     as V
import Foreign.Marshal.Array
import Foreign.Storable
import GI.Gdk
import GI.Gdk.Objects
import GI.GLib.Constants
import GI.Gtk hiding (IsWindow,windowFullscreen,windowUnfullscreen)
import LambdaCube.GL             as LC
import LambdaCube.GL.Mesh        as LC
import Numeric.LinearAlgebra hiding ((<>))
import System.Environment
import System.IO
import System.IO.Error
import Control.Exception
import LambdaCube.GL as LC
import LambdaCube.IR as LC
import LambdaCube.Gtk
import LambdaCube.GL.Data (uploadCubeMapToGPU,uploadTextureBufferToGPU,updateTextureBuffer)
import LambdaCube.GL.Type (TextureCubeData(..),Object(..))
-- import Text.Show.Pretty (ppShow)
import qualified Graphics.Rendering.OpenGL as GL
import Data.Char
import Text.Printf

import CubeMap
import GLWidget (nullableContext, withCurrentGL)
import LambdaCube.GL.HMatrix
import LambdaCubeWidget (loadPipeline,DynamicPipeline(..))
import Animator
import LoadMesh
import InfinitePlane
import MtlParser (ObjMaterial(..))
import Matrix
import TextureBufferRing


prettyDebug :: GL.DebugMessage -> String
prettyDebug (GL.DebugMessage src typ (GL.DebugMessageID mid) severity msg) = unwords ws
 where
    ws = [wsrc,wtyp,wmid,wseverity,msg]
    -- DebugSourceShaderCompiler DebugTypeOther 1 DebugSeverityNotification
    wsrc = filter isUpper $ drop 11 $ show src
    wtyp = take 2 $ drop 9 $ show typ
    wmid = printf "%03i" mid
    wseverity = drop 13 $ show severity

setupGLDebugging :: IO ()
setupGLDebugging = do
    let pdebug m@(GL.DebugMessage src typ mid severity msg) = do
            putStrLn (">> " ++ prettyDebug m)
    GL.debugOutput GL.$= GL.Enabled
    GL.debugOutputSynchronous GL.$= GL.Enabled
    GL.debugMessageControl (GL.MessageGroup Nothing Nothing Nothing) GL.$= GL.Enabled
    GL.debugMessageCallback GL.$= Just pdebug


-- State created by uploadState.
data State = State
    { stAnimator   :: Animator
    , stCamera     :: IORef Camera
    , stFullscreen :: IO ()
    , stSkyboxes   :: Skyboxes
    , stSkybox     :: IORef Int
    , stSkyTexture :: IORef TextureCubeData
    , stDragFrom   :: IORef (Maybe (Vector Float,Camera))
    , stRingBuffer :: Ring
    , stPenDown    :: IORef Bool
    }

data Camera = Camera
    { camHeightAngle   :: Float
    , camTarget        :: Vector Float
    , camDirection     :: Vector Float
    , camDistance      :: Float
    , camWidth         :: Float
    , camHeight        :: Float
    , camUp            :: Vector Float
    , camWorldToScreen :: Maybe (Matrix Float)
    , camScreenToWorld :: Maybe (Matrix Float)
    }

initCamera :: Camera
initCamera = Camera
    { camHeightAngle = pi/6
    , camTarget      = fromList [0,0,0]
    , camDirection   = scale (1/d) $ fromList [-2,-2,-10]
    , camDistance    = d
    , camWidth       = 700
    , camHeight      = 700
    , camUp          = fromList [0,1,0]
    , camWorldToScreen = Nothing
    , camScreenToWorld = Nothing
    }
    where d = realToFrac $ norm_2 $ fromList [2::Float,2,10]

viewProjection :: Camera -> (Camera,(Matrix Float,Vector Float))
viewProjection c
    | Just m <- camWorldToScreen c = (c,(m,pos))
    | otherwise                    = (c { camWorldToScreen = Just m' }, (m',pos))
 where
    m' = proj <> cam
    cam = lookat pos (camTarget c) (camUp c)
    pos = camTarget c - scale (camDistance c) (camDirection c)
    proj = perspective 0.1 100 (camHeightAngle c) (camWidth c / camHeight c)


addOBJToObjectArray :: GLStorage -> String -> [(GPUMesh, Maybe Text)] -> Map Text (ObjMaterial,TextureData) -> IO [LC.Object]
addOBJToObjectArray storage slotName objMesh mtlLib = forM objMesh $ \(mesh,mat) -> do
  obj <- LC.addMeshToObjectArray storage slotName ["diffuseTexture","diffuseColor"] mesh
         -- diffuseTexture and diffuseColor values can change on each model
  case mat >>= flip Map.lookup mtlLib of
    Nothing -> return ()
    Just (ObjMaterial{..},t) -> LC.updateObjectUniforms obj $ do
      "diffuseTexture" @= return t -- set model's diffuse texture
      "diffuseColor" @= let (r,g,b) = mtl_Kd in return (V4 r g b mtl_Tr)
  return obj

mkFullscreenToggle :: IsWindow a => a -> IO (IO ())
mkFullscreenToggle w = do
    full <- newIORef False
    return $ do
        b <- atomicModifyIORef' full $ \b -> (not b, not b)
        if b then windowFullscreen w
             else windowUnfullscreen w


uploadState :: IsWidget glarea => MeshData -> glarea -> GLStorage -> IO State
uploadState obj glarea storage = do
    -- load OBJ geometry and material descriptions
    (objMesh,mtlLib) <- uploadOBJToGPU obj
    -- load materials textures
    gpuMtlLib <- uploadMtlLib mtlLib
    -- add OBJ to pipeline input
    addOBJToObjectArray storage "objects" objMesh gpuMtlLib
    -- grid plane
    uploadMeshToGPU xzplane >>= addMeshToObjectArray storage "plane" []

    ring <- newRing storage 100

    -- setup FrameClock
    w <- toWidget glarea
    tm <- newAnimator w
    cam <- newIORef initCamera

    Just pwidget <- get w #parent
    Just parent <- get pwidget #window
    toggle <- mkFullscreenToggle parent
    skyboxes <- loadSkyboxes
    skybox <- newIORef 0
    skybox_id <- skyboxLoad skyboxes 0 >>= \case
        Right ts -> do
            skybox_id <- uploadCubeMapToGPU ts
            LC.updateUniforms storage $ do
               "CubeMap" @= return skybox_id
            return skybox_id
        Left msg -> do
            putStrLn msg
            return (TextureCubeName 0)
    skytex <- newIORef skybox_id
    mi <- LC.uploadMeshToGPU cubeMesh
    LC.addMeshToObjectArray storage "SkyCube" [] mi

    drag <- newIORef Nothing
    pendown <- newIORef False

    let st = State
            { stAnimator     = tm
            , stCamera       = cam
            , stFullscreen   = toggle
            , stSkyboxes     = skyboxes
            , stSkybox       = skybox
            , stSkyTexture   = skytex
            , stDragFrom     = drag
            , stRingBuffer   = ring
            , stPenDown      = pendown
            }
    -- _ <- addAnimation tm (whirlingCamera st)

    return st


destroyState :: GLArea -> State -> IO ()
destroyState glarea st = do
    -- widgetRemoveTickCallback glarea (stTickCallback st)
    return ()

deg30 :: Float
deg30 = pi/6

whirlingCamera :: State -> Animation
whirlingCamera st = Animation $ \_ t -> do
    let tf = realToFrac t :: Float
        rot = rotMatrixZ (-tf/2) <> rotMatrixX (-tf/pi)
    modifyIORef (stCamera st) $ \cam -> cam
                            { camUp            = fromList [0,1,0] <# rot
                            , camDirection     = (scale (1/camDistance cam) $ fromList [-2,-2,-10]) <# rot
                            , camWorldToScreen = Nothing
                            , camScreenToWorld = Nothing
                            }
    return $ Just (whirlingCamera st)

setUniforms :: glctx -> GLStorage -> State -> IO ()
setUniforms gl storage st = do
    (mvp,pos) <- atomicModifyIORef' (stCamera st) viewProjection
    LC.updateUniforms storage $ do
      "CameraPosition" @= return (pos :: Vector Float)
      "ViewProjection" @= return (mvp :: Matrix Float)
    updateRingUniforms storage (stRingBuffer st)

data MeshSketch = MeshSketch
    { mmWidget   :: GLArea
    , mmRealized :: IORef (Maybe Realized)
    }

data Realized = Realized
    { stStorage  :: GLStorage
    , stRenderer :: GLRenderer
    , stState    :: State
    }

new :: IO GLArea
new = do
    m <- do
        objName <- head . (++ ["cube.obj"]) <$> getArgs
        mobj <- loadOBJ objName
        -- mpipeline <- (\s -> return (Right (DynamicPipeline savedPipeline (makeSchema s)))) $ do
        mpipeline <- loadPipeline "hello_obj2.json" $ do
          defObjectArray "SkyCube" Triangles $ do
            "position"  @: Attribute_V3F
          defObjectArray "objects" Triangles $ do
            "position"  @: Attribute_V4F
            "normal"    @: Attribute_V3F
            "uvw"       @: Attribute_V3F
          defObjectArray "plane" Triangles $ do
            "position"  @: Attribute_V4F
          defObjectArray "Points" Points $ do
            "position"  @: Attribute_Float
          defUniforms $ do
            "PointBuffer"     @: FTextureBuffer
            "CubeMap"         @: FTextureCube
            "CameraPosition"  @: V3F
            "ViewProjection"  @: M44F
            "PointsMax"       @: Int
            "PointsStart"     @: Int
            "diffuseTexture"  @: FTexture2D
            "diffuseColor"    @: V4F
        return $ (,) <$> mobj <*> mpipeline
    either (\e _ -> hPutStrLn stderr e >> throwIO (userError e)) (&) m $ \(obj,pipeline) -> do

        -- putStrLn $ ppShow (dynamicPipeline pipeline)

        ref <- newIORef Nothing
        -- glarea <- newGLWidget return (lambdaRender app glmethods)
        do
            g <- gLAreaNew
            let mm = MeshSketch g ref
            gLAreaSetHasDepthBuffer g True
            st <- return g
            _ <- on g #realize       $ withCurrentGL g (onRealize obj (dynamicPipeline pipeline) (dynamicSchema pipeline) mm)
            _ <- on g #unrealize     $ onUnrealize mm
            -- _ <- on g #createContext $ nullableContext (glCreateContext w st)
            return g

onUnrealize :: MeshSketch -> IO ()
onUnrealize mm = do
    m <- readIORef (mmRealized mm)
    forM_ m $ \st -> do
        LC.disposeStorage (stStorage st)
        LC.disposeRenderer (stRenderer st)
        -- lcDestroyState lc x

onRealize :: MeshData -> Pipeline -> PipelineSchema -> MeshSketch -> IO ()
onRealize mesh pipeline schema mm = do
    onUnrealize mm
    setupGLDebugging
    storage <- LC.allocStorage schema
    renderer <- LC.allocRenderer pipeline
    compat <- LC.setStorage renderer storage -- check schema compatibility
    x <- uploadState mesh (mmWidget mm) storage
    let r = Realized
            { stStorage  = storage
            , stRenderer = renderer
            , stState    = x
            }
        w = mmWidget mm
    set w [ #canFocus := True ] -- For keyboard events.
    widgetAddEvents w
        [ EventMaskPointerMotionMask
        , EventMaskButtonPressMask
        , EventMaskButtonReleaseMask
        , EventMaskTouchMask
        , EventMaskScrollMask
        , EventMaskKeyPressMask -- , EventMaskKeyReleaseMask
        ]
    _ <- on w #event  $ onEvent  w r
    _ <- on w #render $ onRender w r
    _ <- on w #resize $ onResize w r
    writeIORef (mmRealized mm) $ Just r

onRender :: w -> Realized -> GLContext -> IO Bool
onRender w realized gl = do
    r <- fixupRenderTarget (stRenderer realized)
    setUniforms gl (stStorage realized) (stState realized)
    LC.renderFrame r
    return True

onResize :: GLArea -> Realized -> Int32 -> Int32 -> IO ()
onResize glarea realized w h = do
    let storage = stStorage realized
    -- Plenty of options here.  I went with the last one.
    -- 1. gLContextGetWindow :: HasCallStack => GLContext -> IO (Maybe Window)
    -- 2. getGLContextWindow ::                 GLContext -> IO (Maybe Window)
    -- 3. widgetGetWindow    :: HasCallStack => GLArea    -> IO (Maybe Window)
    widgetGetWindow glarea >>= mapM_ (\win -> do
        (wd,ht) <- do wd <- windowGetWidth win
                      ht <- windowGetHeight win
                      return (fromIntegral wd,fromIntegral ht)
        modifyIORef' (stCamera $ stState realized)
                  $ \c -> c { camWidth  = fromIntegral wd
                            , camHeight = fromIntegral ht
                            , camWorldToScreen = Nothing
                            , camScreenToWorld = Nothing
                            }
        LC.setScreenSize (stStorage realized) wd ht)

-- This computes a point in world coordinates on the view screen if
-- we assume the camera is located at the origin.
computeDirection :: Camera -> Double -> Double -> Vector Float
computeDirection cam h k =
    let d̂ = camDirection cam -- forward
        û = camUp cam        -- upward
        r̂ = d̂ `cross` û      -- rightward
        xr = realToFrac h - (camWidth cam / 2)
        xu = (camHeight cam / 2) - realToFrac k
        xd = (camHeight cam / 2) / tan (camHeightAngle cam / 2)
    in scale xr r̂ + scale xu û + scale xd d̂

rotate :: Float -> Vector Float -> Matrix Float
rotate cosθ u = (3><3)
  [ cosθ + ux² mcosθ       , (uy.uy)mcosθ - uz sinθ , (ux.uz)mcosθ + uy sinθ
  , (uy.ux)mcosθ + uz sinθ , cosθ + uy² mcosθ       , (uy.uz)mcosθ - ux sinθ
  , (uz.ux)mcosθ - uy sinθ , (uz.uy)mcosθ + ux sinθ , cosθ + uz² mcosθ
  ]
 where
    sinθ = sqrt (1 - cosθ * cosθ)
    mcosθ = 1 - cosθ
    û = scale (1/realToFrac (norm_2 u)) u
    ux a = (û!0) * a
    uy a = (û!1) * a
    uz a = (û!2) * a
    ux² = ux . ux
    uy² = uy . uy
    uz² = uz . uz

updateCameraRotation w st h k = do
    m <- readIORef (stDragFrom st)
    forM_ m $ \(df,cam) -> do
        let d̂ = camDirection cam -- forward
            û = camUp cam        -- upward
            r̂ = d̂ `cross` û      -- rightward
            -- fr = df `dot` r̂
            -- fu = df `dot` û
            -- fd = df `dot` d̂
            dt = computeDirection cam h k
            -- tr = dt `dot` r̂
            -- tu = dt `dot` û
            -- td = dt `dot` d̂
            cosθ = dot df dt / realToFrac (norm_2 df) / realToFrac (norm_2 dt)
            axis0 = df `cross` dt
            small x = abs x < 0.00001
            axis = let xs = toList axis0
                    in if any isNaN xs || all small xs
                        then fromList [0,1,0]
                        else axis0
            cam' = cam
                    { camDirection     = d̂ <# rotate cosθ axis
                    , camUp            = û <# rotate cosθ axis
                    , camWorldToScreen = Nothing
                    , camScreenToWorld = Nothing
                    }
        writeIORef (stCamera st) cam'
        mwin <- widgetGetWindow w
        forM_ mwin $ \win ->
            windowInvalidateRect win Nothing False

sanitizeCamera st = do
    modifyIORef (stCamera st) $ \cam ->
        let d = camDirection cam
            u = camUp cam
            dd = norm_2 d
            uu = norm_2 u
            e = scale (realToFrac $ 1/dd) d
            d̂ = if any isNaN (toList e)
                    then fromList [0,0,-1]
                    else e
            f = scale (realToFrac $ 1/uu) u
            û = if any isNaN (toList f)
                    then fromList [0,1,0]
                    else f
        in cam
            { camDirection = d̂
            , camUp        = û
            , camWorldToScreen = Nothing
            , camScreenToWorld = Nothing
            }


onEvent :: IsWidget w => w -> Realized -> Event -> IO Bool
onEvent w realized ev = do
    let st = stState realized
    msrc <- eventGetSourceDevice ev
    inputSource <- forM msrc $ \src -> do
        src <- get src #inputSource
        return src
    etype <- get ev #type
    -- putStrLn $ "onEvent! " ++ show (etype,inputSource)
    let put x = putStrLn (show inputSource ++ " " ++ show x)
    let st = stState realized
    case etype of

        EventTypeMotionNotify -> do
            case inputSource of
                Just InputSourcePen -> do
                    isDown <- readIORef (stPenDown st)
                    when isDown $ do
                        mev <- get ev #motion
                        h <- get mev #x
                        k <- get mev #y
                        cam <- readIORef (stCamera st)
                        let d = computeDirection cam h k
                        pushBack (stRingBuffer st) (d!0) (d!1) (d!2)
                        put (etype,(h,k),d)
                _ -> do
                    mev <- get ev #motion
                    h <- get mev #x
                    k <- get mev #y
                    put (h,k)
                    updateCameraRotation w st h k
                    return ()

        EventTypeButtonPress -> do
            case inputSource of
                Just InputSourcePen -> do
                    writeIORef (stPenDown st) True
                    bev <- get ev #button
                    h <- get bev #x
                    k <- get bev #y
                    cam <- readIORef (stCamera st)
                    let d = computeDirection cam h k
                    pushBack (stRingBuffer st) (d!0) (d!1) (d!2)
                    put (etype,(h,k),d)
                _ -> do
                    bev <- get ev #button
                    h <- get bev #x
                    k <- get bev #y
                    cam <- readIORef (stCamera st)
                    let d = computeDirection cam h k
                    writeIORef (stDragFrom st) $ Just (d,cam)
                    put (etype,(h,k),d)
                    return ()

        EventTypeButtonRelease -> do
            case inputSource of
                Just InputSourcePen -> do
                    writeIORef (stPenDown st) False
                    bev <- get ev #button
                    h <- get bev #x
                    k <- get bev #y
                    cam <- readIORef (stCamera st)
                    let d = computeDirection cam h k
                    pushBack (stRingBuffer st) (d!0) (d!1) (d!2)
                _ -> do
                    bev <- get ev #button
                    h <- get bev #x
                    k <- get bev #y
                    updateCameraRotation w st h k
                    sanitizeCamera st
                    writeIORef (stDragFrom st) Nothing

        EventTypeScroll -> do
            sev <- get ev #scroll
            d <- get sev #direction
            let δ = case d of
                    ScrollDirectionDown -> - pi/180
                    ScrollDirectionUp   -> pi/180
                    _                   -> 0
            when (δ /= 0) $ do
                modifyIORef (stCamera st) $ \cam -> cam
                    { camHeightAngle   = δ + camHeightAngle cam
                    , camWorldToScreen = Nothing
                    , camScreenToWorld = Nothing
                    }
                mwin <- widgetGetWindow w
                forM_ mwin $ \win ->
                    windowInvalidateRect win Nothing False
            put d
            return ()

        EventTypeKeyPress -> do
            kev <- get ev #key
            val <- get kev #keyval <&> \k -> if k > 0x5A then k - 0x20 else k
            case val of
                KEY_N -> do
                    modifyIORef' (stSkybox st) $ \n -> (n + 1) `mod` (skyboxCount $ stSkyboxes st)
                    idx <- readIORef (stSkybox st)
                    when (skyboxCount (stSkyboxes st) > 1) $ do
                        Right ts <- skyboxLoad (stSkyboxes st) idx
                        disposeTextureCube =<< readIORef (stSkyTexture st)
                        skybox_id <- uploadCubeMapToGPU ts
                        LC.updateUniforms (stStorage realized) $ do
                            "CubeMap" @= return skybox_id
                        writeIORef (stSkyTexture st) skybox_id
                        put (skyboxNames (stSkyboxes st) !! idx)
                        return ()
                KEY_F -> do
                    put 'F'
                    stFullscreen st
                _ -> return ()

        e -> return ()

    return False