{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | Timestamp stamping functionality for OpenTimestamps.

This module provides functions to create timestamps by submitting
digests to calendar servers and handling the results,
including error handling and timestamp file creation.
-}
module OpenTimestamps.Stamp
  ( TimestampError (..)
  , StampResult (..)
  , stampDigest
  ) where

import Control.Exception (SomeException, try)
import Crypto.Hash (SHA256 (SHA256), hashWith)
import qualified Crypto.Random.Entropy as R
import Data.Binary.Get (runGetOrFail)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy.Char8 as BSL
import Data.Either (lefts, rights)
import Data.Function ((&))
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Network.HTTP.Client (Request (..), RequestBody (..), Response (..))
import Network.HTTP.Simple
  ( getResponseBody
  , getResponseStatusCode
  , httpLBS
  , parseRequest
  , setRequestBody
  , setRequestHeader
  , setRequestMethod
  )
import OpenTimestamps.Config as Config
import qualified OpenTimestamps.DetachedTimestampFile as DTSF
import qualified OpenTimestamps.Op as Op
import qualified OpenTimestamps.Timestamp as TS
import OpenTimestamps.Types (OTsBytes)

-- | Error type for timestamp operations.
data TimestampError where
  TimestampError ::
    { errorDigest :: BS.ByteString
    , errorUrl :: String
    , errorMessage :: String
    } ->
    TimestampError
  deriving (Show, Eq)

-- | Result of a stamping operation containing any errors and the timestamp file.
data StampResult where
  StampResult ::
    { srErrors :: [TimestampError]
    , srTimestampFile :: Maybe DTSF.DetachedTimestampFile
    } ->
    StampResult
  deriving (Show)

-- | Send a timestamp request to a calendar server.
timeStampRequest ::
  BS.ByteString ->
  String ->
  Request ->
  IO (Either TimestampError TS.Timestamp)
timeStampRequest digest url request = do
  let postRequest =
        request
          & setRequestMethod "POST"
          & setRequestHeader "User-Agent" [Config.userAgent]
          & setRequestHeader "Accept" [Config.accept]
          & setRequestHeader "Content-Type" [Config.contentType]
          & setRequestBody (RequestBodyBS digest)
  eResponse <-
    try (httpLBS postRequest) ::
      IO (Either SomeException (Response BSL.ByteString))
  case eResponse of
    Left ex -> do
      let errMsg = "HTTP request failed: " ++ show ex
      pure $ Left $ TimestampError digest url errMsg
    Right response -> do
      let body = getResponseBody response
      if getResponseStatusCode response == 200
        then do
          case runGetOrFail (TS.deserialize digest) body of
            Left err -> do
              let errMsg = "Deserialization failed: " ++ show err
              pure $ Left $ TimestampError digest url errMsg
            Right (_, _, newTimestamp) -> do
              pure $ Right newTimestamp
        else do
          let errMsg =
                "Server returned HTTP "
                  ++ show (getResponseStatusCode response)
                  ++ ": "
                  ++ show body
          pure $ Left $ TimestampError digest url errMsg

{- | Stamp a digest by submitting it to calendar servers.

Takes a list of calendar server URLs and a digest to timestamp.
Returns a StampResult containing any errors and the timestamp file.
-}
stampDigest ::
  [String] ->
  OTsBytes ->
  IO StampResult
stampDigest
  urls
  originalDigest = do
    -- 1. Generate a random nonce (16 bytes)
    nonce <- R.getEntropy 16

    -- 2. Append nonce to the original digest
    let appendedDigest = originalDigest <> nonce

    -- 3. Hash the result of the append operation
    let submissionDigest = BA.convert (hashWith SHA256 appendedDigest)

    -- 4. Submit the new digest to the calendar servers
    allResults <- mapM (getTimestamp submissionDigest) urls
    let timestamps = rights allResults
    let errors = lefts allResults

    if null timestamps
      then
        pure $
          StampResult
            { srErrors = errors
            , srTimestampFile = Nothing
            }
      else do
        -- 5. Merge the proofs from the servers
        let serverTimestamp = foldl1 (\ts1 ts2 -> fst (TS.merge ts1 ts2)) timestamps

        -- 6. Construct the final timestamp with the correct structure
        let finalTimestamp =
              TS.Timestamp
                { TS.timestampMsg = originalDigest
                , TS.attestations = Set.empty
                , TS.ops =
                    Map.singleton (Op.Append nonce) $
                      TS.Timestamp
                        { TS.timestampMsg = appendedDigest
                        , TS.attestations = Set.empty
                        , TS.ops = Map.singleton Op.Sha256 serverTimestamp
                        }
                }

        -- 7. Create the final .ots file
        let newOtsFile = DTSF.DetachedTimestampFile DTSF.DSha256 finalTimestamp
        pure $
          StampResult
            { srErrors = errors
            , srTimestampFile = Just newOtsFile
            }
    where
      -- \| Get a timestamp from a specific calendar server.
      getTimestamp ::
        BS.ByteString ->
        String ->
        IO (Either TimestampError TS.Timestamp)
      getTimestamp digest' url' = do
        let url = url' ++ Config.apiUrlDigest
        eRequest <- try (parseRequest url) :: IO (Either SomeException Request)
        case eRequest of
          Left ex -> do
            let errMsg = "Error parsing request: " ++ show ex
            pure $ Left $ TimestampError digest' url errMsg
          Right request -> timeStampRequest digest' url' request