{-# LANGUAGE
  CPP
, RankNTypes
, 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 Numeric (readHex, showHex)

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

import Text.Julius (jsFile)

import Yesod.Auth.JsPath

-- | type alias
type Username = Text

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

  -- | function for accessing the database (runDB eqivalent).
  -- can be set to 'runHmacPersistDB'
  runHmacDB :: db a -> AuthHandler master a
  -- runHmacDB = runHmacPersistDB

  -- | function to determine a valid username.
  -- Default: 'defaultCheckValidUsername'
  checkValidUsername :: (MonadHandler m, HandlerSite m ~ master)
    => Username -> m (Either Text Username)
  checkValidUsername = defaultCheckValidUsername

  -- | Handler for rendering the registration page.
  -- Default: 'getNewAccountR''
  getNewAccountR :: AuthHandler master Html
  getNewAccountR = getNewAccountR'

  -- | Handler for processing registration.
  -- Default: 'postNewAccountR''
  postNewAccountR :: AuthHandler master Html
  postNewAccountR = postNewAccountR'

  -- | Handler for rendering reactivation request page.
  -- Default: 'getReactivateR''
  getReactivateR :: AuthHandler master Html
  getReactivateR = getReactivateR'

  -- | Handler for processing reactivation requests.
  -- Default: 'postReactivateR''
  postReactivateR :: AuthHandler master Html
  postReactivateR = postReactivateR'

  -- | Function for rendering all messages in this plugin.
  -- Default: 'defaultAccountMsg'
  renderAccountMessage :: master -> [Text] -> AccountMsg -> Text
  renderAccountMessage _ _ = defaultAccountMsg

  -- | Route for providing login without javascript.
  -- Default: 'Nothing'
  rawLoginRoute :: Maybe (Route (HandlerSite (WidgetFor master)))
  rawLoginRoute = Nothing

  -- | Widget for the login page.
  -- Default: 'defaultLoginWidget'
  loginWidget
    :: YesodHmacKeccak db master
    => (Route Auth -> Route master) -> WidgetFor master ()
  loginWidget = defaultLoginWidget

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.

-- | Overridable default login widget
defaultLoginWidget
    :: YesodHmacKeccak db master
    => (Route Auth -> Route master) -> WidgetFor master ()
defaultLoginWidget 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}
  |]

-- | Overridable default check for valid usernames
defaultCheckValidUsername
  :: ( MonadHandler m
     , HandlerSite m ~ master
     , YesodHmacKeccak db master
     )
  => Username -> m (Either Text Username)
defaultCheckValidUsername u
  | T.all C.isAlphaNum u = return $ Right u
  | otherwise = do
    mr <- getMessageRender
    return $ Left $ mr MsgInvalidUsername

postLoginR'
  :: ( YesodHmacKeccak db master
     , YesodAuth master
     )
  => AuthHandler master RepJson
postLoginR' = do
  mr <- getMessageRender
  mUserName <- lookupPostParam "username"
  mHexToken <- lookupPostParam "token"
  mHexResponse <- lookupPostParam "response"
  case (mUserName, mHexToken, mHexResponse) of
    (Just userName, Nothing, Nothing) -> do
      tempUser <- runHmacDB $ loadUser userName
      case tempUser of
        Just u ->
          if userUserActive u
          then do
            let salt = userUserSalt u
            token <- liftIO makeRandomToken
            _ <- 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 <- runHmacDB $ loadLoginToken tempToken
        case savedToken of
          Just token -> do
            queriedUser <- 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 !!
              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
          setCreds False $ Creds "authHmacKeccak" (userUserName au) []
          render <- getUrlRender
          m <- 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) -> WidgetFor master ()
newAccountWidget tm = do
  render <- getUrlRenderParams
  toWidgetHead $ $(jsFile jsPath) render
  ((_, widget), enctype) <- runFormPost $ renderDivs newAccountForm
  [whamlet|
<div .newaccount>
  <form method="post" enctype=#{enctype} action=@{tm newAccountR}>
    ^{widget}
    <input type=submit value=_{MsgRegister}>
  |]

getNewAccountR'
  :: YesodHmacKeccak db master
  => AuthHandler master Html
getNewAccountR' = do
  tm <- getRouteToParent
  authLayout $ do
    setTitleI MsgRegisterLong
    newAccountWidget tm

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

createNewAccount
  :: YesodHmacKeccak db master
  => NewAccountData
  -> AuthHandler master (UserAccount db)
