-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Wire.API.Federation.Endpoint
  ( ApplyMods,
    module Wire.API.Federation.Endpoint,
  )
where

import Control.Lens ((?~))
import Data.ByteString.Conversion (fromByteString)
import Data.Kind
import Data.Metrics.Servant
import Data.Misc (IpAddr)
import Data.OpenApi qualified as S
import Data.Sequence qualified as Seq
import GHC.TypeLits
import Imports
import Network.HTTP.Types qualified as HTTP
import Servant
import Servant.Client
import Servant.Client.Core
import Servant.OpenApi (HasOpenApi (toOpenApi))
import Servant.Server.Internal (MkContextWithErrorFormatter)
import Wire.API.ApplyMods
import Wire.API.Federation.API.Common
import Wire.API.Federation.Domain
import Wire.API.Federation.HasNotificationEndpoint
import Wire.API.Routes.Named

data Versioned v name

instance {-# OVERLAPPING #-} (RenderableSymbol a) => RenderableSymbol (Versioned v a) where
  renderSymbol = renderSymbol @a
  renderOperationId = renderOperationId @a

type family FedPath (name :: k) :: Symbol

type instance FedPath (name :: Symbol) = name

type instance FedPath (Versioned v name) = name

type UnnamedFedEndpointWithMods (mods :: [Type]) path input output =
  ( ApplyMods
      mods
      (path :> OriginDomainHeader :> ReqBody '[JSON] input :> Post '[JSON] output)
  )

type FedEndpointWithMods (mods :: [Type]) name input output =
  Named
    name
    (UnnamedFedEndpointWithMods mods (FedPath name) input output)

type FedEndpoint name input output = FedEndpointWithMods '[] name input output

type NotificationFedEndpointWithMods (mods :: [Type]) name path input =
  Named name (UnnamedFedEndpointWithMods mods path input EmptyResponse)

type NotificationFedEndpoint tag =
  MkNotificationFedEndpoint
    (NotificationMods tag)
    (NotificationPath tag)
    (NotificationVersionTag tag)
    (Payload tag)

type StreamingFedEndpoint name input output =
  Named
    name
    ( name
        :> OriginDomainHeader
        :> OriginIpHeader
        :> ReqBody '[JSON] input
        :> StreamPostWithRemoteIp NoFraming OctetStream output
    )

type family
  MkNotificationFedEndpoint
    (m :: [Type])
    (s :: Symbol)
    (v :: Maybe k)
    (p :: Type)

type instance
  MkNotificationFedEndpoint m s 'Nothing p =
    NotificationFedEndpointWithMods m s s p

type instance
  MkNotificationFedEndpoint m s ('Just v) p =
    NotificationFedEndpointWithMods m (Versioned v s) s p

type OriginIpHeaderName = "Wire-Origin-IP" :: Symbol

-- | The remote backend's origin IP is best-effort forensic metadata only.
--   Do not use for auth, policy, or attribution; cert identity is authoritative.
--   IP reflects only the socket peer at our edge (often LB/NAT/egress) and may be inaccurate.
data OriginIpHeader

instance (RoutesToPaths api) => RoutesToPaths (OriginIpHeader :> api) where
  getRoutes = getRoutes @api

instance (HasClient m api) => HasClient m (OriginIpHeader :> api) where
  type Client m (OriginIpHeader :> api) = Client m api
  clientWithRoute pm _ req = clientWithRoute pm (Proxy @api) req
  hoistClientMonad pm _ = hoistClientMonad pm (Proxy @api)

type OriginIpHeaderHasServer = Header' '[Strict] OriginIpHeaderName IpAddr

instance
  ( HasServer api context,
    HasContextEntry (MkContextWithErrorFormatter context) ErrorFormatters
  ) =>
  HasServer (OriginIpHeader :> api) context
  where
  type ServerT (OriginIpHeader :> api) m = Maybe IpAddr -> ServerT api m
  route _pa = route (Proxy @(OriginIpHeaderHasServer :> api))
  hoistServerWithContext _ pc nt s = hoistServerWithContext (Proxy :: Proxy api) pc nt . s

originIpHeaderName :: (IsString a) => a
originIpHeaderName = fromString $ symbolVal (Proxy @OriginIpHeaderName)

-- Header carrying the best-effort remote peer IP observed by our internal
-- federator when calling a remote backend. Forensic metadata only; do not use
-- for auth/policy or attribution.
type RemoteIpHeaderName = "Wire-Remote-IP" :: Symbol

remoteIpHeaderName :: (IsString a) => a
remoteIpHeaderName = fromString $ symbolVal (Proxy @RemoteIpHeaderName)

instance (HasOpenApi api) => HasOpenApi (OriginIpHeader :> api) where
  toOpenApi _ = desc $ toOpenApi (Proxy @api)
    where
      desc =
        S.allOperations
          . S.description
          ?~ ("Federated endpoints may include optional origin IP header: `" <> originIpHeaderName <> "`")

-- | A streaming POST combinator that behaves like Servant's 'StreamPost',
--   but whose client additionally returns an optional remote IP parsed from
--   the 'Wire-Remote-IP' response header.
data StreamPostWithRemoteIp framing (ct :: Type) a

-- Server-side simply delegates to the standard 'StreamPost' implementation.
instance
  (HasServer (StreamPost framing ct a) context) =>
  HasServer (StreamPostWithRemoteIp framing ct a) context
  where
  type ServerT (StreamPostWithRemoteIp framing ct a) m = ServerT (StreamPost framing ct a) m
  route _ = route (Proxy @(StreamPost framing ct a))
  hoistServerWithContext _ = hoistServerWithContext (Proxy @(StreamPost framing ct a))

-- OpenAPI, metrics and path routing can delegate to the underlying StreamPost
instance RoutesToPaths (StreamPostWithRemoteIp (framing :: Type) (ct :: Type) (a :: Type)) where
  getRoutes = getRoutes @(StreamPost framing ct a)

instance (HasOpenApi (StreamPost framing ct a)) => HasOpenApi (StreamPostWithRemoteIp framing ct a) where
  toOpenApi _ = toOpenApi (Proxy @(StreamPost framing ct a))

-- Client-side: make the streaming request and return the body together with an
-- optional 'IpAddr' parsed from the 'Wire-Remote-IP' header.
instance
  ( RunStreamingClient m,
    Accept ct,
    FromSourceIO ByteString a
  ) =>
  HasClient m (StreamPostWithRemoteIp framing ct a)
  where
  type Client m (StreamPostWithRemoteIp framing ct a) = m (Maybe IpAddr, a)

  clientWithRoute _ _ req =
    withStreamingRequest
      req
        { requestMethod = HTTP.methodPost,
          requestAccept = Seq.singleton (contentType (Proxy @ct))
        }
      $ \resp -> do
        let hdrs = toList (responseHeaders resp)
            mIp = lookup remoteIpHeaderName hdrs >>= fromByteString
        body <- fromSourceIO (responseBody resp)
        pure (mIp, body)

  hoistClientMonad _ _ f c = f c
