module Main (main) where

import           Control.Concurrent (threadDelay)
import           Control.Concurrent.Async (concurrently)
import           Control.Exception (SomeException, try)
import           Control.Monad (replicateM_, unless)
import           Crypto.Hash
  ( Digest
  , SHA256
  , hashFinalize
  , hashInit
  , hashUpdate
  )
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base16 as B16
import qualified Data.ByteString.Char8 as B8
import           System.Directory (getFileSize)
import           System.Environment (getArgs)
import           System.Exit (ExitCode (..))
import           System.IO
  ( IOMode (..)
  , hClose
  , withBinaryFile
  )
import           System.IO.Temp
  ( withSystemTempDirectory
  , withSystemTempFile
  )
import           System.Process (readProcessWithExitCode)
import           Text.Read (readMaybe)
import           Nuthatch.Probe.Assert
import           Nuthatch.Probe.Curl
import           Nuthatch.Probe.Oidc
import           Nuthatch.Probe.ProductClient

main :: IO ()
main =
  withSystemTempDirectory "nuthatch-streaming-probe" $ \directory -> do
    context <- newProbeContext directory
    user <- requireEnv "NUTHATCH_OIDC_USER_A"
    session <- login context False user "streaming-a"
    getArgs >>= \case
      ["large"] -> verifyLarge context session directory
      ["throttle"] -> verifyThrottle context session directory
      ["refill"] -> verifyRefill context session directory
      ["unavailable"] -> verifyUnavailable context session
      _ -> fail "expected large, throttle, refill, or unavailable mode"

verifyLarge :: ProbeContext -> Session -> FilePath -> IO ()
verifyLarge context session directory = do
  largeBytes <- readEnvInteger "NUTHATCH_STREAMING_LARGE_BYTES"
  overLimit <- readEnvInteger "NUTHATCH_STREAMING_OVER_LIMIT_BYTES"
  let largePath = directory <> "/large.bin"
      downloadPath = directory <> "/large-download.bin"
  writeZeros largePath largeBytes
  largeDigest <- hashFile largePath
  uploaded <- uploadFileDeclared
    context session "streaming-large.bin"
    "70000000-0000-4000-8000-000000000001"
    largeBytes largeDigest largePath []
  requireSucceeded "large streaming upload" uploaded
  downloaded <- curlToFile downloadPath
    [ "--cookie", sessionCookies session
    , probeBase context
      <> "/api/download?root=personal&path=streaming-large.bin"
    ]
  assertEqual "large streaming download status" 200 (curlStatus downloaded)
  hashFile downloadPath >>= assertEqual "large download digest" largeDigest
  let tiny = "x"
  mapM_
    (\label -> do
      replay <- uploadBytesDeclared
        context session "streaming-large.bin"
        "70000000-0000-4000-8000-000000000007"
        largeBytes largeDigest tiny []
      requireSucceeded ("prior/desired observation " <> label) replay)
    ["first", "replay"]
  sizeMismatch <- uploadBytesDeclared
    context session "size-mismatch.bin"
    "70000000-0000-4000-8000-000000000002"
    2 (sha256Hex tiny) tiny []
  requireError "size mismatch" 400 "size_mismatch" sizeMismatch
  identityMismatch <- uploadBytesDeclared
    context session "identity-mismatch.bin"
    "70000000-0000-4000-8000-000000000003"
    1 (BS.replicate 64 48) tiny []
  requireError "content identity mismatch"
    400 "content_identity_mismatch" identityMismatch
  tooLarge <- uploadFileDeclared
    context session "too-large.bin"
    "70000000-0000-4000-8000-000000000004"
    overLimit (sha256Hex tiny) largePath
    ["--max-time", "5", "--limit-rate", "64k"]
  requireError "object size limit" 413 "object_too_large" tooLarge
  aborted <- interruptedUpload
    context session "aborted-upload.bin"
    "70000000-0000-4000-8000-000000000005"
    largeBytes largeDigest largePath
  assert "aborted upload unexpectedly completed" (not aborted)
  missing <- download context session "personal" "aborted-upload.bin"
  assertEqual "aborted upload visibility" 404 (curlStatus missing)
  replayAfterAbort <- uploadBytesDeclared
    context session "aborted-upload.bin"
    "70000000-0000-4000-8000-000000000005"
    largeBytes largeDigest tiny []
  assertEqual "aborted upload replay status" 200 (curlStatus replayAfterAbort)
  assert "aborted upload replay lost active or failed identity" $
    any (`BS.isInfixOf` curlBody replayAfterAbort)
      ["\"result\":\"pending\"", "\"result\":\"failed\""]
  abortedOverwrite <- interruptedUpload
    context session "streaming-large.bin"
    "70000000-0000-4000-8000-000000000006"
    largeBytes (BS.replicate 64 48) largePath
  assert "aborted overwrite unexpectedly completed" (not abortedOverwrite)
  preservedPath <- pure (directory <> "/preserved.bin")
  preserved <- curlToFile preservedPath
    [ "--cookie", sessionCookies session
    , probeBase context
      <> "/api/download?root=personal&path=streaming-large.bin"
    ]
  assertEqual "preserved object status" 200 (curlStatus preserved)
  hashFile preservedPath >>= assertEqual "preserved object digest" largeDigest
  let partialPath = directory <> "/partial-download.bin"
  partial <- interruptedDownload context session partialPath
  assert "aborted download unexpectedly completed" (not partial)
  partialBytes <- getFileSize partialPath
  assert "aborted download emitted no bytes" (partialBytes > 0)
  assert "aborted download emitted the complete object" $
    partialBytes < largeBytes
  ready <- curl [probeBase context <> "/readyz"]
  requireStatus "readiness after aborts" 200 ready
  assertContains "readiness after aborts" "\"ready\":true" (curlBody ready)
  putStrLn "Large streaming and abort behavior passed"

