{-# LANGUAGE RecordWildCards           #-}
{-# LANGUAGE TemplateHaskell           #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE ImpredicativeTypes        #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ViewPatterns              #-}
{-# LANGUAGE DataKinds                 #-}

module Thermoprint.Server
       ( thermoprintServer
       , Config(..), QMConfig(..)
       , withPrinters
       , module Data.Default.Class
       , module Servant.Utils.Enter
       , module Thermoprint.Server.Printer
       , module Thermoprint.Server.Queue
       , module Thermoprint.Server.Queue.Utils
       ) where

import Data.Default.Class
import qualified Config.Dyre as Dyre

import Data.Map (Map)
import qualified Data.Map as Map

import Data.Set (Set)
import qualified Data.Set as Set
       
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq

import Data.Time (UTCTime)

import Data.Maybe (maybe)
import Data.Foldable (mapM_, forM_, foldlM)
import Data.Function hiding (id, (.))
import Data.Bifunctor
import Data.Proxy

import Control.Monad.Trans.Resource
import Control.Monad.Trans.Control
import Control.Monad.Trans.Identity
import Control.Monad.Logger
import Control.Monad.Reader
import Control.Monad.IO.Class
import Control.Monad.Morph
import Control.Category
import Control.Monad.Catch (MonadMask(mask), finally)
import Prelude hiding (id, (.))

import qualified Control.Monad as M

import Control.Concurrent
import Control.Concurrent.STM
       
import Data.Text (Text)
import qualified Data.Text as T (pack)

import qualified Network.Wai.Handler.Warp as Warp
import Network.Wai (Application)

import Servant.Server (serve)
import Servant.Utils.Enter (enter, (:~>)(..))
import Servant.API
import Servant.Utils.Links
import Network.URI

import System.Environment (lookupEnv)

import Database.Persist.Sql (runMigrationSilent, ConnectionPool, runSqlPool)


import Thermoprint.API (thermoprintAPI, PrinterStatus, JobStatus)
import qualified Thermoprint.API as API (PrinterId, JobId)

import Thermoprint.Server.Fork

import Thermoprint.Server.Push

import Thermoprint.Server.Database
import Thermoprint.Server.Printer
import Thermoprint.Server.Queue
import Thermoprint.Server.Queue.Utils
import qualified Thermoprint.Server.API as API (thermoprintServer)
import Thermoprint.Server.API hiding (thermoprintServer)

import Debug.Trace

-- | Compile-time configuration for 'thermoprintServer'
data Config m = Config { dyreError     :: Maybe String -- ^ Set by 'Dyre' -- sent to log as an error
                       , warpSettings  :: Warp.Settings -- ^ Configure 'Warp's behaviour
                       , printers      :: Map API.PrinterId Printer
                       , queueManagers :: API.PrinterId -> QMConfig m
                       }

data QMConfig m = forall t. ( MonadTrans t
                            , MFunctor t
                            , Monad (t STM)
                            , MonadIO (t IO)
                            ) => QMConfig { manager  :: QueueManager t
                                          , collapse :: (t IO) :~> m
                                          }
  
instance MonadIO m => Default (Config m) where
  def = Config { dyreError     = Nothing
               , warpSettings  = Warp.defaultSettings
               , printers      = Map.empty
               , queueManagers = const def
               }

instance MonadIO m => Default (QMConfig m) where
  def = QMConfig idQM $ NT (liftIO . runIdentityT)

withPrinters :: MonadResource m => Config m -> [(m PrinterMethod, QMConfig m)] -> m (Config m)
-- ^ Add a list of printers to a 'Config'
withPrinters cfg = fmap updateCfg . foldlM mapInsert (Map.mapWithKey (\k a -> (a, queueManagers cfg k)) $ printers cfg)
  where
    mapInsert map (spec, qm) = Map.insert (nextKey map) <$> ((,) <$> printer spec <*> pure qm) <*> pure map
    updateCfg map = let
      printerMap = fmap fst map
      qmMap      = fmap snd map
      qmMap' id
        | (Just qm) <- (Map.lookup id qmMap) = qm
        | otherwise                          = queueManagers cfg id
      in cfg { printers = printerMap, queueManagers = qmMap' }
    nextKey map
      | Map.null map = 0
      | otherwise    = succ . fst $ Map.findMax map

thermoprintServer :: ( MonadLoggerIO m
                     , MonadReader ConnectionPool m
                     , MonadResourceBase m
                     , MonadMask m
                     ) => Bool -- ^ Invoke 'dyre' to look for and attempt to compile custom configurations (pass 'False' iff testing)
                     -> (m :~> IO) -- ^ 'dyre' controls the base of the monad-transformer-stack ('IO') but we let the user specify much of the rest of it (we handle 'ResourceT' ourselves, since we need it to fork properly). Therefore we require a specification of how to collapse the stack.
                     -> ResourceT m (Config (ResourceT m)) -> IO ()
-- ^ Run the server
thermoprintServer dyre io cfg = do
  cfgDir <- lookupEnv "THERMOPRINT_CONFIG"
  cacheDir <- lookupEnv "THERMOPRINT_CACHE"
  flip Dyre.wrapMain cfg $ Dyre.defaultParams
    { Dyre.projectName = "thermoprint-server"
    , Dyre.realMain    = realMain
    , Dyre.showError   = flip (\msg -> fmap (\cfg -> cfg { dyreError = Just msg }))
    , Dyre.configCheck = dyre
    , Dyre.configDir   = return <$> cfgDir
    , Dyre.cacheDir    = return <$> cacheDir
    }
  where
    realMain cfg = (io . NT runResourceT) $$ do
      tMgr <- threadManager resourceForkIO
      flip finally (cleanup tMgr) $ do
        Config{..} <- cfg
        maybe (return ()) ($(logErrorS) "Dyre" . T.pack) dyreError
        mapM_ ($(logWarnS) "DB") =<< runSqlPool (runMigrationSilent migrateAll) =<< ask
        forM_ printers $ fork tMgr . runPrinter
        gcChan <- liftIO newTChanIO
        fork tMgr $ jobGC gcChan
        let
          runQM' (queueManagers -> QMConfig qm nat) printer = nat $$ runQM gcChan qm printer
        mapM_ (fork tMgr . uncurry runQM') $ Map.toList printers
        nChan <- liftIO $ newBroadcastTChanIO
        let
          printerUrl :: API.PrinterId -> URI
          printerUrl = linkURI . safeLink thermoprintAPI (Proxy :: Proxy ("jobs" :> QueryParam "printer" API.PrinterId :> Get '[JSON] (Seq (API.JobId, UTCTime, JobStatus)))) . Just
        mapM_ (fork tMgr . uncurry (notifyOnChange nChan ((==) `on` fromZipper)) . bimap printerUrl queue) $ Map.toList printers
        liftIO . Warp.runSettings warpSettings . withPush nChan . serve thermoprintAPI . flip enter API.thermoprintServer =<< handlerNat printers nChan