{-# LANGUAGE
  CPP
, OverloadedStrings
, RecordWildCards
, QuasiQuotes
, TemplateHaskell
, TypeFamilies
, TypeOperators
, MultiParamTypeClasses
, FunctionalDependencies
, FlexibleContexts
, FlexibleInstances
, AllowAmbiguousTypes
, UndecidableInstances
, GeneralizedNewtypeDeriving
, ScopedTypeVariables
, TypeFamilyDependencies #-}

module Yesod.Auth.HmacKeccak where

import Yesod.Auth.Import
import qualified Data.Text as T
import qualified Data.Char as C
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Char8 as BC
import Data.Maybe (fromJust)

import qualified Database.Persist as P

import System.Random
import System.IO.Unsafe (unsafePerformIO)

import Numeric (readHex, showHex)

import Yesod.Auth
import Yesod.Auth.Message
import Yesod.Persist hiding (get, replace, insertkey, Entity, entityVal)
import Yesod.Static

import Text.Julius (jsFile)

import Paths_yesod_auth_hmac_keccak as Paths
import Yesod.Auth.JsPath

type Username = Text

-- js_auth_js :: String
-- js_auth_js =
--   unsafePerformIO $ readFile =<< getDataFileName "static/js/auth.js"

hmacPlugin
  :: YesodHmacKeccak db master
  => AuthPlugin master
hmacPlugin = AuthPlugin "authHmacKeccak" dispatch loginWidget
  where
    dispatch "POST" ["login"]       = postLoginR' >>= sendResponse
    dispatch "GET"  ["newaccount"]  = getNewAccountR' >>= sendResponse
    dispatch "POST" ["newaccount"]  = postNewAccountR' >>= sendResponse
    dispatch "GET"  ["resetpasswd"] = getReactivateR' >>= sendResponse
    dispatch "POST" ["resetpasswd"] = postReactivateR' >>= sendResponse
    dispatch "GET"  ["verify", k]   = getVerifyR' (encodeUtf8 k) >>= sendResponse
    dispatch "POST" ["verify", k]   = postVerifyR' (encodeUtf8 k) >>= sendResponse
    dispatch _ _ = notFound

newAccountR :: AuthRoute
newAccountR = PluginR "authHmacKeccak" ["newaccount"]

verifyR :: ByteString -> AuthRoute
verifyR k = PluginR "authHmacKeccak" ["verify", (decodeUtf8 k)]

resetPasswordR :: AuthRoute
resetPasswordR = PluginR "authHmacKeccak" ["resetpasswd"]

loginR :: AuthRoute
loginR = PluginR "authHmacKeccak" ["login"]

-- Login procedure.

loginWidget
  :: YesodHmacKeccak db master
  => (Route Auth -> Route master) -> WidgetT master IO ()
loginWidget tm = do
  render <- getUrlRenderParams
  toWidgetHead $ $(jsFile jsPath) render
  [whamlet|
<div .loginDiv>
  <form #loginform action=@{tm loginR} method="post">
    <div>
      <label for="username">_{MsgUser}:
      <input #username type="text" required>
    <div>
      <label for="password">_{MsgPassword}:
      <input #password type="password" required>
    <div>
      <button type="submit" #login>_{MsgLogin}
    <div #progress>
  <a href="@{tm newAccountR}">_{MsgRegister}
  <a href="@{tm resetPasswordR}">_{MsgForgotPassword}
  $maybe route <- rawLoginRoute
    <a href="@{route}">_{MsgNoJsLogin}
  |]

postLoginR'
  :: ( YesodHmacKeccak db master
     , YesodAuth master
     )
  => HandlerT Auth (HandlerT master IO) RepJson
postLoginR' = do
  mr <- lift getMessageRender
  mUserName <- lookupPostParam "username"
  mHexToken <- lookupPostParam "token"
  mHexResponse <- lookupPostParam "response"
  case (mUserName, mHexToken, mHexResponse) of
    (Just userName, Nothing, Nothing) -> do
      tempUser <- lift $ runHmacDB $ loadUser userName
      case tempUser of
        Just u ->
          if userUserActive u
          then do
            let salt = userUserSalt u
            token <- liftIO makeRandomToken
            lift $ runHmacDB $ insertLoginToken (encodeUtf8 token) userName
            returnJson ["salt" .= toHex salt, "token" .= toHex (encodeUtf8 token)]
          else do
            returnJsonError (mr MsgUserNotActive)
        Nothing ->
          returnJsonError (mr MsgNoSuchUser)
    (Nothing, Just hexToken, Just hexResponse) -> do
      response <- do
        let tempToken = fromHex' $ T.unpack hexToken
        savedToken <- lift $ runHmacDB $ loadLoginToken tempToken
        case savedToken of
          Just token -> do
            queriedUser <- lift $ runHmacDB $ loadUser (tokenTokenUsername token)
            let salted = userUserSalted $ fromJust queriedUser
                hexSalted = toHex salted
                expected =
                  hmacKeccak
                  (encodeUtf8 $ toHex $ tokenTokenToken token)
                  (encodeUtf8 hexSalted)
            if encodeUtf8 hexResponse == expected
            then do
              -- SUCCESS !!
              lift $ runHmacDB $ deleteToken token
              return $ Right $ fromJust queriedUser
            else
              return $ Left (mr MsgWrongPassword)
          Nothing ->
            return $ Left (mr MsgInvalidToken)
      case response of
        Left msg -> returnJsonError msg
        Right au -> do
          lift $ setCreds False $ Creds "authHmacKeccak" (userUserName au) []
          render <- lift getUrlRender
          m <- lift getYesod
          let u = render (loginDest m)
          returnJson ["welcome" .= u]
    _ ->
      returnJsonError (mr MsgProtocolError)

-- New account procedure

data NewAccountData = NewAccountData
  { naUsername :: Username
  , naEmail    :: Text
  } deriving Show

newAccountForm
  :: ( YesodHmacKeccak db master
     , MonadHandler m
     , HandlerSite m ~ master
     ) => AForm m NewAccountData
newAccountForm = NewAccountData
  <$> areq (checkM checkValidUsername textField) userSettings Nothing
  <*> areq emailField emailSettings Nothing
  where
    userSettings  = FieldSettings (SomeMessage MsgUsername) Nothing Nothing Nothing []
    emailSettings = FieldSettings (SomeMessage MsgEmail) Nothing Nothing Nothing []

newAccountWidget
  :: YesodHmacKeccak db master
  => (Route Auth -> Route master) -> WidgetT master IO ()
newAccountWidget tm = do
  render <- getUrlRenderParams
  toWidgetHead $ $(jsFile jsPath) render
  ((_, widget), enctype) <- liftHandlerT $ runFormPost $ renderDivs newAccountForm
  [whamlet|
<div .newaccount>
  <form method="post" enctype=#{enctype} action=@{tm newAccountR}>
    ^{widget}
    <input type=submit value=_{MsgRegister}>
  |]

getNewAccountR'
  :: YesodHmacKeccak db master
  => HandlerT Auth (HandlerT master IO) Html
getNewAccountR' = do
  tm <- getRouteToParent
  lift $ defaultLayout $ do
    setTitleI MsgRegisterLong
    newAccountWidget tm

postNewAccountR'
  :: YesodHmacKeccak db master
  => HandlerT Auth (HandlerT master IO) Html
postNewAccountR' = do
  tm <- getRouteToParent
  ((result, _), _) <- lift $ runFormPost $ renderDivs newAccountForm
  case result of
    FormMissing -> invalidArgs ["Form is missing"]
    FormFailure msg -> do
      setMessage $ toHtml $ T.concat msg
      redirect newAccountR
    FormSuccess d -> do
      lift $ setMessageI MsgActivationSent
      lift $ createNewAccount d tm
      redirect LoginR

createNewAccount
  :: YesodHmacKeccak db master
  => NewAccountData
  -> (Route Auth -> Route master)
  -> HandlerT master IO (UserAccount db)
createNewAccount nad@NewAccountData{..} tm = do
  muser <- runHmacDB $ loadUser naUsername
  case muser of
    Just _ -> do
      setMessageI $ MsgUsernameExists naUsername
      redirect $ tm newAccountR
    Nothing -> return ()
  token <- liftIO $ makeRandomToken
  salt <- liftIO $ makeRandomSalt
  enew <- runHmacDB $
    addNewUser naUsername naEmail (encodeUtf8 salt)
  _ <- runHmacDB $
    insertActivateToken (encodeUtf8 token) naUsername
  render <- getUrlRender
  sendVerifyEmail naUsername naEmail $ render $ tm $
    verifyR $ encodeUtf8 token
  new <- case enew of
    Left err -> do
      setMessage $ toHtml err
      redirect $ tm newAccountR
    Right x -> return x
  return new

-- verification procedure

data PWData = PWData
  { pw1 :: Text
  , pw2 :: Text
  }

passwordForm
  :: ( YesodHmacKeccak db master
     , MonadHandler m
     , HandlerSite m ~ master
     ) => AForm m PWData
passwordForm = PWData
  <$> areq textField pw1Settings Nothing
  <*> areq textField pw2Settings Nothing
  where
    pw1Settings = FieldSettings (SomeMessage MsgPassword1) Nothing Nothing Nothing []
    pw2Settings = FieldSettings (SomeMessage MsgPassword2) Nothing Nothing Nothing []

passwordWidget
  :: YesodHmacKeccak db master
  => (Route Auth -> Route master) -> ByteString -> ByteString -> WidgetT master IO ()
passwordWidget tm token hexSalt= do
  render <- getUrlRenderParams
  toWidgetHead $ $(jsFile jsPath) render
  [whamlet|
<div .password>
  <form #activateform method=post action=@{tm $ verifyR token}>
    <div .required>
      <label for="password1">_{MsgPassword1}:
      <input #password1 type="password" required>
    <div .required>
      <label for="password2">_{MsgPassword2}:
      <input #password2 type="password" required>
    <div #progress>
    <button #activate type="submit" data-token="#{BC.unpack token}" data-salt="#{BC.unpack hexSalt}">
      _{MsgActivate}
  |]

-- activateToken'
--   :: YesodHmacKeccak db master
--   => ByteString
--   -> HandlerT Auth (HandlerT master IO) (Maybe (TokenFoo db))
-- activateToken' = lift . runHmacDB . loadActivateToken

getVerifyR'
  :: YesodHmacKeccak db master
  => ByteString -> HandlerT Auth (HandlerT master IO) Html
getVerifyR' k = do
  mtoken <- lift $ runHmacDB $ loadActivateToken k
  case mtoken of
    Nothing -> do
      lift $ setMessageI MsgInvalidToken
      redirect LoginR
    Just token -> do
      muser <- lift $ runHmacDB $ loadUser $ tokenTokenUsername token
      case muser of
        Nothing -> do
          lift $ setMessageI MsgNoSuchUser
          redirect LoginR
        Just user -> do
          let hexSalt = toHex $ userUserSalt user
          tm <- getRouteToParent
          lift $ defaultLayout $ do
            setTitleI MsgSetPassword
            passwordWidget tm (tokenTokenToken token) (BC.pack $ T.unpack hexSalt)

postVerifyR'
  :: YesodHmacKeccak db master
  => ByteString -> HandlerT Auth (HandlerT master IO) RepJson
postVerifyR' k = do
  mtoken <- lift $ runHmacDB $ loadActivateToken k
  case mtoken of
    Nothing -> do
      lift $ setMessageI MsgInvalidToken
      redirect LoginR
    Just token -> do
      muser <- lift $ runHmacDB $ loadUser $ tokenTokenUsername token
      case muser of
        Nothing -> do
          lift $ setMessageI MsgNoSuchUser
          redirect LoginR
        Just user -> do
          msalted <- lookupPostParam "salted"
          case msalted of
            Nothing -> do
              lift $ setMessageI MsgProtocolError
              redirect LoginR
            Just salted' -> do
              let salted = fromHex' $ T.unpack salted'
              lift $ runHmacDB $ activateUser user salted
              lift $ runHmacDB $ deleteToken token
              lift $ setCreds False $ Creds "authHmacKeccak" (tokenTokenUsername token) []
              render <- lift getUrlRender
              m <- lift getYesod
              let u = render (loginDest m)
              returnJson ["welcome" .= u]

-- reactivation procedure (password reset)

reactivateForm
  :: ( YesodHmacKeccak db master
     , MonadHandler m
     , HandlerSite m ~ master
     ) => AForm m Username
reactivateForm = areq textField userSettings Nothing
  where
    userSettings = FieldSettings (SomeMessage MsgUsername) Nothing (Just "username") Nothing []

reactivateWidget
  :: YesodHmacKeccak db master
  => (Route Auth -> Route master) -> WidgetT master IO ()
reactivateWidget tm = do
  render <- getUrlRenderParams
  toWidgetHead $ $(jsFile jsPath) render
  ((_, widget), enctype) <- liftHandlerT $ runFormPost $ renderDivs reactivateForm
  [whamlet|
<div .reactivate>
  <form method="post" enctype=#{enctype} action=@{tm resetPasswordR}>
    ^{widget}
    <input type="submit" value=_{MsgSend}>
  |]

getReactivateR'
  :: YesodHmacKeccak db master
  => HandlerT Auth (HandlerT master IO) Html
getReactivateR' = do
  tm <- getRouteToParent
  lift $ defaultLayout $ do
    setTitleI MsgPasswordReset
    reactivateWidget tm

postReactivateR'
  :: YesodHmacKeccak db master
  => HandlerT Auth (HandlerT master IO) Html
postReactivateR' = do
  ((result, _), _) <- lift $ runFormPost $ renderDivs reactivateForm
  case result of
    FormMissing -> invalidArgs ["Form is missing"]
    FormFailure msg -> do
      lift $ setMessage $ toHtml $ T.concat msg
      redirect LoginR
    FormSuccess uname -> do
      muser <- lift $ runHmacDB $ loadUser uname
      case muser of
        Nothing -> do
          lift $ setMessageI MsgNoSuchUser
          redirect LoginR
        Just user -> do
          token <- liftIO makeRandomToken
          tm <- getRouteToParent
          lift $ runHmacDB $ insertActivateToken (encodeUtf8 token) uname
          render <- lift getUrlRender
          lift $ sendReactivateEmail uname (userUserEmail user) $
            render $ tm $ verifyR $ encodeUtf8 token
          lift $ setMessageI MsgActivationSent
          redirect LoginR

-- classes and foo

class UserCredentials u where
  userUserName   :: u -> Username
  userUserSalt   :: u -> ByteString
  userUserSalted :: u -> ByteString
  userUserEmail  :: u -> Text
  userUserActive :: u -> Bool

class TokenData t where
  tokenTokenKind     :: t -> Text
  tokenTokenUsername :: t -> Username
  tokenTokenToken    :: t -> ByteString

class PersistUserCredentials u where
  userUsernameF   :: EntityField u Username
  userUserSaltF   :: EntityField u ByteString
  userUserSaltedF :: EntityField u ByteString
  userUserEmailF  :: EntityField u Text
  userUserActiveF :: EntityField u Bool
  uniqueUsername  :: Text -> P.Unique u

  userCreate
    :: Username   -- ^ User name
    -> Text       -- ^ Email
    -> ByteString -- ^ User salt
    -> u

class PersistToken t where
  tokenTokenTokenF    :: EntityField t ByteString
  tokenTokenKindF     :: EntityField t Text
  tokenTokenUsernameF :: EntityField t Username
  uniqueToken         :: ByteString -> P.Unique t

  tokenCreate
    :: ByteString -- ^ actual Token
    -> Username   -- ^ User name
    -> Text       -- ^ Token kind
    -> t

class HmacDB m where
  type UserAccount m

  loadUser :: Username -> m (Maybe (UserAccount m))

  addNewUser :: Username -> Text -> ByteString -> m (Either Text (UserAccount m))

  activateUser :: (UserAccount m) -> ByteString -> m ()

  type TokenFoo m

  loadToken :: ByteString -> Text -> m (Maybe (TokenFoo m))

  insertToken :: ByteString -> Username -> Text -> m (Either Text (TokenFoo m))

  deleteToken :: (TokenFoo m) -> m ()

  insertLoginToken :: ByteString -> Username -> m (Either Text (TokenFoo m))
  insertLoginToken t u = insertToken t u "login"

  loadLoginToken :: ByteString -> m (Maybe (TokenFoo m))
  loadLoginToken = flip loadToken "login"

  insertActivateToken :: ByteString -> Username -> m (Either Text (TokenFoo m))
  insertActivateToken t u = insertToken t u "activate"

  loadActivateToken :: ByteString -> m (Maybe (TokenFoo m))
  loadActivateToken = flip loadToken "activate"

class HmacSendMail master where
  sendVerifyEmail
    :: Username -> Text -> Text -> HandlerT master IO ()

  sendReactivateEmail
    :: Username -> Text -> Text -> HandlerT master IO ()

class ( YesodAuth master
      , HmacSendMail master
      , HmacDB db
      , UserCredentials (UserAccount db)
      , TokenData (TokenFoo db)
      , RenderMessage master FormMessage
      ) => YesodHmacKeccak db master | master -> db where

  runHmacDB :: db a -> HandlerT master IO a

  checkValidUsername :: (MonadHandler m, HandlerSite m ~ master)
    => Username -> m (Either Text Username)
  checkValidUsername u
    | T.all C.isAlphaNum u = return $ Right u
    | otherwise = do
      mr <- getMessageRender
      return $ Left $ mr MsgInvalidUsername

  getNewAccountR :: HandlerT Auth (HandlerT master IO) Html
  getNewAccountR = getNewAccountR'

  postNewAccountR :: HandlerT Auth (HandlerT master IO) Html
  postNewAccountR = postNewAccountR'

  getReactivateR :: HandlerT Auth (HandlerT master IO) Html
  getReactivateR = getReactivateR'

  postReactivateR :: HandlerT Auth (HandlerT master IO) Html
  postReactivateR = postReactivateR'

  renderAccountMessage :: master -> [Text] -> AccountMsg -> Text
  renderAccountMessage _ _ = defaultAccountMsg

  rawLoginRoute :: Maybe (Route (HandlerSite (WidgetT master IO)))
  rawLoginRoute = Nothing

instance YesodHmacKeccak db master => RenderMessage master AccountMsg where
  renderMessage = renderAccountMessage

data PersistHmacFuncs master user token = PersistHmacFuncs
  { puGet :: Text -> HandlerT master IO (Maybe (Entity user))
  , puInsert :: Username -> user -> HandlerT master IO (Either Text (Entity user))
  , puUpdate :: Entity user -> [Update user] -> HandlerT master IO ()
  , ptGet :: ByteString -> Text -> HandlerT master IO (Maybe (Entity token))
  , ptInsert :: ByteString -> token -> HandlerT master IO (Either Text (Entity token))
  , ptUpdate :: Entity token -> [Update token] -> HandlerT master IO ()
  , ptDelete :: Entity token -> HandlerT master IO ()
  }

newtype HmacPersistDB master user token a
  = HmacPersistDB
      ( (ReaderT (PersistHmacFuncs master user token) (HandlerT master IO) a)
      ) deriving (Monad, MonadIO, Functor, Applicative)

instance (Yesod master, PersistUserCredentials user, PersistToken token)
  => HmacDB (HmacPersistDB master user token) where
  type UserAccount (HmacPersistDB master user token) = P.Entity user

  loadUser name = HmacPersistDB $ do
    f <- ask
    lift $ puGet f name

  addNewUser name email salt = HmacPersistDB $ do
    f <- ask
    lift $ puInsert f name $ userCreate name email salt

  activateUser user salted = HmacPersistDB $ do
    f <- ask
    lift $ puUpdate f user
      [ userUserSaltedF P.=. salted
      , userUserActiveF P.=. True
      ]

  type TokenFoo (HmacPersistDB master user token) = P.Entity token

  loadToken token kind = HmacPersistDB $ do
    f <- ask
    lift $ ptGet f token kind

  insertToken token uname kind = HmacPersistDB $ do
    f <- ask
    lift $ ptInsert f token $ tokenCreate token uname kind

  deleteToken t = HmacPersistDB $ do
    f <- ask
    lift $ ptDelete f t

makeRandomToken :: IO Text
makeRandomToken = (toHex . B.pack . take 16 . randoms) <$> newStdGen

makeRandomSalt :: IO Text
makeRandomSalt = (toHex . B.pack . take 8 . randoms) <$> newStdGen

returnJson :: Monad m => [Pair] -> m RepJson
returnJson = return . repJson . object

returnJsonError :: (ToJSON a, Monad m) => a -> m RepJson
returnJsonError = returnJson . (: []) . ("error" .=)

fromHex :: String -> BL.ByteString
fromHex = BL.pack . hexToWords
  where
    hexToWords (c:c':text) =
      let hex = [c, c']
          (word, _):_ = readHex hex
      in word : hexToWords text
    hexToWords _ = []

fromHex' :: String -> ByteString
fromHex' = B.concat . BL.toChunks . fromHex

toHex :: ByteString -> T.Text
toHex = T.pack . concatMap mapByte . B.unpack
  where
    mapByte = pad 2 '0' . flip showHex ""
    pad len padding s
      | length s < len = pad len padding $ padding:s
      | otherwise      = s

hmacKeccak :: ByteString -> ByteString -> ByteString
hmacKeccak key msg = BC.pack $ show $ hmacGetDigest (hmac key msg :: HMAC Keccak_512)

runHmacPersistDB
  :: ( Yesod master
     , PersistQueryRead b
     , PersistToken token
     , YesodPersist master
     , P.PersistEntity user
     , PersistUserCredentials user
     , P.PersistEntity token
     , PersistToken token
     , b ~ YesodPersistBackend master
#if MIN_VERSION_persistent(2,1,0)
     , b ~ PersistEntityBackend user
     , b ~ PersistEntityBackend token
     , PersistUnique b
#else
     , PersistMonadBackend (b (HandlerT master IO)) ~ P.PersistEntityBackend user
     , PersistMonadBackend (b (HandlerT master IO)) ~ P.PersistEntityBackend token
     , P.PersistUnique (b (HandlerT master IO))
     , P.PersistQuery (b (HandlerT master IO))
#endif
#if MIN_VERSION_persistent(2,5,0)
     , b ~ BaseBackend b
#endif
     , YesodHmacKeccak db master
     , db ~ HmacPersistDB master user token
     ) => HmacPersistDB master user token a -> HandlerT master IO a
runHmacPersistDB (HmacPersistDB m) = runReaderT m funcs
  where funcs =
          PersistHmacFuncs
            { puGet = runDB . P.getBy . uniqueUsername
            , puInsert = \name u -> do { mentity <- runDB $ P.insertBy u;
                mr <- getMessageRender;
                case mentity of
                  Left _ -> return $ Left $ mr $ MsgUsernameExists name;
                  Right k -> return $ Right $ P.Entity k u;
                }
            , puUpdate = \(P.Entity key _) u -> runDB $ P.update key u
            , ptGet = \token kind -> runDB $ P.selectFirst
                [ tokenTokenKindF ==. kind
                , tokenTokenTokenF ==. token
                ] []
            , ptInsert = \name t -> do { mentity <- runDB $ P.insertBy t;
                mr <- getMessageRender;
                case mentity of
                  Left _ -> return $ Left $ mr $ MsgInvalidToken;
                  Right k -> return $ Right $ P.Entity k t;
                }
            , ptUpdate = \(P.Entity key _) t -> runDB $ P.update key t
            , ptDelete = \(P.Entity key _) -> runDB $ P.delete key
            }