{-# LANGUAGE DataKinds, TypeOperators, OverloadedStrings, ScopedTypeVariables, TypeApplications, RecordWildCards, AllowAmbiguousTypes #-}

module Server
{-( insecureOAuthMock'
, runMockServer
-- , runMockServer'
)-} where

import AuthCode
import User

import Control.Applicative ((<|>))
import Control.Concurrent
import Control.Concurrent.STM (atomically)
import Control.Concurrent.STM.TVar (newTVarIO, readTVar)
import Control.Exception (bracket)
import Control.Monad (unless)
import Control.Monad.IO.Class
import Control.Monad.Trans.Reader

import Data.Aeson
import Data.ByteString (ByteString (..), fromStrict, toStrict)
import Data.List (find, elemIndex)
import Data.Maybe (fromMaybe, isJust)
import Data.String (IsString (..))
import Data.Text hiding (elem, find, head, length, map, null, splitAt, tail, words)
import Data.Text.Encoding (decodeUtf8)
import Data.Time.Clock (NominalDiffTime (..), nominalDay, UTCTime(..), getCurrentTime, addUTCTime)

import qualified Data.ByteString.Char8 as BS
import qualified Data.Map.Strict as Map

import GHC.Read (readPrec, lexP)

import Jose.Jwa
import Jose.Jwe
import Jose.Jwk (generateRsaKeyPair, generateSymmetricKey, KeyUse(Enc), KeyId)
import Jose.Jwt hiding (decode, encode)

import Network.HTTP.Client (newManager, defaultManagerSettings)
import Network.Wai.Handler.Warp

import Servant
import Servant.Client
import Servant.API

import Text.ParserCombinators.ReadPrec (look, pfail)

import qualified Text.Read.Lex as Lex


testUsers :: [User] -- TODO move to db
testUsers =
  [ User {name = "Fallback User", email = "foo@bar.com", password = "0000", uID = "0"}
  , User {name = "Tina Tester", email = "t@t.tt", password = "1111", uID = "1"}
  , User {name = "Max Muster", email = "m@m.mm", password = "2222", uID = "2"}]

data AuthClient = Client
  { ident :: Text
  , secret :: Text
  } deriving (Show, Eq)

trustedClients :: [AuthClient] -- TODO move to db
trustedClients = [Client "42" "shhh"]

data ResponseType = Code -- ^ authorisation code grant
                  | Token -- ^ implicit grant via access token
                  | IDToken -- ^ implicit grant via access token & ID token
                  deriving (Eq, Show)
instance Read ResponseType where
  readPrec = do
    Lex.Ident str <- lexP
    Lex.EOF <- lexP
    case str of
      "code" -> return Code
      "token" -> return Token
      "id_token" -> return IDToken
      _ -> pfail

 ------------------------------
---- Authorisation endpoint ----
 ------------------------------

type QScope    = String
type QClient   = String
type QResType  = String
type QRedirect = String
type QState    = Text

type QParam = QueryParam' [Required, Strict]

type Auth user userData = BasicAuth "login" user
                       :> "auth"
                       :> QParam "scope" QScope
                       :> QParam "client_id" QClient
                       :> QParam "response_type" QResType
                       :> QParam "redirect_uri" QRedirect
                       :> QueryParam "state" QState
                       :> Get '[JSON] userData

-- type Insert = "insert" :> Post '[JSON] User


type AuthHandler = ReaderT AuthState Handler
type AuthServer a = ServerT a AuthHandler

toHandler :: AuthState -> AuthHandler a -> Handler a
toHandler s h = runReaderT h s

authServer :: forall user userData . UserData user userData => AuthServer (Auth user userData)
authServer = handleAuth
  where
    handleAuth :: user
               -> QScope
               -> QClient
               -> QResType
               -> QRedirect
               -> Maybe QState
               -> AuthHandler userData
    handleAuth u scopes client responseType url mState = do
      unless (isJust $ find (\c -> ident c == pack client) trustedClients) . -- TODO fetch trusted clients from db | TODO also check if the redirect url really belongs to the client
        throwError $ err404 { errBody = "Not a trusted client."}
      let responseType' = read @ResponseType responseType
      liftIO $ print responseType'
      unless (responseType' == Code) $ throwError err500 { errBody = "Unsupported response type" }  
      mAuthCode <- asks (genUnencryptedCode client url 600) >>= liftIO
      liftIO $ print mAuthCode
      -- liftIO . putStrLn $ "user: " ++ show u ++ " | scopes: " ++ show (map (showScope @user @userData) scopes')
      redirect $ addParams url mAuthCode mState
    redirect :: Maybe ByteString -> AuthHandler userData
    redirect (Just url) = liftIO (print url) >> throwError err303 { errHeaders = [("Location", url)]}
    redirect Nothing = throwError err500 { errBody = "Could not generate authorisation code."}
    addParams :: String -> Maybe Text -> Maybe Text -> Maybe ByteString
    addParams url Nothing _ = Nothing
    addParams url (Just code) mState =
      let qPos = fromMaybe (length url) $ elemIndex '?' url
          (pre, post) = splitAt qPos url
          rState = case mState of {Just s -> "&state=" ++ (unpack . replace "/" "%2F" $ replace "=" "%3D" s); Nothing -> ""}
          post' = if not (null post) then '&' : tail post else post
      in Just . fromString $ pre ++ "?code=" ++ (unpack code) ++ post' ++ rState


 ----------------------
---- Token Endpoint ----
 ----------------------


data ClientData = ClientData
  { grantType :: GrantType
  , grant :: String
  , userName :: Maybe String
  , clientID :: String
  , clientSecret :: String
  , redirect :: String
  } deriving Show

instance FromJSON ClientData where
  parseJSON (Object o) = ClientData
    <$> o .:  "grant_type"
    <*>  (o .: ps PassCreds --TODO add alternatives
      <|> o .: ps AuthCode
      <|> o .: ps Implicit
      <|> o .: ps ClientCreds)
    <*> o .:? "username"
    <*> o .:  "client_id"
    <*> o .:  "client_secret"
    <*> o .:  "redirect_url"
      where ps = fromString @Key . show

data GrantType = PassCreds | AuthCode | Implicit | ClientCreds deriving (Eq)

instance Show GrantType where
  show PassCreds = "password"
  show AuthCode  = "authorization_code"
  show _         = undefined --TODO support other flows

instance FromJSON GrantType where
  parseJSON (String s)
    | s == pack (show AuthCode) = pure AuthCode
    | otherwise = error $ show s ++ " grant type not supported yet"

data JWT = JWT
  { issuer :: Text
  , expiration :: UTCTime
  } deriving (Show, Eq)

instance ToJSON JWT where
  toJSON (JWT i e) = object ["iss" .= i, "exp" .= e]

data JWTWrapper = JWTW
  { token :: String
  , expiresIn :: NominalDiffTime
  } deriving (Show)

instance ToJSON JWTWrapper where
  toJSON (JWTW t e) = object ["access_token" .= t, "token_type" .= ("JWT" :: Text), "expires_in" .= e]

instance FromJSON JWTWrapper where
  parseJSON (Object o) = JWTW
    <$> o .: "access_token"
    <*> o .: "expires_in"

instance FromHttpApiData JWTWrapper where
  parseHeader bs = case decode (fromStrict bs) of
    Just x -> Right x
    Nothing -> Left "Invalid JWT wrapper"

type Token = "token"
          :> ReqBody '[JSON] ClientData
          :> Post '[JSON] JWTWrapper

tokenEndpoint :: AuthServer Token
tokenEndpoint = provideToken
  where
    provideToken :: ClientData -> AuthHandler JWTWrapper
    provideToken client = case (grantType client) of
      AuthCode -> do
        unless (Client (pack $ clientID client) (pack $ clientSecret client) `elem` trustedClients) .
          throwError $ err500 { errBody = "Invalid client" }
        valid <- asks (verify (pack $ grant client) (clientID client)) >>= liftIO
        unless valid . throwError $ err500 { errBody = "Invalid authorisation code" }
        -- return JWT {token = "", tokenType = "JWT", expiration = 0.25 * nominalDay}
        token <- asks mkToken >>= liftIO
        liftIO . putStrLn $ "token: " ++ show token
        return token 
      x -> error $ show x ++ " not supported yet"


mkToken :: AuthState -> IO JWTWrapper
mkToken state = do
  privateKey <- atomically $ readTVar state >>= return . privateKey
  now <- getCurrentTime
  let
    lifetime = nominalDay / 4 -- TODO make configurable
    jwt = JWT "Oauth2MockServer" (lifetime `addUTCTime` now)
  encoded <- jwkEncode RSA_OAEP_256 A128GCM privateKey (Nested . Jwt . toStrict $ encode jwt)
  case encoded of
    Right (Jwt token) -> return $ JWTW (BS.unpack token) lifetime
    Left e -> error $ show e


 ----------------------
---- Users Endpoint ----
 ----------------------


type Users = "users"

type Me userData = Users
                :> "me"
                :> Header "Authorization" JWTWrapper
                :> Get '[JSON] userData

type UserList userData = Users
                      :> "query"
                      :> Header "Authorization" JWTWrapper
                      :> Get '[JSON] [userData] -- TODO support query params


userEndpoint :: forall user userData . UserData user userData => AuthServer (Me userData)
userEndpoint = handleUserData
  where
    handleUserData :: Maybe JWTWrapper -> AuthHandler userData
    handleUserData jwtw = do
      undefined
      -- let
      --   scopes' = map (readScope @user @userData) $ words scopes
      --   uData = mconcat $ map (userScope @user @userData u) scopes'
      -- liftIO . putStrLn $ "user: " ++ show u ++ " | scopes: " ++ show (map (showScope @user @userData) scopes')
      -- return uData

userListEndpoint :: forall user userData . UserData user userData => AuthServer (UserList userData)
userListEndpoint = handleUserData
  where
    handleUserData :: Maybe JWTWrapper -> AuthHandler [userData]
    handleUserData jwtw = do
      undefined


 -------------------
---- Server Main ----
 -------------------

type Routing user userData = Auth user userData
                        :<|> Token
                        :<|> Me userData
                        :<|> UserList userData

routing :: forall user userData . UserData user userData => AuthServer (Routing user userData)
routing = authServer @user @userData
     :<|> tokenEndpoint
     :<|> userEndpoint @user @userData
     :<|> userListEndpoint @user @userData

exampleAuthServer :: AuthServer (Routing User (Map.Map Text Text))
exampleAuthServer = routing


authAPI :: Proxy (Routing User (Map.Map Text Text))
authAPI = Proxy

-- insecureOAuthMock :: Application
-- insecureOAuthMock = authAPI `serve` exampleAuthServer

insecureOAuthMock' :: [User] -> AuthState -> Application
insecureOAuthMock' testUsers s = serveWithContext authAPI c $ hoistServerWithContext authAPI p (toHandler s) exampleAuthServer
  where
    c = authenticate testUsers :. EmptyContext
    p = Proxy :: Proxy '[BasicAuthCheck User]

authenticate :: [User] -> BasicAuthCheck User
authenticate users = BasicAuthCheck $ \authData -> do
  let
    (uEmail, uPass) = (,) <$> (decodeUtf8 . basicAuthUsername) <*> (decodeUtf8 . basicAuthPassword) $ authData
  case (find (\u -> email u == uEmail) users) of
    Nothing -> return NoSuchUser
    Just u  -> return $ if uPass == password u then Authorized u else BadPassword

-- frontend :: BasicAuthData -> ClientM (Map.Map Text Text)
-- frontend ba = client authAPI ba "[ID]" "42" "code" ""

runMockServer :: Int -> IO ()
runMockServer port = do
  state <- mkState
  run port $ insecureOAuthMock' testUsers state

-- runMockServer' :: Int -> IO ()
-- runMockServer' port = do
--   mgr <- newManager defaultManagerSettings
--   state <- mkState
--   bracket (forkIO . run port $ insecureOAuthMock' testUsers state) killThread $ \_ ->
--     runClientM (frontend $ BasicAuthData "foo@bar.com" "0000") (mkClientEnv mgr (BaseUrl Http "localhost" port ""))
--       >>= print

mkState :: IO AuthState
mkState = do
  (publicKey, privateKey) <- generateRsaKeyPair 256 (KeyId "Oauth2MockKey") Enc Nothing
  let activeCodes = Map.empty
  newTVarIO State{..}