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.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 directory
    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\""
    (curlBody 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` curlBody response
      conflicted response =
        "\"result\":\"conflict\"" `BS.isInfixOf` curlBody 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 <- downloadFrom
    (probeBase context)
    (resolveArgs lane ipA)
    session
    "personal"
    "idempotency.txt"
  assertEqual "new epoch object body" winningBody (curlBody 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\""
    (curlBody pending)
  locked <- targetedUploadAt
    context
    session
    lane
    ipB
    "active-window.bin"
    "73000000-0000-4000-8000-000000000003"
    "first-epoch"
  assertEqual "active key lock status" 200 (curlStatus locked)
  assertContains "active key lock"
    "\"result\":\"rejected\""
    (curlBody locked)
  completed <- wait active
  assertContains "active request completion"
    "\"result\":\"succeeded\""
    (curlBody completed)

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

targetedUploadAt
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> String
  -> BS.ByteString
  -> IO CurlResponse
targetedUploadAt context session lane ip path operationId =
  uploadTo
    (probeBase context)
    (resolveArgs lane ip)
    session
    "personal"
    path
    operationId

uploadDeclared
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> String
  -> String
  -> Int
  -> BS.ByteString
  -> BS.ByteString
  -> IO CurlResponse
uploadDeclared context session lane ip path operationId size digest body =
  withSystemTempFile "nuthatch-declared-upload" $ \bodyPath handle -> do
    BS.hPut handle body
    hClose handle
    curl
      ( resolveArgs lane ip
      <> [ "--cookie", sessionCookies session
         , "--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 CurlResponse
slowUpload context session lane ip directory operationId body = do
  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
        , "--cookie", sessionCookies session
        , "--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 CurlResponse
    { curlStatus = status
    , curlHeaders = []
    , curlBody = 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 -> CurlResponse -> IO ()
requireSucceededWindow label response = do
  assertEqual (label <> " status") 200 (curlStatus response)
  assertContains label "\"result\":\"succeeded\"" (curlBody response)
  assertContains (label <> " window")
    "\"idempotency_window_seconds\":2"
    (curlBody 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\"}}"
    ]