verifyThrottle :: ProbeContext -> Session -> FilePath -> IO ()
verifyThrottle context session directory = do
  userB <- requireEnv "NUTHATCH_OIDC_USER_B"
  sessionB <- login context False userB "streaming-b"
  let jointPath = directory <> "/joint-upload.bin"
      throttlePath = directory <> "/throttle-upload.bin"
  writeZeros jointPath 65536
  jointDigest <- hashFile jointPath
  joint <- uploadFileDeclared
    context session "joint-throttle.bin"
    "70000000-0000-4000-8000-000000000013"
    65536 jointDigest jointPath []
  requireSucceeded "joint throttle setup" joint
  writeZeros throttlePath 1048576
  uploadDigest <- hashFile throttlePath
  throttled <- uploadFileDeclared
    context sessionB "throttled-upload.bin"
    "71000000-0000-4000-8000-000000000001"
    1048576 uploadDigest throttlePath []
  requireError "incremental upload throttle"
    429 "bandwidth_exceeded" throttled
  prestart <- uploadBytesDeclared
    context sessionB "prestart.bin"
    "71000000-0000-4000-8000-000000000002"
    1 uploadDigest "x" []
  assertEqual "pre-start upload throttle" 429 (curlStatus prestart)
  let firstPath = directory <> "/parallel-a"
      secondPath = directory <> "/parallel-b"
      downloadTo path =
        try (curlToFile path
          [ "--cookie", sessionCookies session
          , probeBase context
            <> "/api/download?root=personal&path=streaming-large.bin"
          ]) :: IO (Either SomeException CurlResponse)
  BS.writeFile firstPath BS.empty
  BS.writeFile secondPath BS.empty
  (first, second) <- concurrently (downloadTo firstPath) (downloadTo secondPath)
  assert "first throttled download completed" (isLeft first)
  assert "second throttled download completed" (isLeft second)
  firstBytes <- getFileSize firstPath
  secondBytes <- getFileSize secondPath
  burst <- readEnvInteger "NUTHATCH_STREAMING_THROTTLE_BURST_BYTES"
  let moved = firstBytes + secondBytes
      jointLimit = burst - 65536 + 4096
  assert "parallel throttled downloads emitted no bytes" (moved > 0)
  assert "parallel downloads exceeded their joint token budget" $
    moved <= jointLimit
  exhausted <- download context session "personal" "streaming-large.bin"
  requireStatus "post-exhaustion download" 429 exhausted
  putStrLn "Shared incremental throttle behavior passed"

verifyRefill :: ProbeContext -> Session -> FilePath -> IO ()
verifyRefill context session directory = do
  let bodyPath = directory <> "/refill-upload.bin"
  writeZeros bodyPath 65536
  digest <- hashFile bodyPath
  initial <- uploadFileDeclared
    context session "refill-initial.bin"
    "72000000-0000-4000-8000-000000000001"
    65536 digest bodyPath []
  requireSucceeded "initial token use" initial
  exhausted <- uploadFileDeclared
    context session "refill-exhausted.bin"
    "72000000-0000-4000-8000-000000000003"
    65536 digest bodyPath []
  requireStatus "exhausted token bucket" 429 exhausted
  threadDelay 2000000
  refilled <- uploadFileDeclared
    context session "refill-refilled.bin"
    "72000000-0000-4000-8000-000000000002"
    65536 digest bodyPath []
  requireSucceeded "refilled token bucket" refilled
  putStrLn "Token-bucket refill behavior passed"

