{-# 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 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)
import LambdaCube.GL.Type (TextureData(..))
-- import Text.Show.Pretty (ppShow)

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

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

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" []

    -- 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
    Right ts <- skyboxLoad skyboxes 0
    skybox_id <- uploadCubeMapToGPU ts
    LC.updateUniforms storage $ do
       "CubeMap" @= return skybox_id
    skytex <- newIORef skybox_id
    mi <- LC.uploadMeshToGPU cubeMesh
    LC.addMeshToObjectArray storage "SkyCube" [] mi

    drag <- newIORef Nothing

    let st = State
            { stAnimator     = tm
            , stCamera       = cam
            , stFullscreen   = toggle
            , stSkyboxes     = skyboxes
            , stSkybox       = skybox
            , stSkyTexture   = skytex
            , stDragFrom     = drag
            }
    -- _ <- 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)

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
          defUniforms $ do
            "CubeMap"         @: FTextureCube
            "CameraPosition"  @: V3F
            "ViewProjection"  @: M44F
            "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
    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)

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
            mev <- get ev #motion
            h <- get mev #x
            k <- get mev #y
            {-
            cam <- readIORef (stCamera st)
            {-
            let o = fromList [ camWidth cam / 2, camHeight cam / 2 ]
                r = camHeight cam / (2 * sin (camHeight cam / 2))

                c = fromList [realToFrac h, realToFrac k]  - o :: Vector Float
                d = realToFrac $ norm_2 c
                τ = asin (d / r) -- angle from center
                axis = fromList [c!1, - (c!0)] :: Vector Float
            -}
            let d̂ = camDirection cam -- forward
                û = camUp cam        -- upward
                r̂ = d̂ `cross` û      -- rightward
                x_r = realToFrac h - (camWidth cam / 2)
                x_u = (camHeight cam / 2) - realToFrac k
                x_d = (camHeight cam / 2) / tan (camHeightAngle cam / 2)
                x = fromList [x_r,x_u,x_d]
            -}

            updateCameraRotation w st h k
            return ()

        EventTypeButtonPress -> 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
            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)
                    Right ts <- skyboxLoad (stSkyboxes st) idx
                    disposeTexture =<< 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