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.Http
import           Nuthatch.Probe.Oidc
import           Nuthatch.Probe.ProductClient

main :: IO ()
main =
  withSystemTempDirectory "nuthatch-streaming-probe" $ \directory -> do
    context <- newProbeContext
    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 defaultRequestOptions
  requireSucceeded "large streaming upload" uploaded
  downloaded <- downloadTo
    (probeClient context)
    defaultRequestOptions
      { optionCookies = Just (sessionCookies session) }
    downloadPath
    (probeBase context
      <> "/api/download?root=personal&path=streaming-large.bin")
  assertEqual "large streaming download status" 200 (probeStatus 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 defaultRequestOptions
      requireSucceeded ("prior/desired observation " <> label) replay)
    ["first", "replay"]
  sizeMismatch <- uploadBytesDeclared
    context session "size-mismatch.bin"
    "70000000-0000-4000-8000-000000000002"
    2 (sha256Hex tiny) tiny defaultRequestOptions
  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 defaultRequestOptions
  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
    defaultRequestOptions { optionTimeoutMicros = Just 5000000 }
  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 (probeStatus missing)
  replayAfterAbort <- uploadBytesDeclared
    context session "aborted-upload.bin"
    "70000000-0000-4000-8000-000000000005"
    largeBytes largeDigest tiny defaultRequestOptions
  assertEqual "aborted upload replay status" 200 (probeStatus replayAfterAbort)
  assert "aborted upload replay lost active or failed identity" $
    any (`BS.isInfixOf` probeBody 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 <- downloadTo
    (probeClient context)
    defaultRequestOptions
      { optionCookies = Just (sessionCookies session) }
    preservedPath
    (probeBase context
      <> "/api/download?root=personal&path=streaming-large.bin")
  assertEqual "preserved object status" 200 (probeStatus 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 <- get (probeClient context) (probeBase context <> "/readyz")
  requireStatus "readiness after aborts" 200 ready
  assertContains "readiness after aborts" "\"ready\":true" (probeBody 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 <- curlUploadFileDeclared
    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 <- curlUploadFileDeclared
    context sessionB "throttled-upload.bin"
    "71000000-0000-4000-8000-000000000001"
    1048576 uploadDigest throttlePath
  requireError "incremental upload throttle"
    429 "bandwidth_exceeded" throttled
  prestart <- curlUploadBytesDeclared
    context sessionB "prestart.bin"
    "71000000-0000-4000-8000-000000000002"
    65536 uploadDigest "x"
  assertEqual
    ("pre-start upload throttle: " <> B8.unpack (probeBody prestart))
    429
    (probeStatus prestart)
  let firstPath = directory <> "/parallel-a"
      secondPath = directory <> "/parallel-b"
      downloadAttempt path = do
        cookies <- cookieHeader (sessionCookies session)
        try
          (curlToFile path
            [ "--header", "Cookie: " <> B8.unpack cookies
            , probeBase context
              <> "/api/download?root=personal&path=streaming-large.bin"
            ])
          :: IO (Either SomeException ProbeResponse)
  BS.writeFile firstPath BS.empty
  BS.writeFile secondPath BS.empty
  (first, second) <-
    concurrently (downloadAttempt firstPath) (downloadAttempt 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"
      refillPath = directory <> "/refill-byte.bin"
  writeZeros bodyPath 65536
  BS.writeFile refillPath "x"
  digest <- hashFile bodyPath
  refillDigest <- hashFile refillPath
  initial <- curlUploadFileDeclared
    context session "refill-initial.bin"
    "72000000-0000-4000-8000-000000000001"
    65536 digest bodyPath
  requireSucceeded "initial token use" initial
  exhausted <- curlUploadFileDeclared
    context session "refill-exhausted.bin"
    "72000000-0000-4000-8000-000000000003"
    65536 digest bodyPath
  requireStatus "exhausted token bucket" 429 exhausted
  threadDelay 2000000
  refilled <- curlUploadFileDeclared
    context session "refill-refilled.bin"
    "72000000-0000-4000-8000-000000000002"
    1 refillDigest refillPath
  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\":" (probeBody listing)
  let path =
        if "cross-user.txt" `BS.isInfixOf` probeBody 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) (probeBody 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
  -> RequestOptions
  -> IO ProbeResponse
uploadFileDeclared context session path operationId size digest bodyPath options =
  postWith
    (probeClient context)
    options
      { optionCookies = Just (sessionCookies session)
      , optionHeaders =
          [ ("Content-Type", "application/octet-stream")
          , ("X-Declared-Size", B8.pack (show size))
          , ("X-Content-SHA256", digest)
          , ("X-CSRF-Token", sessionCsrf session)
          ] <> optionHeaders options
      , optionBody = FileBody bodyPath
      }
    (uploadUrl context path operationId)

uploadBytesDeclared
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> Integer
  -> BS.ByteString
  -> BS.ByteString
  -> RequestOptions
  -> IO ProbeResponse
uploadBytesDeclared context session path operationId size digest body options =
  postWith
    (probeClient context)
    options
      { optionCookies = Just (sessionCookies session)
      , optionHeaders =
          [ ("Content-Type", "application/octet-stream")
          , ("X-Declared-Size", B8.pack (show size))
          , ("X-Content-SHA256", digest)
          , ("X-CSRF-Token", sessionCsrf session)
          ] <> optionHeaders options
      , optionBody = ByteBody body
      }
    (uploadUrl context path operationId)

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

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

curlUploadBytesDeclared
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> Integer
  -> BS.ByteString
  -> BS.ByteString
  -> IO ProbeResponse
curlUploadBytesDeclared context session path operationId size digest body =
  withSystemTempFile "nuthatch-streaming-body" $ \bodyPath handle -> do
    BS.hPut handle body
    hClose handle
    curlUploadFileDeclared
      context session path operationId size digest bodyPath

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

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

uploadUrl :: ProbeContext -> String -> String -> String
uploadUrl context path operationId =
  probeBase context
    <> "/api/upload?root=personal&path="
    <> path
    <> "&operation_id="
    <> operationId

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 -> ProbeResponse -> IO ()
requireSucceeded label response = do
  requireStatus label 200 response
  assertContains label "\"result\":\"succeeded\"" (probeBody response)

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

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

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