createNewAccount NewAccountData{..} = do
  muser <- runHmacDB $ loadUser naUsername
  tm <- getRouteToParent
  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 -> WidgetFor master ()
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 -> AuthHandler master Html
getVerifyR' k = do
  mtoken <- runHmacDB $ loadActivateToken k
  tm <- getRouteToParent
  case mtoken of
    Nothing -> do
      setMessageI MsgInvalidToken
      redirect $ tm LoginR
    Just token -> do
      muser <- runHmacDB $ loadUser $ tokenTokenUsername token
      case muser of
        Nothing -> do
          setMessageI MsgNoSuchUser
          redirect $ tm LoginR
        Just user -> do
          let hexSalt = toHex $ userUserSalt user
          authLayout $ do
            setTitleI MsgSetPassword
            passwordWidget tm (tokenTokenToken token) (BC.pack $ T.unpack hexSalt)

postVerifyR'
  :: YesodHmacKeccak db master
  => ByteString -> AuthHandler master RepJson
postVerifyR' k = do
  mtoken <- runHmacDB $ loadActivateToken k
  tm <- getRouteToParent
  case mtoken of
    Nothing -> do
      setMessageI MsgInvalidToken
      redirect $ tm LoginR
    Just token -> do
      muser <- runHmacDB $ loadUser $ tokenTokenUsername token
      case muser of
        Nothing -> do
          setMessageI MsgNoSuchUser
          redirect $ tm LoginR
        Just user -> do
          msalted <- lookupPostParam "salted"
          case msalted of
            Nothing -> do
              setMessageI MsgProtocolError
              redirect $ tm LoginR
            Just salted' -> do
              let salted = fromHex' $ T.unpack salted'
              runHmacDB $ activateUser user salted
              runHmacDB $ deleteToken token
              setCreds False $ Creds "authHmacKeccak" (tokenTokenUsername token) []
              render <- getUrlRender
              m <- 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) -> WidgetFor master ()
reactivateWidget tm = do
  render <- getUrlRenderParams
  toWidgetHead $ $(jsFile jsPath) render
  ((_, widget), enctype) <- runFormPost $ renderDivs reactivateForm
  [whamlet|
<div .reactivate>
  <form method="post" enctype=#{enctype} action=@{tm resetPasswordR}>
    ^{widget}
    <input type="submit" value=_{MsgSend}>
  |]

getReactivateR'
  :: YesodHmacKeccak db master
  => AuthHandler master Html
getReactivateR' = do
  tm <- getRouteToParent
  authLayout $ do
    setTitleI MsgPasswordReset
    reactivateWidget tm

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

-- classes and foo

-- | Class for providing user credentials to the plugin. A user type is required
-- to have all these fields.
class UserCredentials u where
  userUserName   :: u -> Username
  userUserSalt   :: u -> ByteString
  userUserSalted :: u -> ByteString
  userUserEmail  :: u -> Text
  userUserActive :: u -> Bool

-- | Class for providing tokens for user activation.
class TokenData t where
  tokenTokenKind     :: t -> Text
  tokenTokenUsername :: t -> Username
  tokenTokenToken    :: t -> ByteString

-- | Class for defining the accessor functions of the database for users
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

  -- | create a new user ready for activation from provided data
  userCreate
    :: Username   -- ^ User name
    -> Text       -- ^ Email
    -> ByteString -- ^ User salt
    -> u

-- | Class for defining the accessor functions of the database for tokens
class PersistToken t where
  tokenTokenTokenF    :: EntityField t ByteString
  tokenTokenKindF     :: EntityField t Text
  tokenTokenUsernameF :: EntityField t Username
  uniqueToken         :: ByteString -> P.Unique t

  -- | create a new token from provided data
  tokenCreate
    :: ByteString -- ^ actual Token
    -> Username   -- ^ User name
    -> Text       -- ^ Token kind
    -> t

-- | This class lets you define you own database behaviour, but it comes pre-
-- defined with sane defaults.
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 -> AuthHandler master ()

  sendReactivateEmail
    :: Username -> Text -> Text -> AuthHandler master ()

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

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

newtype HmacPersistDB master user token a
  = HmacPersistDB
      ( (ReaderT (PersistHmacFuncs master user token) (HandlerFor master) 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
  :: ( PersistEntityBackend token ~ BaseBackend (YesodPersistBackend master)
     , PersistEntityBackend user ~ BaseBackend (YesodPersistBackend master)
     , PersistToken token, PersistEntity token, PersistEntity user
     , YesodPersist master
     , PersistUserCredentials user
     , PersistUniqueWrite (YesodPersistBackend master)
     , YesodHmacKeccak db master
     , PersistQueryRead (YesodPersistBackend master))
  => HmacPersistDB master user token a -> HandlerFor master a
runHmacPersistDB (HmacPersistDB master) = runReaderT master 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 = \_ 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
            }