verifyUnavailable :: ProbeContext -> Session -> IO ()
verifyUnavailable context session = do
  listing <- browse context session "?root=personal"
  requireStatus "browse without Redis" 200 listing
  assertContains "browse without Redis" "\"path\":" (curlBody listing)
  let path =
        if "cross-user.txt" `BS.isInfixOf` curlBody listing
          then "cross-user.txt"
          else "streaming-large.bin"
      query =
        if path == "cross-user.txt" then "CROSS-USER" else "STREAMING-LARGE"
  searched <- browse context session ("?root=personal&q=" <> query)
  assertContains "search without Redis" (B8.pack path) (curlBody searched)
  unavailableUpload <- upload context session "personal" "unavailable.bin"
    "72000000-0000-4000-8000-000000000001" "x"
  requireError "upload without Redis"
    503 "throttle_unavailable" unavailableUpload
  unavailableDownload <-
    retryTransport 4 (download context session "personal" path)
  requireError "download without Redis"
    503 "throttle_unavailable" unavailableDownload
  putStrLn "Redis-unavailable transfer behavior passed"

retryTransport :: Int -> IO value -> IO value
retryTransport attempts action = do
  result <- try action
  case result of
    Right value -> pure value
    Left exception
      | attempts > 1 -> do
          threadDelay 100000
          retryTransport (attempts - 1) action
      | otherwise -> fail (show (exception :: SomeException))

uploadFileDeclared
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> Integer
  -> BS.ByteString
  -> FilePath
  -> [String]
  -> IO CurlResponse
uploadFileDeclared context session path operationId size digest bodyPath extra =
  curl
    ( extra
    <> uploadArguments context session path operationId size digest
    <> ["--data-binary", "@" <> bodyPath]
    )

uploadBytesDeclared
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> Integer
  -> BS.ByteString
  -> BS.ByteString
  -> [String]
  -> IO CurlResponse
uploadBytesDeclared context session path operationId size digest body extra =
  withSystemTempFile "nuthatch-streaming-body" $ \bodyPath handle -> do
    BS.hPut handle body
    hClose handle
    uploadFileDeclared
      context session path operationId size digest bodyPath extra

uploadArguments
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> Integer
  -> BS.ByteString
  -> [String]
uploadArguments context session path operationId size digest =
  [ "--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)
  , probeBase context
    <> "/api/upload?root=personal&path="
    <> path
    <> "&operation_id="
    <> operationId
  ]

interruptedUpload
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> Integer
  -> BS.ByteString
  -> FilePath
  -> IO Bool
interruptedUpload context session path operationId size digest bodyPath = do
  let arguments =
        [ "--insecure"
        , "--silent"
        , "--show-error"
        , "--max-time", "1"
        , "--limit-rate", "64k"
        , "--request", "POST"
        , "--upload-file", bodyPath
        ]
        <> uploadArguments context session path operationId size digest
  (exitCode, _, _) <- readProcessWithExitCode "curl" arguments ""
  pure (exitCode == ExitSuccess)

interruptedDownload
  :: ProbeContext
  -> Session
  -> FilePath
  -> IO Bool
interruptedDownload context session outputPath = do
  let arguments =
        [ "--insecure"
        , "--silent"
        , "--show-error"
        , "--max-time", "1"
        , "--limit-rate", "64k"
        , "--cookie", sessionCookies session
        , "--output", outputPath
        , probeBase context
          <> "/api/download?root=personal&path=streaming-large.bin"
        ]
  (exitCode, _, _) <- readProcessWithExitCode "curl" arguments ""
  pure (exitCode == ExitSuccess)

writeZeros :: FilePath -> Integer -> IO ()
writeZeros path byteCount =
  withBinaryFile path WriteMode $ \handle -> do
    let chunk = BS.replicate 1048576 0
        (fullChunks, remainder) = byteCount `divMod` 1048576
    replicateM_ (fromIntegral fullChunks) (BS.hPut handle chunk)
    unless (remainder == 0) $
      BS.hPut handle (BS.take (fromIntegral remainder) chunk)

hashFile :: FilePath -> IO BS.ByteString
hashFile path =
  withBinaryFile path ReadMode $ \handle ->
    go handle hashInit
  where
    go handle context = do
      chunk <- BS.hGetSome handle 1048576
      if BS.null chunk
        then
          pure
            (B16.encode
              (BA.convert (hashFinalize context :: Digest SHA256)))
        else do
          let updated = hashUpdate context chunk
          updated `seq` go handle updated

readEnvInteger :: String -> IO Integer
readEnvInteger name = do
  value <- requireEnv name
  maybe (fail ("invalid " <> name)) pure (readMaybe value)

requireSucceeded :: String -> CurlResponse -> IO ()
requireSucceeded label response = do
  requireStatus label 200 response
  assertContains label "\"result\":\"succeeded\"" (curlBody response)

requireError :: String -> Int -> BS.ByteString -> CurlResponse -> IO ()
requireError label status code response = do
  requireStatus label status response
  assertContains label ("\"error\":\"" <> code <> "\"") (curlBody response)

requireStatus :: String -> Int -> CurlResponse -> IO ()
requireStatus label expected response =
  assertEqual (label <> " status") expected (curlStatus response)

isLeft :: Either left right -> Bool
isLeft (Left _) = True
isLeft (Right _) = False