module Main (main) where

import           Control.Concurrent (threadDelay)
import           Control.Concurrent.Async (async, concurrently, wait)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import           Nuthatch.Probe.Assert
import           Nuthatch.Probe.Curl
import           Nuthatch.Probe.Http
import           Nuthatch.Probe.Oidc
import           Nuthatch.Probe.ProductClient
import           System.Exit (ExitCode (..))
import           System.IO (hClose)
import           System.IO.Temp
  ( withSystemTempDirectory
  , withSystemTempFile
  )
import           System.Process
  ( CreateProcess (..)
  , StdStream (..)
  , createProcess
  , proc
  , waitForProcess
  )

main :: IO ()
main =
  withSystemTempDirectory "nuthatch-idempotency-probe" $ \directory -> do
    context <- newProbeContext
    user <- requireEnv "NUTHATCH_OIDC_USER_A"
    session <- login context False user "idempotency"
    lane <- requireEnv "NUTHATCH_FIXTURE_LANE"
    ipA <- requireEnv "NUTHATCH_IDEMPOTENCY_IP_A"
    ipB <- requireEnv "NUTHATCH_IDEMPOTENCY_IP_B"
    verifyEpochs context session lane ipA ipB
    verifyActiveIdentity context session lane ipA ipB directory
    emitRefinementTrace
    putStrLn "Cross-process idempotency epochs passed"

verifyEpochs
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> IO ()
verifyEpochs context session lane ipA ipB = do
  let operationId = "73000000-0000-4000-8000-000000000001"
      firstBody = "first-epoch"
  first <- targetedUpload context session lane ipA operationId firstBody
  requireSucceededWindow "first epoch" first
  replay <- targetedUpload context session lane ipB operationId firstBody
  requireSucceededWindow "cross-process replay" replay
  conflict <- targetedUpload context session lane ipB operationId "conflict"
  assertContains "conflicting reuse"
    "\"result\":\"conflict\""
    (probeBody conflict)
  threadDelay 3000000
  (responseA, responseB) <- concurrently
    (targetedUpload context session lane ipA operationId "second-epoch-a")
    (targetedUpload context session lane ipB operationId "second-epoch-b")
  let outcomes = [responseA, responseB]
      succeeded response =
        "\"result\":\"succeeded\"" `BS.isInfixOf` probeBody response
      conflicted response =
        "\"result\":\"conflict\"" `BS.isInfixOf` probeBody response
  assertEqual "new epoch winner count" 1 (length (filter succeeded outcomes))
  assertEqual "new epoch conflict count" 1 (length (filter conflicted outcomes))
  let (winningBody, replayIp) =
        if succeeded responseA
          then ("second-epoch-a", ipB)
          else ("second-epoch-b", ipA)
  secondReplay <-
    targetedUpload context session lane replayIp operationId winningBody
  requireSucceededWindow "new epoch replay" secondReplay
  downloaded <-
    targetedDownload context session lane ipA "personal" "idempotency.txt"
  assertEqual "new epoch object body" winningBody (probeBody downloaded)

verifyActiveIdentity
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> FilePath
  -> IO ()
verifyActiveIdentity context session lane ipA ipB directory = do
  let activeBody = BS.replicate 1048576 0
      activeOperation = "73000000-0000-4000-8000-000000000002"
  active <- async $
    slowUpload context session lane ipA directory activeOperation activeBody
  threadDelay 3000000
  pending <- uploadDeclared
    context
    session
    lane
    ipB
    "active-window.bin"
    activeOperation
    (BS.length activeBody)
    (sha256Hex activeBody)
    "first-epoch"
  assertContains "active replay after terminal window"
    "\"result\":\"pending\""
    (probeBody pending)
  locked <- targetedUploadAt
    context
    session
    lane
    ipB
    "active-window.bin"
    "73000000-0000-4000-8000-000000000003"
    "first-epoch"
  assertEqual "active key lock status" 200 (probeStatus locked)
  assertContains "active key lock"
    "\"result\":\"rejected\""
    (probeBody locked)
  completed <- wait active
  assertContains "active request completion"
    "\"result\":\"succeeded\""
    (probeBody completed)

targetedUpload
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> BS.ByteString
  -> IO ProbeResponse
targetedUpload context session lane ip =
  targetedUploadAt context session lane ip "idempotency.txt"

targetedUploadAt
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> String
  -> BS.ByteString
  -> IO ProbeResponse
targetedUploadAt context session lane ip path operationId body =
  uploadDeclared
    context
    session
    lane
    ip
    path
    operationId
    (BS.length body)
    (sha256Hex body)
    body

targetedDownload
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> String
  -> IO ProbeResponse
targetedDownload context session lane ip root path = do
  cookies <- cookieHeader (sessionCookies session)
  curl
    ( resolveArgs lane ip
    <> [ "--header", "Cookie: " <> B8.unpack cookies
       , probeBase context
         <> "/api/download?root="
         <> root
         <> "&path="
         <> path
       ]
    )

uploadDeclared
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> String
  -> Int
  -> BS.ByteString
  -> BS.ByteString
  -> IO ProbeResponse
