{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}

module Endpoints.LoginEndpoint (LoginAPI, loginServer) where

import Servant
import Data.Aeson as A
import Database (runDb)
import Database.Persist
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import Crypto.BCrypt (validatePassword)
import qualified Data.ByteString.Base16 as Base16
import Jose.Jws (hmacEncode)
import Jose.Jwa (JwsAlg(HS256))
import Jose.Jwt (Jwt (Jwt))
import Data.ByteString.Lazy as BL
import Data.ByteString as B
import Crypto.Random (getRandomBytes)
import Data.ByteString.Base64.URL as BURL
import qualified Data.ByteString.UTF8 as UTF8
import Control.Monad.IO.Class (MonadIO(..), liftIO)

import Model.Login
import Model.MatrixErrorResponse
import Data.User
import Data.Device
import Model.AuthenticationHolder (AuthenticationHolder (AuthenticationHolder), server_password)

----------------------------------------------------------------------------------------------------
type LoginAPI = GetLogin :<|> PostLogin

loginServer :: Server LoginAPI
loginServer = handleLoginGet :<|> handleLoginPost

--- GET /_matrix/client/v3/login -------------------------------------------------------------------
type GetLogin = "_matrix" :> "client" :> "v3" :> "login"
              :> Get '[JSON] LoginFlowsResponse

handleLoginGet :: Handler LoginFlowsResponse
handleLoginGet = return $ LoginFlowsResponse [LoginFlow "m.login.password"]

--- POST /_matrix/client/v3/login ------------------------------------------------------------------
type PostLogin = "_matrix" :> "client" :> "v3" :> "login" :> ReqBody '[JSON] LoginRequest
               :> Post '[JSON] LoginResponse

handleLoginPost :: LoginRequest -> Handler LoginResponse
handleLoginPost LoginRequest
  { device_id = maybe_device_id
  , password = input_password
  , initial_device_display_name = _
  , identifier
  } = do
  resolved_device_id <- resolveDeviceId maybe_device_id
  let username = user identifier

  Jwt token <- eitherToHandler failed_token_generation $
    hmacEncode HS256 server_password $
      BL.toStrict $
        A.encode $
          AuthenticationHolder username resolved_device_id Nothing

  Entity _ db_user <- maybeToHandler invalid_username_error =<<
    liftIO (runDb $ getBy $ UniqueName username)

  hashedPasswordBytes <- eitherToHandler password_decoding_error $
    Base16.decode $
      T.encodeUtf8 $
        userPassword db_user

  unless (validatePassword hashedPasswordBytes (T.encodeUtf8 input_password)) $
    throw403 invalid_credentials_error

  pure LoginResponse
    { user_id      = T.concat ["@", username, ":localhost"]
    , access_token = T.decodeUtf8 token
    , home_server  = "localhost"
    , device_id    = resolved_device_id
    }
  where
    invalid_username_error = MatrixErrorResponse "M_FORBIDDEN" "Invalid username"
    password_decoding_error = MatrixErrorResponse "M_UNKNOWN" "Password hash decoding failed"
    invalid_credentials_error = MatrixErrorResponse "M_FORBIDDEN" "Invalid username or password"
    failed_token_generation = MatrixErrorResponse "M_UNKNOWN" "Access Token generation failed"

    resolveDeviceId :: MonadIO m => Maybe T.Text -> m T.Text
    resolveDeviceId device_id = case device_id of
      Just did -> pure did
      Nothing  -> liftIO $ generateDeviceId

    generateDeviceId :: IO T.Text
    generateDeviceId = do
      bytes <- getRandomBytes 12
      let encoded = BURL.encode bytes
      let generated_id = T.pack $ UTF8.toString $ B.take 16 encoded
      _ <- liftIO $ runDb $ insert $ Device generated_id maybe_device_name
      pure generated_id

maybeToHandler :: ServerErrorBody -> Maybe a -> Handler a
maybeToHandler err =
  maybe (throwError $ err403 { errBody = A.encode err }) pure

eitherToHandler :: ServerErrorBody -> Either b a -> Handler a
eitherToHandler err =
  either (const $ throwError $ err500 { errBody = A.encode err }) pure

throw403 :: ServerErrorBody -> Handler a
throw403 err =
  throwError $ err403 { errBody = A.encode err }