{-# LANGUAGE OverloadedStrings #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 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 SAML2.WebSSO.Servant where

import Data.EitherR
import Data.Function
import Data.List (nubBy)
import Data.Map qualified as Map
import Data.Proxy
import Data.String.Conversions
import Network.HTTP.Media ((//))
import Network.HTTP.Types
import Network.HTTP.Types.Header qualified as HttpTypes
import Network.Wai hiding (Response)
import Network.Wai.Internal as Wai
import SAML2.WebSSO.XML
import Servant.API as Servant hiding (URI (..))
import Text.Hamlet.XML
import Text.XML

type GetRedir = Verb 'GET 307

type PostRedir = Verb 'POST 303

-- | There is a tiny package `servant-xml`, which does essentially what this type and its
-- 'Mime{,Un}Render' instances do, but inlining this package seems easier.
--
-- FUTUREWORK: there is also `RawXML` in wire-api, we should probably drop one of the two.
data XML

instance Accept XML where
  contentType Proxy = "application" // "xml"

instance {-# OVERLAPPABLE #-} (HasXMLRoot a) => MimeRender XML a where
  mimeRender Proxy = cs . encode

instance {-# OVERLAPPABLE #-} (HasXMLRoot a) => MimeUnrender XML a where
  mimeUnrender Proxy = fmapL show . decode . cs

data HTML

instance Accept HTML where
  contentType Proxy = "text" // "html"

instance MimeRender HTML ST where
  mimeRender Proxy msg =
    mkHtml
      [xml|
      <body>
        <p>
          #{msg}
    |]

mkHtml :: [Node] -> LBS
mkHtml nodes = renderLBS def doc
  where
    doc = Document (Prologue [] (Just doctyp) []) root []
    doctyp = Doctype "html" (Just $ PublicID "-//W3C//DTD XHTML 1.1//EN" "http://www.w3.org/TR/xhtml11/DTD/xhtml11.dtd")
    root = Element "html" rootattr nodes
    rootattr = Map.fromList [("xmlns", "http://www.w3.org/1999/xhtml"), ("xml:lang", "en")]

-- | [3.5.5.1] Caching
setHttpCachePolicy :: Middleware
setHttpCachePolicy ap rq respond = ap rq $ respond . addHeadersToResponse httpCachePolicy
  where
    httpCachePolicy :: HttpTypes.ResponseHeaders
    httpCachePolicy = [("Cache-Control", "no-cache, no-store"), ("Pragma", "no-cache")]
    addHeadersToResponse :: HttpTypes.ResponseHeaders -> Wai.Response -> Wai.Response
    addHeadersToResponse extraHeaders resp = case resp of
      ResponseFile status hdrs filepath part -> ResponseFile status (updH hdrs) filepath part
      ResponseBuilder status hdrs builder -> ResponseBuilder status (updH hdrs) builder
      ResponseStream status hdrs body -> ResponseStream status (updH hdrs) body
      ResponseRaw action resp' ->
        ResponseRaw action $
          addHeadersToResponse extraHeaders resp'
      where
        updH hdrs = nubBy ((==) `on` fst) $ extraHeaders ++ hdrs
