{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}

module Aws.Lambda.Wai
  ( runWaiAsLambda,
    runWaiAsProxiedHttpLambda,
    WaiLambdaProxyType (..),
    apiGatewayWaiHandler,
    ApiGatewayWaiHandler,
    albWaiHandler,
    ALBWaiHandler,
    ignoreALBPathPart,
    ignoreNothing,
    waiHandler,
    runMultipleWaiApplications,
  )
where

import Aws.Lambda
import Control.Concurrent.MVar
import Control.Monad (forM, forM_)
import Data.Aeson
import Data.Aeson.Types
import Data.Bifunctor (Bifunctor (bimap))
import qualified Data.Binary.Builder as Binary
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.CaseInsensitive as CI
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HMap
import Data.IORef
import qualified Data.IP as IP
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Text.Encoding (decodeUtf8', encodeUtf8)
import qualified Data.Text.Encoding as T
import qualified Data.Vault.Lazy as Vault
import qualified Network.HTTP.Types as H
import qualified Network.Socket as Socket
import Network.Wai (Application)
import qualified Network.Wai as Wai
import qualified Network.Wai.Internal as Wai
import qualified System.IO as IO
import Text.Read (readMaybe)

type ApiGatewayWaiHandler = ApiGatewayRequest Text -> Context Application -> IO (Either (ApiGatewayResponse Text) (ApiGatewayResponse Text))

type ALBWaiHandler = ALBRequest Text -> Context Application -> IO (Either (ALBResponse Text) (ALBResponse Text))

type GenericWaiHandler = Value -> Context Application -> IO (Either Value Value)

newtype ALBIgnoredPathPortion = ALBIgnoredPathPortion {unALBIgnoredPathPortion :: Text}

data WaiLambdaProxyType
  = APIGateway
  | ALB (Maybe ALBIgnoredPathPortion)

runWaiAsProxiedHttpLambda ::
  DispatcherOptions ->
  Maybe ALBIgnoredPathPortion ->
  HandlerName ->
  IO Application ->
  IO ()
runWaiAsProxiedHttpLambda options ignoredAlbPath handlerName mkApp =
  runLambdaHaskellRuntime options mkApp id $
    addStandaloneLambdaHandler handlerName (waiHandler ignoredAlbPath)

runMultipleWaiApplications ::
  DispatcherOptions ->
  HashMap HandlerName (Maybe ALBIgnoredPathPortion, IO Application) ->
  IO ()
runMultipleWaiApplications options handlersAndApps = do
  runLambdaHaskellRuntime options initializeApplications id $
    forM_ (HMap.keys handlersAndApps) $ \handler ->
      addStandaloneLambdaHandler handler $ \request context ->
        multiApplicationWaiHandler handler request context
  where
    initializeApplications :: IO (HashMap HandlerName (Maybe ALBIgnoredPathPortion, Application))
    initializeApplications = do
      HMap.fromList
        <$> forM
          (HMap.toList handlersAndApps)
          (\(handler, (alb, mkApp)) -> mkApp >>= \app -> return (handler, (alb, app)))

multiApplicationWaiHandler ::
  HandlerName ->
  Value ->
  Context (HashMap HandlerName (Maybe ALBIgnoredPathPortion, Application)) ->
  IO (Either Value Value)
multiApplicationWaiHandler handlerName request context = do
  appMay <- HMap.lookup handlerName <$> readIORef (customContext context)
  case appMay of
    Just (ignoredAlbPart, app) -> do
      applicationRef <- newIORef app
      waiHandler ignoredAlbPart request (context {customContext = applicationRef})
    Nothing ->
      fail $ "No application was registered for handler '" <> T.unpack (unHandlerName handlerName) <> "'."

waiHandler ::
  Maybe ALBIgnoredPathPortion ->
  GenericWaiHandler
waiHandler ignoredAlbPath request context =
  case parse parseIsAlb request of
    Success isAlb -> do
      if isAlb
        then case fromJSON @(ALBRequest Text) request of
          Success albRequest ->
            bimap toJSON toJSON <$> albWaiHandler ignoredAlbPath albRequest context
          Error err -> error $ "Could not parse the request as a valid ALB request: " <> err
        else case fromJSON @(ApiGatewayRequest Text) request of
          Success apiGwRequest ->
            bimap toJSON toJSON <$> apiGatewayWaiHandler apiGwRequest context
          Error err -> error $ "Could not parse the request as a valid API Gateway request: " <> err
    Error err ->
      error $
        "Could not parse the request as a valid API Gateway or ALB proxy request: " <> err
  where
    parseIsAlb :: Value -> Parser Bool
    parseIsAlb = withObject "Request" $ \obj -> do
      requestContextMay <- obj .:? "requestContext"
      case requestContextMay of
        Just requestContext -> do
          elb <- requestContext .:? "elb"
          case elb of
            Just (_ :: Value) -> pure True
            Nothing -> pure False
        Nothing -> pure False

runWaiAsLambda ::
  WaiLambdaProxyType ->
  DispatcherOptions ->
  HandlerName ->
  IO Application ->
  IO ()
runWaiAsLambda proxyType options handlerName mkApp = do
  case proxyType of
    APIGateway -> do
      IO.print $ "Starting Lambda using API gateway handler '" <> unHandlerName handlerName <> "'."
      runLambdaHaskellRuntime options mkApp id $ do
        addAPIGatewayHandler handlerName apiGatewayWaiHandler
    (ALB ignoredPath) -> do
      IO.print $ "Starting Lambda using ALB handler '" <> unHandlerName handlerName <> "'."
      runLambdaHaskellRuntime options mkApp id $ do
        addALBHandler handlerName (albWaiHandler ignoredPath)

ignoreALBPathPart :: Text -> Maybe ALBIgnoredPathPortion
ignoreALBPathPart = Just . ALBIgnoredPathPortion

ignoreNothing :: Maybe ALBIgnoredPathPortion
ignoreNothing = Nothing

albWaiHandler :: Maybe ALBIgnoredPathPortion -> ALBWaiHandler
albWaiHandler ignoredPathPortion request context = do
  waiApplication <- readIORef (customContext context)
  waiRequest <- mkWaiRequestFromALB ignoredPathPortion request

  (status, headers, body) <- processRequest waiApplication waiRequest >>= readResponse

  if BS.null body
    then return . pure . mkALBResponse (H.statusCode status) headers $ mempty
    else case decodeUtf8' body of
      Right responseBodyText ->
        return . pure . mkALBResponse (H.statusCode status) headers $ responseBodyText
      Left err -> error $ "Expected a response body that is valid UTF-8: " <> show err

apiGatewayWaiHandler :: ApiGatewayWaiHandler
apiGatewayWaiHandler request context = do
  waiApplication <- readIORef (customContext context)
  waiRequest <- mkWaiRequestFromApiGw request

  (status, headers, body) <- processRequest waiApplication waiRequest >>= readResponse

  if BS.null body
    then return . pure . mkApiGatewayResponse (H.statusCode status) headers $ mempty
    else case decodeUtf8' body of
      Right responseBodyText ->
        return . pure . mkApiGatewayResponse (H.statusCode status) headers $ responseBodyText
      Left err -> error $ "Expected a response body that is valid UTF-8: " <> show err

mkWaiRequestFromALB :: Maybe ALBIgnoredPathPortion -> ALBRequest Text -> IO Wai.Request
mkWaiRequestFromALB (fmap unALBIgnoredPathPortion -> pathPortionToIgnore) ALBRequest {..} = do
  let sourceIpMay = albRequestHeaders >>= HMap.lookup "x-forwarded-for"

  ip <- parseIp sourceIpMay

  let requestPath =
        case pathPortionToIgnore of
          Just toIgnore ->
            let toIgnoreSafe = "/" <> T.dropWhile (\c -> c == '/' || c == '\\') toIgnore
                throwPathError =
                  error $
                    "Given path piece to ignore '"
                      <> T.unpack toIgnoreSafe
                      <> "' is longer than the received request path "
                      <> T.unpack albRequestPath
                      <> "!"
             in fromMaybe throwPathError $ T.stripPrefix toIgnoreSafe albRequestPath
          Nothing -> albRequestPath

  -- TODO: Duplication
  let pathInfo = H.decodePathSegments (encodeUtf8 requestPath)

  let requestBodyRaw = maybe mempty T.encodeUtf8 albRequestBody
  let requestBodyLength = Wai.KnownLength $ fromIntegral $ BS.length requestBodyRaw

  requestBodyMVar <- newMVar requestBodyRaw

  let requestBody = takeRequestBodyChunk requestBodyMVar
  let headers = fromMaybe HMap.empty albRequestHeaders
  let requestHeaderHost = encodeUtf8 <$> HMap.lookup "host" headers
  let requestHeaderRange = encodeUtf8 <$> HMap.lookup "range" headers
  let requestHeaderReferer = encodeUtf8 <$> HMap.lookup "referer" headers
  let requestHeaderUserAgent = encodeUtf8 <$> HMap.lookup "User-Agent" headers

  let queryParameters = toQueryStringParameters albRequestQueryStringParameters
      rawQueryString = H.renderQuery True queryParameters
      httpVersion = H.http11 -- ALB converts even HTTP/2 requests to 1.1
  let result =
        Wai.Request
          (encodeUtf8 albRequestHttpMethod)
          httpVersion
          (encodeUtf8 requestPath)
          rawQueryString
          (map toHeader $ HMap.toList headers)
          True -- We assume it's always secure as we're passing through API Gateway
          ip
          pathInfo
          queryParameters
          requestBody
          Vault.empty
          requestBodyLength
          requestHeaderHost
          requestHeaderRange
          requestHeaderReferer
          requestHeaderUserAgent

  return result

mkWaiRequestFromApiGw :: ApiGatewayRequest Text -> IO Wai.Request
mkWaiRequestFromApiGw ApiGatewayRequest {..} = do
  let ApiGatewayRequestContext {..} = apiGatewayRequestRequestContext
      ApiGatewayRequestContextIdentity {..} = apiGatewayRequestContextIdentity

  ip <- parseIp apiGatewayRequestContextIdentitySourceIp

  let requestPath =
        -- We prefer the proxied path because apiGatewayRequestPath also
        -- includes the resource which we don't need
        case apiGatewayRequestPathParameters of
          Just pathParametersMap ->
            fromMaybe
              apiGatewayRequestPath
              (HMap.lookup "proxy" pathParametersMap)
          Nothing -> apiGatewayRequestPath

  let pathInfo = H.decodePathSegments (encodeUtf8 requestPath)

  let requestBodyRaw = maybe mempty T.encodeUtf8 apiGatewayRequestBody
  let requestBodyLength = Wai.KnownLength $ fromIntegral $ BS.length requestBodyRaw

  requestBodyMVar <- newMVar requestBodyRaw

  let requestBody = takeRequestBodyChunk requestBodyMVar
  let headers = fromMaybe HMap.empty apiGatewayRequestHeaders
  let requestHeaderHost = encodeUtf8 <$> HMap.lookup "host" headers
  let requestHeaderRange = encodeUtf8 <$> HMap.lookup "range" headers
  let requestHeaderReferer = encodeUtf8 <$> HMap.lookup "referer" headers
  let requestHeaderUserAgent = encodeUtf8 <$> HMap.lookup "User-Agent" headers

  let queryParameters = toQueryStringParameters apiGatewayRequestQueryStringParameters
      rawQueryString = H.renderQuery True queryParameters
      httpVersion = getHttpVersion apiGatewayRequestContextProtocol

  let result =
        Wai.Request
          (encodeUtf8 apiGatewayRequestHttpMethod)
          httpVersion
          (encodeUtf8 requestPath)
          rawQueryString
          (map toHeader $ HMap.toList headers)
          True -- We assume it's always secure as we're passing through API Gateway
          ip
          pathInfo
          queryParameters
          requestBody
          Vault.empty
          requestBodyLength
          requestHeaderHost
          requestHeaderRange
          requestHeaderReferer
          requestHeaderUserAgent

  return result

getHttpVersion :: Text -> H.HttpVersion
getHttpVersion protocol
  | "0.9" `T.isSuffixOf` protocol = H.http09
  | "1.0" `T.isSuffixOf` protocol = H.http10
  | "1.1" `T.isSuffixOf` protocol = H.http11
  | "2.0" `T.isSuffixOf` protocol = H.http20
  | otherwise = H.http11

takeRequestBodyChunk :: MVar ByteString -> IO ByteString
takeRequestBodyChunk requestBodyMVar = do
  result <- tryTakeMVar requestBodyMVar
  case result of
    Just bs -> pure bs
    Nothing -> pure BS.empty

toQueryStringParameters :: Maybe (HMap.HashMap Text Text) -> [H.QueryItem]
toQueryStringParameters (Just params) =
  let toQueryItem (key, value) = (encodeUtf8 key, Just $ encodeUtf8 value)
   in map toQueryItem $ HMap.toList params
toQueryStringParameters _ = []

parseIp :: Maybe Text -> IO Socket.SockAddr
parseIp sourceIpText =
  case sourceIpText of
    Just sourceIp ->
      case readMaybe (T.unpack sourceIp) of
        Just ip ->
          pure $ case ip of
            IP.IPv4 ip4 ->
              Socket.SockAddrInet
                0 -- default port
                (IP.toHostAddress ip4)
            IP.IPv6 ip6 ->
              Socket.SockAddrInet6
                0 -- default port
                0 -- flow info
                (IP.toHostAddress6 ip6)
                0 -- scope id
        Nothing -> error "Could not parse source ip."
    Nothing -> error "Missing source ip."

processRequest :: Application -> Wai.Request -> IO Wai.Response
processRequest app req = do
  mvar <- newEmptyMVar
  Wai.ResponseReceived <- app req $ \resp -> do
    putMVar mvar resp
    pure Wai.ResponseReceived
  takeMVar mvar

readResponse :: Wai.Response -> IO (H.Status, H.ResponseHeaders, ByteString)
readResponse (Wai.responseToStream -> (st, hdrs, mkBody)) = do
  body <- mkBody drainBody
  pure (st, hdrs, body)
  where
    drainBody :: Wai.StreamingBody -> IO ByteString
    drainBody body = do
      ioRef <- newIORef Binary.empty
      body
        (\b -> atomicModifyIORef ioRef (\b' -> (b <> b', ())))
        (pure ())
      BL.toStrict . Binary.toLazyByteString <$> readIORef ioRef

toHeader :: (Text, Text) -> H.Header
toHeader (name, val) = (CI.mk . encodeUtf8 $ name, encodeUtf8 val)