{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE UndecidableInstances #-}

module Engine.Vulkan.Pipeline.External
  ( Process
  , spawn
  , spawnReflect

  , loadConfig
  , loadConfigReflect

  , Observer

  , newObserverGraphics
  , observeGraphics

  , newObserverCompute
  , observeCompute

  , type (^)
  , ConfigureGraphics
  , ConfigureCompute
  , Observers
  , observeField

  , dumpPipelines
  ) where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Data.List (maximum)
import Data.Tagged (Tagged(..))
import Data.Type.Equality (type (~))
import RIO.ByteString qualified as ByteString
import RIO.Directory (createDirectoryIfMissing, getModificationTime, doesFileExist)
import RIO.FilePath ((</>), (<.>))
import RIO.Map qualified as Map
import RIO.Text qualified as Text
import RIO.Time (UTCTime, getCurrentTime)
import UnliftIO.Resource (MonadResource, ReleaseKey)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk

import Render.Code (Code(..))
import Engine.SpirV.Reflect qualified as Reflect
import Engine.Types (StageFrameRIO, StageRIO)
import Engine.Vulkan.Pipeline.Compute (Compute)
import Engine.Vulkan.Pipeline.Compute qualified as Compute
import Engine.Vulkan.Pipeline.Graphics qualified as Graphics
import Engine.Vulkan.Pipeline.Stages (StageInfo(..))
import Engine.Vulkan.Shader qualified as Shader
import Engine.Vulkan.Types (DsLayoutBindings, HasRenderPass)
import Engine.Worker qualified as Worker

type Process config = Worker.Timed () config

spawn
  :: ( Foldable stages
     , MonadReader env m
     , HasLogFunc env
     , MonadResource m
     , MonadUnliftIO m
     )
  => (stages (Maybe FilePath) -> m stuff)
  -> Text
  -> stages (Maybe FilePath)
  -> (stuff -> config)
  -> m (Process config)
spawn loader label stageFiles makeConfig =
  Worker.spawnTimed startActive dtF initF stepF ()
  where
    startActive = True

    dtF = Left 1e6

    initF () = do
      logDebug $ "Starting pipeline watch on " <> display label

      result <- loader stageFiles
      initialTime <- getCurrentTime
      pure (makeConfig result, initialTime)

    stepF oldTime () = do
      checkTime oldTime (catMaybes $ toList stageFiles) >>= \case
        Nothing ->
          -- logDebug $ "Skipping pipeline update from " <> display label
          pure (Nothing, oldTime)
        Just newTime -> do
          logInfo $ "Updating pipeline from " <> display label
          try (loader stageFiles) >>= \case
            Left (SomeException err) -> do
              logError $ displayShow err
              pure (Nothing, newTime)
            Right result ->
              pure
                ( Just $ makeConfig result
                , newTime
                )

spawnReflect
  :: ( MonadResource m
     , MonadUnliftIO m
     , MonadReader env m
     , HasLogFunc env
     , StageInfo stages
     )
  => Text
  -> stages (Maybe FilePath)
  -> ((stages (Maybe ByteString), Reflect.Reflect stages) -> config)
  -> m (Process config)
spawnReflect = spawn loadConfigReflect

checkTime
  :: MonadIO io
  => UTCTime
  -> [FilePath]
  -> io (Maybe UTCTime)
checkTime oldTime = fmap collect . traverse getModificationTime
  where
    collect = \case
      [] ->
        Nothing
      (maximum -> maxTime) ->
        if maxTime <= oldTime then
          Nothing
        else
          Just maxTime

loadConfig
  :: ( Traversable stages
     , MonadIO io
     )
  => stages (Maybe FilePath)
  -> io (stages (Maybe ByteString))
loadConfig stageFiles =
  for stageFiles $
    traverse ByteString.readFile

loadConfigReflect
  :: ( StageInfo stages
     , MonadIO io
     , MonadReader env io
     , HasLogFunc env
     )
  => stages (Maybe FilePath)
  -> io (stages (Maybe ByteString), Reflect.Reflect stages)
loadConfigReflect stageFiles = do
  stageCode <- loadConfig stageFiles

  -- TODO: use stageCode?
  stageRefl <- for stageFiles $
    traverse Reflect.invoke

  reflDS <- Reflect.stagesBindMap stageRefl

  let reflIS = Reflect.stagesInterfaceMap stageRefl
  case Reflect.interfaceCompatible reflIS of
    Right ok ->
      logDebug $ displayShow ok
    Left (inputStage, outputStage, location, err) -> do
      let
        between = "Between " <> outputStage <> " and " <> inputStage
        locInfo = "(location=" <> show location <> ")"
      case err of
        Nothing ->
          throwString $
            unwords [ between, "missing output ", locInfo ]
        Just (sigRequested, sigProvided) ->
          throwString $ unlines
            [ unwords [ between, locInfo <> ":", "incompatible signatures" ]
            , "  requested: " <> show sigRequested
            , "  provided: " <> show sigProvided
            ]

  (inputStage, inputs) <-
    case Reflect.inputStageInterface reflIS of
      Nothing ->
        throwString "No active stage"
      Just found ->
        pure found

  let
    reflect = Reflect.Reflect
      { bindMap    = reflDS
      , interfaces = reflIS
      , inputStage = inputStage
      , inputs     = inputs
      }
  pure
    ( stageCode
    , reflect
    )

type Observer pipeline = Worker.ObserverIO (ReleaseKey, pipeline)

newObserverGraphics
  :: ( pipeline ~ Graphics.Pipeline dsl vertices instances
    , Worker.HasOutput worker
    , Shader.Specialization (Graphics.Specialization pipeline)
    , HasRenderPass renderpass
    , Worker.GetOutput worker ~ Graphics.Configure pipeline
    )
  => renderpass
  -> Vk.SampleCountFlagBits
  -> worker
  -> ResourceT (StageRIO rs) (Observer pipeline)
newObserverGraphics rp msaa process = do
  initialConfig <- Worker.getOutputData process

  initial <- Graphics.allocate
    Nothing
    msaa
    initialConfig
    rp

  Worker.newObserverIO initial

observeGraphics
  :: ( HasRenderPass renderpass
     , Worker.HasOutput output
     , Worker.GetOutput output ~ Graphics.Configure pipeline
     , pipeline ~ Graphics.Pipeline dsl vertices instances
     , spec ~ Graphics.Specialization pipeline
     , Shader.Specialization spec
     )
  => renderpass
  -> Vk.SampleCountFlagBits
  -> Tagged dsl [DsLayoutBindings]
  -> output
  -> Worker.ObserverIO (ReleaseKey, pipeline)
  -> StageFrameRIO rp p fr rs ()
observeGraphics rp msaa sceneBinds configP output =
  void $! Worker.observeIO configP output \(oldKey, _old) config -> do
    logDebug "Rebuilding pipeline"
    Resource.release oldKey
    mapRIO fst $ Graphics.allocate
      Nothing
      msaa
      ( config
          { Graphics.cDescLayouts = sceneBinds
          }
      )
      rp

newObserverCompute
  :: ( config ~ Compute.Configure pipeline ()
     , pipeline ~ Compute.Pipeline dsl Compute Compute
     )
  => Process config
  -> ResourceT (StageRIO rs) (Observer pipeline)
newObserverCompute process = do
  initialConfig <- Worker.getOutputData process

  initial <- Compute.allocate initialConfig

  Worker.newObserverIO initial

observeCompute
  :: ( Worker.HasOutput output
     , Worker.GetOutput output ~ config
     , Shader.Specialization spec
     , config ~ Compute.Configure pipeline spec
     , pipeline ~ Compute.Pipeline dsl Compute Compute
     )
  => Tagged dsl [DsLayoutBindings]
  -> output
  -> Worker.ObserverIO (ReleaseKey, pipeline)
  -> StageFrameRIO rp p fr rs ()
observeCompute binds configP output =
  void $! Worker.observeIO configP output \(oldKey, _old) config -> do
    logDebug "Rebuilding pipeline"
    Resource.release oldKey
    mapRIO fst $ Compute.allocate config
      { Compute.cDescLayouts = binds
      }

-- * HKD wrappers

data ConfigureGraphics p
data ConfigureCompute p
data Observers p

type family f ^ p where
  Identity ^ p = p
  ConfigureGraphics ^ p = Process (Graphics.Configure p)
  ConfigureCompute ^ p = Process (Compute.Configure p ())
  Observers ^ p = Observer p
  f ^ p = f p

observeField
  :: forall
      pf
      p
      renderpass
      dsl
      s vs is
      rps ps fr rs
  .   ( p ~ Graphics.Pipeline s vs is
      , Shader.Specialization (Graphics.Specialization p)
      , HasRenderPass renderpass
      )
  => renderpass
  -> Vk.SampleCountFlagBits
  -> Tagged dsl DsLayoutBindings
  -> pf ConfigureGraphics
  -> pf Observers
  -> (forall a . pf a -> a ^ p)
  -> StageFrameRIO rps ps fr rs ()
observeField rp msaa binds workers observers field =
  observeGraphics
    rp
    msaa
    (Tagged [unTagged binds])
    (field @ConfigureGraphics workers)
    (field @Observers observers)

dumpPipelines
  :: StageInfo t
  => MonadIO io
  => FilePath
  -> Map Text (t (Maybe Code))
  -> io ()
dumpPipelines prefix pipelines = do
  createDirectoryIfMissing True prefix
  for_ (Map.toList pipelines) \(pipeline, stageCode) -> do
    let stages = (,) <$> stageNames <*> stageCode
    for_ stages \(stage, mcode) ->
      for_ mcode \(Code code) -> do
        let
          file = prefix </> Text.unpack pipeline <.> stage
          bytes = encodeUtf8 code
        exists <- doesFileExist file
        if exists then do
          old <- ByteString.readFile file
          unless (bytes == old) $
            ByteString.writeFile file bytes
        else
          ByteString.writeFile file bytes