uploadDeclared context session lane ip path operationId size digest body =
  withSystemTempFile "nuthatch-declared-upload" $ \bodyPath handle -> do
    BS.hPut handle body
    hClose handle
    cookies <- cookieHeader (sessionCookies session)
    curl
      ( resolveArgs lane ip
      <> [ "--header", "Cookie: " <> B8.unpack cookies
         , "--header", "Content-Type: application/octet-stream"
         , "--header", "X-Declared-Size: " <> show size
         , "--header", "X-Content-SHA256: " <> B8.unpack digest
         , "--header"
         , "X-CSRF-Token: " <> B8.unpack (sessionCsrf session)
         , "--data-binary", "@" <> bodyPath
         , probeBase context
           <> "/api/upload?root=personal&path="
           <> path
           <> "&operation_id="
           <> operationId
         ]
      )

slowUpload
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> FilePath
  -> String
  -> BS.ByteString
  -> IO ProbeResponse
slowUpload context session lane ip directory operationId body = do
  cookies <- cookieHeader (sessionCookies session)
  let outputPath = directory <> "/active-upload-response"
      statusPath = directory <> "/active-upload-status"
      firstChunk = BS.take 65536 body
      remaining = BS.drop 65536 body
      arguments =
        [ "--insecure"
        , "--silent"
        , "--show-error"
        , "--resolve", "nuthatch-" <> lane <> ":8080:" <> ip
        , "--header", "Cookie: " <> B8.unpack cookies
        , "--header", "Content-Type: application/octet-stream"
        , "--header", "X-Declared-Size: " <> show (BS.length body)
        , "--header", "X-Content-SHA256: " <> B8.unpack (sha256Hex body)
        , "--header"
        , "X-CSRF-Token: " <> B8.unpack (sessionCsrf session)
        , "--request", "POST"
        , "--upload-file", "-"
        , "--output", outputPath
        , "--write-out", "%{http_code}"
        , probeBase context
          <> "/api/upload?root=personal&path=active-window.bin"
          <> "&operation_id="
          <> operationId
        ]
  (Just input, Just statusOutput, _, processHandle) <-
    createProcess
      (proc "curl" arguments)
        { std_in = CreatePipe
        , std_out = CreatePipe
        , std_err = Inherit
        }
  BS.hPut input firstChunk
  threadDelay 5000000
  BS.hPut input remaining
  hClose input
  statusBytes <- BS.hGetContents statusOutput
  exitCode <- waitForProcess processHandle
  case exitCode of
    ExitSuccess -> pure ()
    ExitFailure code -> fail ("slow upload curl failed: " <> show code)
  BS.writeFile statusPath statusBytes
  responseBody <- BS.readFile outputPath
  status <- case reads (B8.unpack statusBytes) of
    [(value, "")] -> pure value
    _ -> fail "slow upload returned invalid HTTP status"
  pure ProbeResponse
    { probeStatus = status
    , probeHeaders = []
    , probeBody = responseBody
    }

resolveArgs :: String -> String -> [String]
resolveArgs lane ip =
  [ "--connect-timeout", "2"
  , "--retry", "10"
  , "--retry-connrefused"
  , "--retry-delay", "1"
  , "--resolve", "nuthatch-" <> lane <> ":8080:" <> ip
  ]

requireSucceededWindow :: String -> ProbeResponse -> IO ()
requireSucceededWindow label response = do
  assertEqual (label <> " status") 200 (probeStatus response)
  assertContains label "\"result\":\"succeeded\"" (probeBody response)
  assertContains (label <> " window")
    "\"idempotency_window_seconds\":2"
    (probeBody response)

emitRefinementTrace :: IO ()
emitRefinementTrace =
  mapM_ putStrLn
    [ "{\"epoch\":1,\"action\":\"Init\"}"
    , "{\"epoch\":1,\"action\":\"FirstRequest\",\"observation\":{\"result\":\"pending\",\"disposition\":\"first\"}}"
    , "{\"epoch\":1,\"action\":\"PrepareEffect\"}"
    , "{\"epoch\":1,\"action\":\"DispatchEffect\"}"
    , "{\"epoch\":1,\"action\":\"ExternalEffectApplied\"}"
    , "{\"epoch\":1,\"action\":\"FinalizeSuccess\"}"
    , "{\"epoch\":1,\"action\":\"ReplayRequest\",\"observation\":{\"result\":\"succeeded\",\"disposition\":\"replay\"}}"
    , "{\"epoch\":1,\"action\":\"ConflictRequest\",\"observation\":{\"result\":\"conflict\",\"disposition\":\"conflict\"}}"
    , "{\"after_epoch\":1,\"action\":\"TraceCut\",\"reason\":\"terminal_identity_expired\"}"
    , "{\"epoch\":2,\"action\":\"Init\"}"
    , "{\"epoch\":2,\"action\":\"FirstRequest\",\"observation\":{\"result\":\"pending\",\"disposition\":\"first\"}}"
    , "{\"epoch\":2,\"action\":\"PrepareEffect\"}"
    , "{\"epoch\":2,\"action\":\"DispatchEffect\"}"
    , "{\"epoch\":2,\"action\":\"ExternalEffectApplied\"}"
    , "{\"epoch\":2,\"action\":\"FinalizeSuccess\"}"
    , "{\"epoch\":2,\"action\":\"ReplayRequest\",\"observation\":{\"result\":\"succeeded\",\"disposition\":\"replay\"}}"
    ]