{-# OPTIONS_GHC -O0 -fasm -fprof-auto #-}

module Foundation.Yesod.Middleware
  ( yesodMiddleware
  , updateFavourites
  ) where

import Import.NoFoundation hiding (yesodMiddleware)

import Foundation.Type
import Foundation.Routes
import Foundation.Authorization
import Foundation.I18n

import Utils.Metrics
import Utils.Workflow

import Handler.Utils.Workflow.CanonicalRoute

import qualified Network.Wai as W
import qualified Data.Aeson as JSON
import qualified Data.CaseInsensitive as CI
  
import Yesod.Core.Types (GHState(..), HandlerData(..), RunHandlerEnv(rheSite, rheChild))

import qualified Data.Map as Map


yesodMiddleware :: ( BearerAuthSite UniWorX
                   , BackendCompatible SqlReadBackend (YesodPersistBackend UniWorX)
                   , BackendCompatible SqlBackend (YesodPersistBackend UniWorX)
                   )
                => HandlerFor UniWorX res -> HandlerFor UniWorX res
yesodMiddleware = cacheControlMiddleware . storeBearerMiddleware . csrfMiddleware . dryRunMiddleware . observeYesodCacheSizeMiddleware . languagesMiddleware appLanguages . headerMessagesMiddleware . securityMiddleware . normalizeRouteMiddleware . updateFavouritesMiddleware . setActiveAuthTagsMiddleware . normalizeApprootMiddleware
  where
    dryRunMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    dryRunMiddleware handler = do
      dryRun <- isDryRun
      if | dryRun -> do
             hData <- ask
             prevState <- readIORef (handlerState hData)
             let
               restoreSession =
                 modifyIORef (handlerState hData) $
                   \hst -> hst { ghsSession = ghsSession prevState
                               , ghsCache = ghsCache prevState
                               , ghsCacheBy = ghsCacheBy prevState
                               }
               site' = (rheSite $ handlerEnv hData) { appMemcached = Nothing }
               handler' = local (\hd -> hd { handlerEnv = (handlerEnv hd) { rheSite = site', rheChild = site' } }) handler

             addCustomHeader HeaderDryRun $ toPathPiece True

             handler' `finally` restoreSession
         | otherwise -> handler
    updateFavouritesMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    updateFavouritesMiddleware handler = (*> handler) . runMaybeT $ do
      route <- MaybeT getCurrentRoute
      case route of -- update Course Favourites here
        CourseR tid ssh csh _  -> do
          void . lift . runDB . runMaybeT $ do
            guardM . lift . hasReadAccessTo $ CourseR tid ssh csh CShowR
            lift . updateFavourites $ Just (tid, ssh, csh)
        _other -> return ()
    normalizeRouteMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    normalizeRouteMiddleware handler = (*> handler) . runMaybeT $ do
      route <- MaybeT getCurrentRoute
      (route', getAny -> changed) <- lift . runDB . runWriterT $ foldM (&) route routeNormalizers
      when changed $ do
        $logDebugS "normalizeRouteMiddleware" [st|Redirecting to #{tshow route'}|]
        redirectWith movedPermanently301 route'
    headerMessagesMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    headerMessagesMiddleware handler = (handler `finally`) . runMaybeT $ do
      isModal <- hasCustomHeader HeaderIsModal
      dbTableShortcircuit <- hasCustomHeader HeaderDBTableShortcircuit
      massInputShortcircuit <- hasCustomHeader HeaderMassInputShortcircuit
      $logDebugS "headerMessagesMiddleware" $ tshow (isModal, dbTableShortcircuit, massInputShortcircuit)
      guard $ or
        [ isModal
        , dbTableShortcircuit
        , massInputShortcircuit
        ]

      lift . bracketOnError getMessages (mapM_ addMessage') $
        addCustomHeader HeaderAlerts . decodeUtf8 . urlEncode True . toStrict . JSON.encode
    observeYesodCacheSizeMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    observeYesodCacheSizeMiddleware handler = handler `finally` observeYesodCacheSize
    csrfMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    csrfMiddleware handler = do
      hasBearer <- is _Just <$> lookupBearerAuth
      reqHost <- W.requestHeaderHost <$> waiRequest
      userGeneratedHost <- getsYesod $ \app ->
        guardOnM (views _appRoot ($ ApprootDefault) app /= views _appRoot ($ ApprootUserGenerated) app) $ approotScopeHost ApprootUserGenerated app
      mCurrentRoute <- getCurrentRoute
      let isError = case mCurrentRoute of
            Just ErrorR -> True
            _other -> False

      if | hasBearer 
           -> local (\HandlerData{..} -> HandlerData{ handlerRequest = handlerRequest { reqToken = Nothing }, .. }) handler
         | fromMaybe False ((==) <$> reqHost <*> userGeneratedHost) || isError -> do
             whenIsJust mCurrentRoute $ \currentRoute -> do
               isWrite <- isWriteRequest currentRoute
               when isWrite $
                 permissionDeniedI MsgUnauthorizedCsrfDisabled
             handler
         | otherwise
           -> csrfSetCookieMiddleware' . defaultCsrfCheckMiddleware $ handler
      where
        csrfSetCookieMiddleware' handler' = do
          mcsrf <- reqToken <$> getRequest
          whenIsJust mcsrf $ setRegisteredCookie CookieXSRFToken
          handler'
    storeBearerMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    storeBearerMiddleware handler = do
      askBearer >>= \case
        Just (Jwt bs) -> setSessionBS (toPathPiece SessionBearer) bs
        Nothing       -> return ()

      handler
    setActiveAuthTagsMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    setActiveAuthTagsMiddleware handler = do
      mtagActive <- lookupSessionJson SessionActiveAuthTags :: HandlerFor UniWorX (Maybe AuthTagActive)
      when (is _Nothing mtagActive) $ do
        mAuthTagActive <- lookupRegisteredCookieJson CookieActiveAuthTags
        for_ mAuthTagActive $ setSessionJson SessionActiveAuthTags . review _ReducedActiveAuthTags

      handler
    securityMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    securityMiddleware handler = do
      addHeader "X-XSS-Protection" "1; mode=block"
      addHeader "X-Frame-Options" "sameorigin"
      addHeader "X-Content-Type-Options" "nosniff"
      authorizationCheck
      handler
    cacheControlMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    cacheControlMiddleware = (addHeader "Vary" "Accept, Accept-Language" *>)
    normalizeApprootMiddleware :: HandlerFor UniWorX a -> HandlerFor UniWorX a
    normalizeApprootMiddleware handler = maybeT handler $ do
      route <- MaybeT getCurrentRoute

      case route of
        MetricsR -> mzero
        HealthR -> mzero
        InstanceR -> mzero
        AdminCrontabR -> mzero
        _other -> return ()
  
      reqHost <- MaybeT $ W.requestHeaderHost <$> waiRequest
      let rApproot = authoritiveApproot route 
      app <- getYesod
      approotHost <- hoistMaybe $ approotScopeHost rApproot app
      let doRedirect = do
            url <- approotRender rApproot route
            $logDebugS "normalizeApprootMiddleware" url
            redirect url
      if | approotHost /= reqHost
         , rApproot /= ApprootUserGenerated
           -> doRedirect
         | approotHost /= reqHost -> do
             resp <- try $ lift handler
             $logDebugS "normalizeApprootMiddleware" $ tshow (is _Right resp, preview _Left resp)
             case resp of
               Right _ -> doRedirect
               Left sc | is _HCRedirect sc -> throwM sc
               Left _ -> doRedirect
         | otherwise -> lift handler
    


updateFavourites :: forall m backend.
                    ( MonadHandler m, HandlerSite m ~ UniWorX
                    , BackendCompatible SqlBackend backend
                    , YesodAuth UniWorX
                    , UserId ~ AuthId UniWorX
                    )
                 => Maybe (TermId, SchoolId, CourseShorthand) -- ^ Insert course into favourites, as appropriate
                 -> ReaderT backend m ()
updateFavourites cData = void . withReaderT projectBackend . runMaybeT $ do
  $logDebugS "updateFavourites" "Updating favourites"

  now <- liftIO getCurrentTime
  uid <- MaybeT $ liftHandler maybeAuthId
  mcid  <- (for cData $ \(tid, ssh, csh) -> MaybeT . getKeyBy $ TermSchoolCourseShort tid ssh csh) :: MaybeT (SqlPersistT m) (Maybe CourseId)
  User{userMaxFavourites} <- MaybeT $ get uid

  -- update Favourites
  -- no need to store them with userMaxFavourites==0, since they will be removed in the pruning step anyway!
  when (userMaxFavourites > 0) $ for_ mcid $ \cid ->
    void . lift $ upsertBy
      (UniqueCourseFavourite uid cid)
      (CourseFavourite uid cid FavouriteVisited now)
      [CourseFavouriteLastVisit =. now]
  -- prune Favourites to user-defined size
  oldFavs <- lift $ selectList [CourseFavouriteUser ==. uid] []
  let deleteFavs = oldFavs
        & sortOn ((courseFavouriteReason &&& Down . courseFavouriteLastVisit) . entityVal)
        & drop userMaxFavourites
        & filter ((<= FavouriteVisited) . courseFavouriteReason . entityVal)
        & map entityKey
  unless (null deleteFavs) $
    lift $ deleteWhere [CourseFavouriteId <-. deleteFavs]

  
routeNormalizers :: forall m backend.
                    ( BackendCompatible SqlReadBackend backend
                    , MonadHandler m, HandlerSite m ~ UniWorX
                    , BearerAuthSite UniWorX
                    ) => [Route UniWorX -> WriterT Any (ReaderT backend m) (Route UniWorX)]
routeNormalizers = map (hoist (hoist liftHandler . withReaderT projectBackend) .)
  [ normalizeRender
  , ncSchool
  , ncAllocation
  , ncCourse
  , ncSheet
  , ncMaterial
  , ncTutorial
  , ncExam
  , ncExternalExam
  , ncAdminWorkflowDefinition
  , ncWorkflowInstance
  , ncWorkflowPayloadLabel
  , verifySubmission
  , verifyCourseApplication
  , verifyCourseNews
  , verifyWorkflowWorkflow
  , verifyMaterialVideo
  , verifyAllocationMatchingLog
  ]
  where
    normalizeRender :: Route UniWorX -> WriterT Any (ReaderT SqlReadBackend (HandlerFor UniWorX)) (Route UniWorX)
    normalizeRender route = route <$ do
      YesodRequest{..} <- liftHandler getRequest
      let original = (W.pathInfo reqWaiRequest, reqGetParams)
          rendered = renderRoute route
      if
        | (isSuffixOf `on` fst) original rendered -> do -- FIXME: this breaks when subsite prefixes are dynamic
            $logDebugS "normalizeRender" [st|#{tshow rendered} matches #{tshow original}|]
        | otherwise -> do
            $logDebugS "normalizeRender" [st|Redirecting because #{tshow rendered} does not match #{tshow original}|]
            tell $ Any True

    maybeOrig :: (Route UniWorX -> MaybeT (WriterT Any (ReaderT SqlReadBackend (HandlerFor UniWorX))) (Route UniWorX))
              -> Route UniWorX -> WriterT Any (ReaderT SqlReadBackend (HandlerFor UniWorX)) (Route UniWorX)
    maybeOrig f route = maybeT (return route) $ f route

    caseChanged :: (Eq a, Show a) => CI a -> CI a -> MaybeT (WriterT Any (ReaderT SqlReadBackend (HandlerFor UniWorX))) ()
    caseChanged a b
      | ((/=) `on` CI.original) a b = do
          $logDebugS "routeNormalizers" [st|#{tshow a} /= #{tshow b}|]
          tell $ Any True
      | otherwise                   = return ()

    ncSchool = maybeOrig . typesUsing @RouteChildren @SchoolId $ \ssh -> $cachedHereBinary ssh $ do
      let schoolShort :: SchoolShorthand
          schoolShort = unSchoolKey ssh
      Entity ssh' _  <- MaybeT . lift . getBy $ UniqueSchoolShorthand schoolShort
      (caseChanged `on` unSchoolKey) ssh ssh'
      return ssh'
    ncAllocation = maybeOrig $ \route -> do
      AllocationR tid ssh ash _ <- return route
      Entity _ Allocation{..} <- MaybeT . $cachedHereBinary (tid, ssh, ash) . lift . getBy $ TermSchoolAllocationShort tid ssh ash
      caseChanged ash allocationShorthand
      return $ route & typesUsing @RouteChildren @AllocationShorthand . filtered (== ash) .~ allocationShorthand
    ncCourse = maybeOrig $ \route -> do
      CourseR tid ssh csh _ <- return route
      Entity _ Course{..} <- MaybeT . $cachedHereBinary (tid, ssh, csh) . lift . getBy $ TermSchoolCourseShort tid ssh csh
      caseChanged csh courseShorthand
      return $ route & typesUsing @RouteChildren @CourseShorthand . filtered (== csh) .~ courseShorthand
    ncSheet = maybeOrig $ \route -> do
      CSheetR tid ssh csh shn _ <- return route
      cid <- MaybeT . $cachedHereBinary (tid, ssh, csh) . lift . getKeyBy $ TermSchoolCourseShort tid ssh csh
      Entity _ Sheet{..} <- MaybeT . $cachedHereBinary (cid, shn) . lift . getBy $ CourseSheet cid shn
      caseChanged shn sheetName
      return $ route & typesUsing @RouteChildren @SheetName . filtered (== shn) .~ sheetName
    ncMaterial = maybeOrig $ \route -> do
      CMaterialR tid ssh csh mnm _ <- return route
      cid <- MaybeT  . $cachedHereBinary (tid, ssh, csh) . lift . getKeyBy $ TermSchoolCourseShort tid ssh csh
      Entity _ Material{..} <- MaybeT . $cachedHereBinary (cid, mnm) . lift . getBy $ UniqueMaterial cid mnm
      caseChanged mnm materialName
      return $ route & typesUsing @RouteChildren @MaterialName . filtered (== mnm) .~ materialName
    ncTutorial = maybeOrig $ \route -> do
      CTutorialR tid ssh csh tutn _ <- return route
      cid <- MaybeT  . $cachedHereBinary (tid, ssh, csh) . lift . getKeyBy $ TermSchoolCourseShort tid ssh csh
      Entity _ Tutorial{..} <- MaybeT . $cachedHereBinary (cid, tutn) . lift . getBy $ UniqueTutorial cid tutn
      caseChanged tutn tutorialName
      return $ route & typesUsing @RouteChildren @TutorialName . filtered (== tutn) .~ tutorialName
    ncExam = maybeOrig $ \route -> do
      CExamR tid ssh csh examn _ <- return route
      cid <- MaybeT  . $cachedHereBinary (tid, ssh, csh) . lift . getKeyBy $ TermSchoolCourseShort tid ssh csh
      Entity _ Exam{..} <- MaybeT . $cachedHereBinary (cid, examn) . lift . getBy $ UniqueExam cid examn
      caseChanged examn examName
      return $ route & typesUsing @RouteChildren @ExamName . filtered (== examn) .~ examName
    ncExternalExam = maybeOrig $ \route -> do
      EExamR tid ssh coursen examn _ <- return route
      Entity _ ExternalExam{..} <- MaybeT . $cachedHereBinary (tid, ssh, coursen, examn) . lift . getBy $ UniqueExternalExam tid ssh coursen examn
      caseChanged coursen externalExamCourseName
      caseChanged examn externalExamExamName
      return $ route
        & typesUsing @RouteChildren @CourseName . filtered (== coursen) .~ externalExamCourseName
        & typesUsing @RouteChildren @ExamName . filtered (== examn) .~ externalExamExamName
    ncAdminWorkflowDefinition = maybeOrig $ \route -> do
      AdminWorkflowDefinitionR wds wdn _ <- return route
      Entity _ WorkflowDefinition{..} <- MaybeT . $cachedHereBinary (wds, wdn) . lift . getBy $ UniqueWorkflowDefinition wdn wds
      caseChanged wdn workflowDefinitionName
      return $ route
        & typesUsing @RouteChildren @WorkflowDefinitionName . filtered (== wdn) .~ workflowDefinitionName
    ncWorkflowInstance = maybeOrig $ \route -> do
      (rScope, WorkflowInstanceR win _) <- hoistMaybe $ route ^? _WorkflowScopeRoute
      dbScope <- fmap (view _DBWorkflowScope) . hoist lift $ fromRouteWorkflowScope rScope
      Entity _ WorkflowInstance{..} <- lift . lift . getBy404 $ UniqueWorkflowInstance win dbScope
      caseChanged win workflowInstanceName
      return $ route
        & typesUsing @RouteChildren @WorkflowInstanceName . filtered (== win) .~ workflowInstanceName
    ncWorkflowPayloadLabel = maybeOrig $ \route -> do
      (_, WorkflowWorkflowR cID (WWFilesR wpl _)) <- hoistMaybe $ route ^? _WorkflowScopeRoute
      wwId <- decrypt cID
      WorkflowWorkflow{..} <- MaybeT . $cachedHereBinary wwId . lift $ get wwId
      wwGraph <- lift . lift $ getSharedDBWorkflowGraph workflowWorkflowGraph
      [wpl'] <- return . filter (== wpl) . sortOn (CI.original . unWorkflowPayloadLabel) . foldMap Map.keys $ wgnPayloadView <$> wgNodes wwGraph
      (caseChanged `on` unWorkflowPayloadLabel) wpl wpl'
      return $ route
        & typesUsing @RouteChildren @WorkflowPayloadLabel . filtered (== wpl) .~ wpl'
    verifySubmission = maybeOrig $ \route -> do
      CSubmissionR _tid _ssh _csh _shn cID sr <- return route
      sId <- decrypt cID
      Submission{submissionSheet} <- MaybeT . $cachedHereBinary cID . lift $ get sId
      Sheet{sheetCourse, sheetName} <- MaybeT . $cachedHereBinary submissionSheet . lift $ get submissionSheet
      Course{courseTerm, courseSchool, courseShorthand} <- MaybeT . $cachedHereBinary sheetCourse . lift $ get sheetCourse
      let newRoute = CSubmissionR courseTerm courseSchool courseShorthand sheetName cID sr
      tell . Any $ route /= newRoute
      return newRoute
    verifyCourseApplication = maybeOrig $ \route -> do
      CApplicationR _tid _ssh _csh cID sr <- return route
      aId <- decrypt cID
      CourseApplication{courseApplicationCourse} <- lift . lift $ get404 aId
      Course{courseTerm, courseSchool, courseShorthand} <- lift . lift $ get404 courseApplicationCourse
      let newRoute = CApplicationR courseTerm courseSchool courseShorthand cID sr
      tell . Any $ route /= newRoute
      return newRoute
    verifyCourseNews = maybeOrig $ \route -> do
      CNewsR _tid _ssh _csh cID sr <- return route
      aId <- decrypt cID
      CourseNews{courseNewsCourse} <- lift . lift $ get404 aId
      Course{courseTerm, courseSchool, courseShorthand} <- lift . lift $ get404 courseNewsCourse
      let newRoute = CNewsR courseTerm courseSchool courseShorthand cID sr
      tell . Any $ route /= newRoute
      return newRoute
    verifyWorkflowWorkflow = maybeOrig $ \route -> do
      (_, WorkflowWorkflowR cID wwR) <- hoistMaybe $ route ^? _WorkflowScopeRoute
      wwId <- decrypt cID
      WorkflowWorkflow{..} <- lift . lift $ get404 wwId 
      rScope <- hoist lift . toRouteWorkflowScope $ _DBWorkflowScope # workflowWorkflowScope
      let newRoute = _WorkflowScopeRoute # (rScope, WorkflowWorkflowR cID wwR)
      tell . Any $ route /= newRoute
      return newRoute
    verifyMaterialVideo = maybeOrig $ \route -> do
      CMaterialR _tid _ssh _csh _mnm (MVideoR cID) <- return route
      mfId <- decrypt cID
      MaterialFile{materialFileMaterial} <- lift . lift $ get404 mfId
      Material{materialName, materialCourse} <- lift . lift $ get404 materialFileMaterial
      Course{courseTerm, courseSchool, courseShorthand} <- lift . lift $ get404 materialCourse
      let newRoute = CMaterialR courseTerm courseSchool courseShorthand materialName $ MVideoR cID
      tell . Any $ route /= newRoute
      return newRoute
    verifyAllocationMatchingLog = maybeOrig $ \route -> do
      AllocationR _tid _ssh _ash (AMatchingR cID AMLogR) <- return route
      amId <- decrypt cID
      AllocationMatching{allocationMatchingAllocation} <- lift . lift $ get404 amId
      Allocation{allocationTerm, allocationSchool, allocationShorthand} <- lift . lift $ get404 allocationMatchingAllocation
      let newRoute = AllocationR allocationTerm allocationSchool allocationShorthand $ AMatchingR cID AMLogR
      tell . Any $ route /= newRoute
      return newRoute
