module Handler.Sheet.PersonalisedFiles
  ( sinkPersonalisedSheetFiles
  , getSPersonalFilesR, getCPersonalFilesR
  , PersonalisedSheetFilesKeyException(..)
  , sourcePersonalisedSheetFiles, resolvePersonalisedSheetFiles
  , PersonalisedSheetFilesDownloadAnonymous(..)
  , PersonalisedSheetFileUnresolved(..)
  , _PSFUnresolved, _PSFUnresolvedCollatable, _PSFUnresolvedDirectory
  ) where

import Import hiding (StateT(..))

import Handler.Utils
import Handler.Sheet.PersonalisedFiles.Meta
import Handler.Sheet.PersonalisedFiles.Types

import qualified Data.Conduit.Combinators as C
import qualified Data.Conduit.List as C (mapMaybeM)
import Data.Conduit.ResumableSink

import qualified Crypto.MAC.KMAC as Crypto
import qualified Data.ByteArray as BA
import qualified Data.Binary as Binary
import Crypto.Hash.Algorithms (SHAKE256)

import Data.ByteString.Lazy.Base32

import qualified Data.CaseInsensitive as CI

import Language.Haskell.TH (nameBase)

import qualified Data.CryptoID.ByteString as CryptoID
import qualified Data.CryptoID.Class.ImplicitNamespace as I

import qualified Database.Esqueleto as E

import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text as Text
import Data.List (inits, tails)
 
import Text.Unidecode (unidecode)
import Data.Char (isAlphaNum)

import qualified System.FilePath as FilePath (joinPath)
import System.FilePath.Glob

import Control.Monad.Trans.State.Strict (StateT, runStateT)
import qualified Control.Monad.State as State
import Control.Monad.Memo (MemoStateT, MonadMemo(..), for2)
import Utils.Memo

{-# ANN module ("HLint: ignore Use newtype instead of data" :: String) #-}

  
data PersonalisedSheetFileUnresolved a
  = PSFUnresolvedDirectory a
  | PSFUnresolvedCollatable Text a
  | PSFUnresolved a
  deriving (Eq, Ord, Read, Show, Functor, Foldable, Traversable, Generic, Typeable)

makePrisms ''PersonalisedSheetFileUnresolved


data PersonalisedSheetFilesRestriction
  = PSFRExamRegistered { psfrExam :: ExamId }
  deriving (Eq, Ord, Read, Show, Generic, Typeable)

makeLenses_ ''PersonalisedSheetFilesRestriction

data PersonalisedSheetFilesForm = PersonalisedSheetFilesForm
  { psffAnonymous :: PersonalisedSheetFilesDownloadAnonymous
  , psffRestrictions :: Set PersonalisedSheetFilesRestriction
  } deriving (Eq, Ord, Read, Show, Generic, Typeable)


personalisedSheetFileTypes :: [SheetFileType]
personalisedSheetFileTypes = filter (/= SheetMarking) universeF

  
resolvePersonalisedSheetFiles
  :: forall m a.
     ( MonadHandler m
     , HandlerSite m ~ UniWorX
     , MonadCatch m, MonadRandom m
     )
  => Lens' a FilePath
  -> (a -> Bool) -- ^ @isDirectory@
  -> CourseId
  -> SheetId
  -> ConduitT a (Either (PersonalisedSheetFileUnresolved a) (a, FileReferenceResidual PersonalisedSheetFile)) (SqlPersistT m) ()
resolvePersonalisedSheetFiles fpL isDir cid sid = do
  app <- getYesod
  C.mapM $ \fRef -> exceptT (return . Left . ($ fRef)) (return . Right . swap) . flip runStateT fRef $ do
    let
      genRefOptions :: ConduitT () (UserId, SheetFileType, FilePath) (StateT FilePath (ExceptT _ (SqlPersistT m))) ()
      genRefOptions = evalMemoStateC Map.empty $
             transPipe lift (yieldMany <=< State.gets $ resolvePersonalisedFilesDirectory app)
          .| C.mapMaybeM (runMaybeT . filterRefOption)
          where
            filterRefOption :: _ -> MaybeT (MemoStateT _ _ _ (StateT FilePath (ExceptT _ (SqlPersistT m)))) (UserId, SheetFileType, FilePath)
            filterRefOption (mbIdx, cID, sfType, fPath) = hoist (hoistStateCache $ lift . lift) $ do
                  let
                    getUid :: _ -> _ -> MemoStateT _ _ _ (SqlPersistT m) (Maybe UserId)
                    getUid mbIdx' cID' = runMaybeT $ do
                      cIDKey <- catchMPlus (Proxy @PersonalisedSheetFilesKeyException) . lift . lift $ getPersonalisedFilesKey cid (Just sid) mbIdx'
                      uid <- either (const mzero) return . (runReaderT ?? cIDKey) $ I.decrypt cID'
                      guardM . lift . lift $ exists [CourseParticipantUser ==. uid, CourseParticipantCourse ==. cid]
                      return uid
                        
                  fmap (, sfType, fPath) . hoistMaybeM . lift $ for2 memo getUid mbIdx cID
                  
    mbRef <- zoom fpL . runConduit $ genRefOptions .| C.head
    case mbRef of
      Just (uid, sfType, fPath) -> PersonalisedSheetFileResidual sid uid sfType <$ (fpL .= fPath)
      Nothing -> do
        isDirectory <- State.gets isDir
        fPath <- use fpL
        if | isDirectory
             -> lift $ throwE PSFUnresolvedDirectory
           | lstPtn : _ <- Map.keys $ Map.filter (`match'` fPath) personalisedSheetFilesCollatable
             -> lift . throwE $ PSFUnresolvedCollatable lstPtn
           | otherwise
             -> lift $ throwE PSFUnresolved
  where match' = matchWith $ matchDefault { matchDotsImplicitly = True }
    

sinkPersonalisedSheetFiles :: forall m.
                              ( MonadHandler m
                              , HandlerSite m ~ UniWorX
                              , MonadCatch m, MonadRandom m
                              )
                           => CourseId
                           -> SheetId
                           -> Bool -- ^ Keep existing?
                           -> ConduitT FileReference Void (SqlPersistT m) ()
sinkPersonalisedSheetFiles cid sid keep
  = resolvePersonalisedSheetFiles _fileReferenceTitle (hasn't $ _fileReferenceContent . _Just) cid sid
 .| evalRWSC () Map.empty fanoutReferences >>= msgUnreferenced
 where
   fanoutReferences = do
     C.mapM_ $ \case
       Left unresolved -> tell $ Set.singleton unresolved
       Right (fRef, residual) -> do
         let PersonalisedSheetFileResidual{..} = residual
             redResidual = (personalisedSheetFileResidualSheet, personalisedSheetFileResidualUser)
         mSinks <- State.gets $ Map.lookup redResidual
         let mkSinks
               | Just sinks' <- mSinks
               = Left sinks'
               | keep
               = Right $ \residual' -> newResumableSink $ sinkFileReferences residual'
               | otherwise
               = Right $ \residual' -> newResumableSink . void $ replaceFileReferences' mkFilter residual'
             sinks = case mkSinks of
               Left sinks' -> sinks'
               Right mkSinks' -> Map.fromList
                 [ (residual', mkSinks' residual')
                 | sfType <- personalisedSheetFileTypes
                 , let residual' = PersonalisedSheetFileResidual{ personalisedSheetFileResidualType = sfType, .. }
                 ]
             sink = Map.findWithDefault (error "No sink for SheetFileType") residual sinks
         sink' <- lift $ yield fRef ++$$ sink
         case sink' of
           Left _ -> error "sinkFileReferences/replaceFileReferences returned prematurely"
           Right nSink -> State.modify . Map.insert redResidual $ Map.insert residual nSink sinks
     openSinks <- State.get
     lift . lift . mapM_ closeResumableSink $ openSinks ^.. folded . folded
     let (nub -> sinkSheets, nub -> sinkUsers) = unzip $ Map.keys openSinks
     unless keep $
       lift . lift $ deleteWhere [ PersonalisedSheetFileSheet  <-. sinkSheets
                                 , PersonalisedSheetFileUser  /<-. sinkUsers
                                 ]

   msgUnreferenced ((), unreferenced) = unless (null collated && null uncollated) $
     addMessageModal msgStatus msgTrigger $ Right msgWidget
     where collated = Map.fromListWith (<>)
             [ (ptn, Sum 1)
             | PSFUnresolvedCollatable ptn _fRef <- Set.toList unreferenced
             ]
           collatedL = Map.toList collated
           uncollated = [ fileReferenceTitle | PSFUnresolved FileReference{..} <- Set.toList unreferenced ]

           Sum c = Sum (fromIntegral $ length uncollated) <> fold collated

           msgStatus | null uncollated = Info
                     | otherwise       = Warning

           msgTrigger = i18n $ MsgPersonalisedSheetFilesIgnored c
           msgWidget = $(widgetFile "messages/personalisedSheetFilesIgnored")

   mkFilter :: FileReferenceResidual PersonalisedSheetFile -> [Filter PersonalisedSheetFile]
   mkFilter PersonalisedSheetFileResidual{..} = [ PersonalisedSheetFileSheet ==. personalisedSheetFileResidualSheet
                                                , PersonalisedSheetFileUser  ==. personalisedSheetFileResidualUser
                                                , PersonalisedSheetFileType  ==. personalisedSheetFileResidualType
                                                ]

   sinkFileReferences :: FileReferenceResidual PersonalisedSheetFile -> ConduitT FileReference Void (SqlPersistT m) ()
   sinkFileReferences residual' = C.mapM_ $ \fRef -> void . put $ _FileReference # (fRef, residual')
           


sourcePersonalisedSheetFiles :: forall m.
                                ( MonadHandler m
                                , HandlerSite m ~ UniWorX
                                , MonadThrow m
                                , MonadRandom m
                                )
                             => CourseId
                             -> Maybe SheetId
                             -> Maybe (Set UserId)
                             -> PersonalisedSheetFilesDownloadAnonymous
                             -> Set PersonalisedSheetFilesRestriction
                             -> ConduitT () (Either PersonalisedSheetFile DBFile) (SqlPersistT m) ()
sourcePersonalisedSheetFiles cId mbsid mbuids anonMode restrs = do
  (mbIdx, cIDKey) <- lift . newPersonalisedFilesKey $ maybe (Left cId) Right mbsid
  let
    genSuffixes uid = case anonMode of
      PersonalisedSheetFilesDownloadGroups -> do
        subGroups <- E.select . E.from $ \(submissionGroup `E.InnerJoin` submissionGroupUser) -> do
          E.on $ submissionGroup E.^. SubmissionGroupId E.==. submissionGroupUser E.^. SubmissionGroupUserSubmissionGroup
          E.where_ $ submissionGroup E.^. SubmissionGroupCourse E.==. E.val cId
          E.where_ $ submissionGroupUser E.^. SubmissionGroupUserUser E.==. E.val uid
          return $ submissionGroup E.^. SubmissionGroupName
        return . nub . sort $ map (filter isAlphaNum . foldMap unidecode . unpack . CI.original . E.unValue) subGroups
      otherAnon
        | Just f <- userFeature otherAnon -> do
            features <- E.select . E.from $ \user -> do
              E.where_ $ user E.^. UserId E.==. E.val uid
              return $ f user
            return . sort $ mapMaybe (fmap (filter isAlphaNum . foldMap unidecode . unpack) . E.unValue) features
      _other -> return mempty
      where userFeature PersonalisedSheetFilesDownloadSurnames
              = Just $ E.just . (E.^. UserSurname)
            userFeature PersonalisedSheetFilesDownloadMatriculations
              = Just $ E.castString . (E.^. UserMatrikelnummer)
            userFeature _
              = Nothing
    
    sqlSource = E.selectSource . E.from $ \(courseParticipant `E.LeftOuterJoin` personalisedSheetFile) -> do
      E.on $ E.just (courseParticipant E.^. CourseParticipantUser) E.==. personalisedSheetFile E.?. PersonalisedSheetFileUser
       E.&&. E.val mbsid E.==. personalisedSheetFile E.?. PersonalisedSheetFileSheet
      E.where_ $ courseParticipant E.^. CourseParticipantCourse E.==. E.val cId
      case mbuids of
        Just uids -> E.where_ $ courseParticipant E.^. CourseParticipantUser `E.in_` E.valList (Set.toList uids)
        Nothing -> E.where_ $ courseParticipant E.^. CourseParticipantState E.==. E.val CourseParticipantActive
      forM_ restrs $ \case
        PSFRExamRegistered{..} -> E.where_ . E.exists . E.from $ \examRegistration ->
          E.where_ $ examRegistration E.^. ExamRegistrationExam E.==. E.val psfrExam
               E.&&. examRegistration E.^. ExamRegistrationUser E.==. courseParticipant E.^. CourseParticipantUser
      return (courseParticipant, personalisedSheetFile)

    toRefs = awaitForever $ \(Entity _ cPart@CourseParticipant{..}, mbPFile) -> do
      MsgRenderer mr <- getMsgRenderer
      suffix <- do
        sufCache <- uses _sufCache $ Map.lookup courseParticipantUser
        case sufCache of
          Just suf -> return suf
          Nothing -> do
            suf <- lift . lift $ genSuffixes courseParticipantUser
            _sufCache %= Map.insert courseParticipantUser suf
            return suf
      cID <- either throwM return . (runReaderT ?? cIDKey) $ I.encrypt courseParticipantUser
      let dirName = unpack . Text.intercalate "_" . map pack $ suffix `snoc` mkPersonalisedFilesDirectory mbIdx cID
      unlessM (uses _dirCache $ Set.member dirName) $ do
        yield $ Right File
          { fileTitle = dirName
          , fileContent = Nothing
          , fileModified = courseParticipantRegistration
          }
        forM_ [SheetExercise, SheetHint, SheetSolution] $ \sfType ->
          yield $ Right File
            { fileTitle = dirName <//> unpack (mr $ SheetArchiveFileTypeDirectory sfType)
            , fileContent = Nothing
            , fileModified = courseParticipantRegistration
            }
        yieldM . fmap Right $ do
          fileContent' <- lift $ formatPersonalisedSheetFilesMeta anonMode cPart cID
          let fileTitle = (dirName <//>) . ensureExtension "yaml" . unpack . mr $ MsgPersonalisedSheetFilesMetaFilename cID
              fileModified = courseParticipantRegistration
              fileContent = Just $ C.sourceLazy fileContent'
          return File{..}
        _dirCache %= Set.insert dirName
      whenIsJust mbPFile $ \(Entity _ pFile@PersonalisedSheetFile{..}) -> do
        let dirName' = dirName <//> unpack (mr $ SheetArchiveFileTypeDirectory personalisedSheetFileType)
        yield . Left $ over (_FileReference . _1 . _fileReferenceTitle) (dirName' <//>) pFile
      where
        _sufCache :: Lens' _ _
        _sufCache = _1
        _dirCache :: Lens' _ _
        _dirCache = _2
            

  sqlSource .| evalStateC (Map.empty, Set.empty) toRefs


data PersonalisedSheetFilesKeyException
  = PersonalisedSheetFilesKeyCouldNotDecodeRandom
  | FallbackPersonalisedSheetFilesKeysExhausted
  | PersonalisedSheetFilesKeyInsufficientContext
  | PersonalisedSheetFilesKeyNotFound
  deriving (Eq, Ord, Read, Show, Generic, Typeable)
  deriving anyclass (Exception)
  
newPersonalisedFilesKey :: forall m.
                           ( MonadHandler m
                           , HandlerSite m ~ UniWorX
                           , MonadThrow m, MonadRandom m
                           )
                        => Either CourseId SheetId -> SqlPersistT m (Maybe Word24, CryptoIDKey)
newPersonalisedFilesKey (Right shId) = cryptoIDKey $ \cIDKey -> fmap (Nothing,) $
  either (const $ throwM PersonalisedSheetFilesKeyCouldNotDecodeRandom) (views _3 return) . Binary.decodeOrFail . fromStrict . BA.convert $
  Crypto.kmac @(SHAKE256 CryptoIDCipherKeySize) (encodeUtf8 . (pack :: String -> Text) $ nameBase 'newPersonalisedFilesKey) (toStrict $ Binary.encode shId) cIDKey
newPersonalisedFilesKey (Left cId) = do
  now <- liftIO getCurrentTime
  secret <- CryptoID.genKey
  let secret' = toStrict $ Binary.encode secret
  firstN <- getRandom

  let loop :: Word24 -> SqlPersistT m (Maybe Word24, CryptoIDKey)
      loop n = do
        didInsert <- is _Just <$> insertUnique (FallbackPersonalisedSheetFilesKey cId n secret' now)
        if | didInsert
             -> return (Just n, secret)
           |    (firstN == minBound && n == maxBound)
             || n == pred firstN
             -> throwM FallbackPersonalisedSheetFilesKeysExhausted
           | n == maxBound
             -> loop minBound
           | otherwise
             -> loop $ succ n
   in loop firstN

getPersonalisedFilesKey :: forall m.
                           ( MonadHandler m
                           , HandlerSite m ~ UniWorX
                           , MonadThrow m, MonadRandom m
                           )
                        => CourseId -> Maybe SheetId -> Maybe Word24 -> SqlPersistT m CryptoIDKey
getPersonalisedFilesKey _    Nothing     Nothing   = throwM PersonalisedSheetFilesKeyInsufficientContext
getPersonalisedFilesKey _   (Just shId)  Nothing   = view _2 <$> newPersonalisedFilesKey (Right shId)
getPersonalisedFilesKey cId  _          (Just idx) = maybeT (throwM PersonalisedSheetFilesKeyNotFound) $ do
  Entity _ FallbackPersonalisedSheetFilesKey{..} <- MaybeT . getBy $ UniqueFallbackPersonalisedSheetFilesKey cId idx
  either (const $ throwM PersonalisedSheetFilesKeyCouldNotDecodeRandom) (views _3 return) . Binary.decodeOrFail . fromStrict $ BA.convert fallbackPersonalisedSheetFilesKeySecret

mkPersonalisedFilesDirectory :: Maybe Word24 -> CryptoFileNameUser -> FilePath
mkPersonalisedFilesDirectory  Nothing cID = unpack $ toPathPiece cID
mkPersonalisedFilesDirectory (Just idx) cID = unpack $ toPathPiece cID <> "-" <> CI.foldCase (toStrict . encodeBase32Unpadded $ Binary.encode idx)

resolvePersonalisedFilesDirectory :: forall master.
                                     RenderMessage master SheetArchiveFileTypeDirectory
                                  => master
                                  -> FilePath
                                  -> [(Maybe Word24, CryptoFileNameUser, SheetFileType, FilePath)]
resolvePersonalisedFilesDirectory foundation (splitPath -> fPath) = do
  (fPath', remFPath) <- inits fPath `zip` tails fPath
  guard . not $ null remFPath
  (SheetArchiveFileTypeDirectory sfType, fPath'') <- foldMap (\(seg, rest) -> (, rest) <$> unRenderMessageLenient foundation (pack seg)) $ foci fPath'
  guard $ sfType `elem` personalisedSheetFileTypes
  let cryptSegments = foldMap (filter (not . Text.null) . Text.split (flip Set.notMember cryptChars . CI.mk) . Text.pack) fPath''
  (mIdx, cryptSegments') <- foldMap (\(inp, rest) -> (, rest) . Just <$> hoistMaybe (decodeIdx inp)) (foci cryptSegments) <|> pure (Nothing, cryptSegments)
  cID <- foldMap (hoistMaybe . fromPathPiece) cryptSegments'
  return (mIdx, cID, sfType, FilePath.joinPath remFPath)
  where
    foci :: forall a. [a] -> [(a, [a])]
    foci [] = []
    foci (x:xs) = (x, xs) : map (over _2 (x:)) (foci xs)

    cryptoIdChars, base32Chars, cryptChars :: Set (CI Char)
    cryptChars = base32Chars <> cryptoIdChars
    cryptoIdChars = mappend base32Chars . Set.fromList $ map CI.mk "uwb"
    base32Chars = Set.fromList $ map CI.mk "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"

    decodeIdx :: Text -> Maybe Word24
    decodeIdx inp
      | Right inp' <- decodeBase32Unpadded . fromStrict $ encodeUtf8 inp
      , Right (remInp, _, idx) <- Binary.decodeOrFail inp'
      , null remInp
      = Just idx
      | otherwise = Nothing
  


getPersonalFilesR :: CourseId -> Maybe SheetId -> Handler TypedContent
getPersonalFilesR cId mbsid = do
  (Course{..}, mbSheet) <- runDB $ (,)
    <$> get404 cId
    <*> traverse get404 mbsid
  
  cRoute <- getCurrentRoute

  let toRestrictions = maybe Set.empty $ Set.singleton . PSFRExamRegistered
  ((psfRes, psfWdgt), psfEnctype) <- runFormGet . renderAForm FormStandard $ PersonalisedSheetFilesForm
    <$> apopt (selectField optionsFinite) (fslI MsgPersonalisedSheetFilesDownloadAnonymousField & setTooltip MsgPersonalisedSheetFilesDownloadAnonymousFieldTip) (Just PersonalisedSheetFilesDownloadAnonymous)
    <*> fmap toRestrictions (aopt (examField (Just $ SomeMessage MsgPersonalisedSheetFilesDownloadRestrictByExamNone) cId) (fslI MsgPersonalisedSheetFilesDownloadRestrictByExam & setTooltip MsgPersonalisedSheetFilesDownloadRestrictByExamTip) (Just $ mbSheet ^? _Just . _sheetType . _examPart . from _SqlKey))
  
  formResult psfRes $ \PersonalisedSheetFilesForm{..} -> do
    archiveName <- fmap (flip addExtension (unpack extensionZip) . unpack) . ap getMessageRender . pure $ case mbSheet of
      Nothing -> MsgCoursePersonalisedSheetFilesArchiveName courseTerm courseSchool courseShorthand
      Just Sheet{..} -> MsgPersonalisedSheetFilesArchiveName courseTerm courseSchool courseShorthand sheetName
    sendResponse <=< serveZipArchive' archiveName $ sourcePersonalisedSheetFiles cId mbsid Nothing psffAnonymous psffRestrictions
  
  isModal <- hasCustomHeader HeaderIsModal

  fmap toTypedContent . siteLayoutMsg MsgMenuSheetPersonalisedFiles $ do
    setTitleI MsgMenuSheetPersonalisedFiles
    wrapForm psfWdgt def
      { formMethod = GET
      , formAction = SomeRoute <$> cRoute
      , formEncoding = psfEnctype
      , formAttrs = formAttrs def <> bool mempty [("uw-no-navigate-away-prompt", ""), ("target", "_blank")] isModal
      }
  
getSPersonalFilesR :: TermId -> SchoolId -> CourseShorthand -> SheetName -> Handler TypedContent
getSPersonalFilesR tid ssh csh shn = do
  Entity shId Sheet{..} <- runDB $ fetchSheet tid ssh csh shn
  getPersonalFilesR sheetCourse $ Just shId

getCPersonalFilesR :: TermId -> SchoolId -> CourseShorthand -> Handler TypedContent
getCPersonalFilesR tid ssh csh = do
  cId <- runDB . getKeyBy404 $ TermSchoolCourseShort tid ssh csh
  getPersonalFilesR cId Nothing
