module Main (main) where

import           Control.Concurrent.Async (concurrently, mapConcurrently)
import           Control.Exception (SomeException, try)
import           Control.Monad (unless)
import           Crypto.Hash (Digest, SHA256, hashlazy)
import qualified Data.ByteArray as BA
import qualified Data.ByteString.Base16 as B16
import qualified Data.ByteString.Lazy as BL
import           Data.List (sort)
import           Data.Text (Text)
import qualified Data.Text.Encoding as TE
import           Database.PostgreSQL.Simple (Only (..))
import           Nuthatch.Config
  ( Config (..)
  , loadConfig
  , revealSecret
  , secret
  )
import           Nuthatch.Db
  ( Db
  , configurePoolLimit
  , ensurePersonalPool
  , migrate
  , withDb
  , withTransaction
  )
import qualified Nuthatch.Db as Db
import           Nuthatch.Protocol
import           Nuthatch.S3
import           Nuthatch.Types

main :: IO ()
main = do
  config <- loadConfig
  withDb config $ \database -> do
    migrate database "sql" config
    migrate database "sql" config
    verifySchema database
    verifyAdmission database config
    verifyRowDecoder database
  store <- newS3 config
  verifyStorage config store
  putStrLn "Storage foundation integration gate passed"

verifySchema :: Db -> IO ()
verifySchema database = do
  tables <- Db.query_ database
    "SELECT table_name FROM information_schema.tables \
    \WHERE table_schema = 'public' \
    \AND table_name IN ('pools','mutations','key_locks','objects_meta','schema_migrations') \
    \ORDER BY table_name"
    :: IO [Only Text]
  assertEqual "authoritative schema tables"
    ["key_locks", "mutations", "objects_meta", "pools", "schema_migrations"]
    (map fromOnly tables)
  constraints <- Db.query_ database
    "SELECT count(*) FROM pg_constraint \
    \WHERE conname = 'pools_live_bytes_within_limit' \
    \AND conrelid = 'pools'::regclass"
    :: IO [Only Int]
  assertEqual "pool limit constraint" [Only 1] constraints
  migrations <- Db.query_ database
    "SELECT filename FROM schema_migrations ORDER BY filename"
    :: IO [Only Text]
  assertEqual "migration history"
    ["001_initial_schema.sql"]
    (map fromOnly migrations)

verifyAdmission :: Db -> Config -> IO ()
verifyAdmission database config = do
  resetAuthority database 100
  replayOutcomes <- mapConcurrently
    (const (admit database "same-operation" "same-key" 10))
    [1 .. 8 :: Int]
  assertEqual "one first request" 1 (countPrepared replayOutcomes)
  assertEqual "stable concurrent replay" 7 (countReplay replayOutcomes)
  assertAuthority database 1 10

  resetAuthority database 100
  conflictOutcomes <- mapConcurrently
    (\size -> admit database "conflicting-operation" "conflicting-key" size)
    [10, 11]
  assertEqual "one conflicting first request" 1 (countPrepared conflictOutcomes)
  assertEqual "conflicting reuse" 1 (countConflict conflictOutcomes)
  assertAuthority database 1 (mrGrowth (preparedRecord conflictOutcomes))

  resetAuthority database 100
  keyOutcomes <- mapConcurrently
    (\operation -> admit database operation "locked-key" 10)
    ["key-operation-a", "key-operation-b"]
  assertEqual "one same-key authority" 1 (countPrepared keyOutcomes)
  assertEqual "one same-key rejection" 1 (countRejected keyOutcomes)
  assertAuthority database 1 10

  resetAuthority database 10
  quotaOutcomes <- mapConcurrently
    (\(operation, key) -> admit database operation key 8)
    [("quota-operation-a", "quota-key-a"), ("quota-operation-b", "quota-key-b")]
  assertEqual "one quota reservation" 1 (countPrepared quotaOutcomes)
  assertEqual "one quota rejection" 1 (countRejected quotaOutcomes)
  assertAuthority database 1 8
  lower <- try (configurePoolLimit database sharedPool 7)
    :: IO (Either SomeException ())
  assert "unsafe pool limit reduction rejected" (isLeft lower)
  configurePoolLimit database sharedPool 20
  pool <- Db.query database
    "SELECT limit_bytes, used_bytes, reserved_bytes FROM pools WHERE pool_id = ?"
    (Only (unPoolId sharedPool))
    :: IO [(Integer, Integer, Integer)]
  assertEqual "safe pool limit reconfiguration" [(20, 0, 8)] pool

  let personalUser = UserId "personal-limit-user"
      personal = personalPool personalUser
  ensurePersonalPool database
    config { cfgPersonalQuotaBytes = 30 }
    personalUser
  ensurePersonalPool database
    config { cfgPersonalQuotaBytes = 40 }
    personalUser
  _ <- Db.execute database
    "UPDATE pools SET reserved_bytes = 35 WHERE pool_id = ?"
    (Only (unPoolId personal))
  personalLower <- try
    (ensurePersonalPool database
      config { cfgPersonalQuotaBytes = 34 }
      personalUser)
    :: IO (Either SomeException ())
  assert "unsafe personal limit reduction rejected" (isLeft personalLower)
  personalPoolRow <- Db.query database
    "SELECT limit_bytes, used_bytes, reserved_bytes FROM pools WHERE pool_id = ?"
    (Only (unPoolId personal))
    :: IO [(Integer, Integer, Integer)]
  assertEqual "personal limit reconfiguration" [(40, 0, 35)] personalPoolRow

  resetAuthority database 100
  (racingAdmission, racingLimit) <- concurrently
    (admit database "limit-race-operation" "limit-race-key" 80)
    (try (configurePoolLimit database sharedPool 50)
      :: IO (Either SomeException ()))
  racingPool <- Db.query database
    "SELECT limit_bytes, used_bytes, reserved_bytes FROM pools WHERE pool_id = ?"
    (Only (unPoolId sharedPool))
    :: IO [(Integer, Integer, Integer)]
  case (racingAdmission, racingLimit, racingPool) of
    (AdmitPrepared _, Left _, [(100, 0, 80)]) ->
      assertAuthority database 1 80
    (AdmitRejected _ _, Right (), [(50, 0, 0)]) ->
      assertAuthority database 0 0
    other -> fail ("unsafe admission/limit race result: " <> show other)

  resetAuthority database 100
  observationOutcome <- admit database
    "observation-operation"
    "observation-key"
    10
  let observationRecord = preparedRecord [observationOutcome]
      thirdState =
        ObjectPresent 10
          (Just (ContentSha256
            "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"))
  settled <- withTransaction database $ \connection -> do
    applying <- markApplying connection observationRecord
    pending <- markExternalPending
      connection observationRecord (DispatchTiming 600 90)
    externalSettled <- markExternalSettled connection observationRecord
    pure (applying && pending && externalSettled)
  assert "observation fixture reached applying/settled" settled
  unavailable <- withTransaction database $ \connection ->
    finalizeFromObservation connection observationRecord
      (ObjectUnavailable "timeout")
  ambiguous <- withTransaction database $ \connection ->
    finalizeFromObservation connection observationRecord thirdState
  assertEqual "unavailable observation remains non-terminal" Nothing unavailable
  assertEqual "third-state observation remains non-terminal" Nothing ambiguous
  assertAuthority database 1 10
  durableState <- Db.query database
    "SELECT state, external_state FROM mutations \
    \WHERE user_id = ? AND kind = ? AND operation_id = ?"
    ( "integration-user" :: Text
    , "put" :: Text
    , "observation-operation" :: Text
    )
    :: IO [(Text, Text)]
  assertEqual "ambiguous observation retains mutation authority"
    [("applying", "settled")]
    durableState

  verifyIdempotencyExpiry database

verifyIdempotencyExpiry :: Db -> IO ()
verifyIdempotencyExpiry database = do
  resetAuthority database 100
  first <- admitWithWindow database "expiring-operation" "expiring-key" 10 1
  let rec = preparedRecord [first]
  finalized <- withTransaction database $ \connection -> do
    applying <- markApplying connection rec
    pending <- markExternalPending connection rec (DispatchTiming 600 90)
    settled <- markExternalSettled connection rec
    succeeded <- finalizeSuccess connection rec 10 (Just fixedIdentity)
    pure (applying && pending && settled && succeeded)
  assert "terminal expiry fixture finalized" finalized
  _ <- Db.execute database
    "UPDATE mutations SET updated_at = now() - interval '5 seconds' \
    \WHERE user_id = ? AND kind = ? AND operation_id = ?"
    ( "integration-user" :: Text
    , "put" :: Text
    , "expiring-operation" :: Text
    )
  outcomes <- mapConcurrently
    (\size ->
      admitWithWindow database "expiring-operation" "expiring-key" size 1)
    [11, 12]
  assertEqual "one first request after terminal expiry" 1 (countPrepared outcomes)
  assertEqual "atomic conflicting reuse after terminal expiry" 1
    (countConflict outcomes)
  assertAuthority database 1 (mrGrowth (preparedRecord outcomes))

  resetAuthority database 100
  active <- admitWithWindow database "active-operation" "active-key" 10 1
  _ <- Db.execute database
    "UPDATE mutations SET updated_at = now() - interval '5 seconds' \
    \WHERE user_id = ? AND kind = ? AND operation_id = ?"
    ( "integration-user" :: Text
    , "put" :: Text
    , "active-operation" :: Text
    )
  replay <- admitWithWindow database "active-operation" "active-key" 10 1
  blocked <- admitWithWindow database "other-operation" "active-key" 10 1
  assertEqual "aged active identity remains replayable" 1
    (countReplay [active, replay])
  assertEqual "aged active key remains locked" 1 (countRejected [blocked])
  assertAuthority database 1 10

verifyRowDecoder :: Db -> IO ()
verifyRowDecoder database = do
  resetAuthority database 100
  outcome <- admit database "invalid-phase-operation" "invalid-phase-key" 10
  let record = preparedRecord [outcome]
  roundTripped <- withTransaction database $ \connection ->
    lookupMutation
      connection
      (mrUserId record)
      (mrKind record)
      (mrOperationId record)
  assertEqual "typed mutation row round-trip" (Just record) roundTripped
  _ <- Db.execute database
    "UPDATE mutations SET external_state = 'pending', \
    \dispatch_started_at = now(), request_deadline_at = now(), \
    \worker_lease_until = now() \
    \WHERE user_id = ? AND kind = ? AND operation_id = ?"
    ( unUserId (mrUserId record)
    , kindLabel (mrKind record)
    , unOperationId (mrOperationId record)
    )
  invalidDecode <-
    ( try $ withTransaction database $ \connection ->
        lookupMutation
          connection
          (mrUserId record)
          (mrKind record)
          (mrOperationId record)
    ) :: IO (Either SomeException (Maybe MutationRecord))
  case invalidDecode of
    Left _ -> pure ()
    Right value ->
      fail ("invalid persisted mutation phase decoded as " <> show value)

verifyStorage :: Config -> S3 -> IO ()
verifyStorage config store = do
  let fallbackKey = ObjectKey "shared/fallback.bin"
      fallbackBody = BL.replicate (5 * 1024 * 1024) 0
      fallbackIdentity =
        ContentSha256
          "c036cbb7553a909f8b8877d4461924307f27ecb66cff928eeeafd569c3887e29"
  expectRight "put object without identity metadata"
    =<< putObject store fallbackKey
      (toInteger (BL.length fallbackBody))
      Nothing
      fallbackBody
  fallbackHead <- expectRight "observe object without identity metadata"
    =<< observeObjectFully store fallbackKey
  assertEqual "streaming digest fallback size"
    (toInteger (BL.length fallbackBody))
    (oiSize fallbackHead)
  assertEqual "streaming digest fallback identity"
    (Just fallbackIdentity)
    (oiSha256 fallbackHead)
  expectRight "delete fallback object" =<< deleteObject store fallbackKey

  let unicodeKey = ObjectKey "shared/문서/βeta.txt"
      firstBody = "same-size-a"
      secondBody = "same-size-b"
      firstIdentity = contentIdentity firstBody
      secondIdentity = contentIdentity secondBody
  expectRight "put Unicode key"
    =<< putObject store unicodeKey (toInteger (BL.length firstBody))
      (Just firstIdentity) firstBody
  fetched <- expectRight "get Unicode key" =<< getObjectBytes store unicodeKey
  assertEqual "Unicode key round trip" firstBody fetched
  firstHead <- expectJust "head Unicode key" =<< headObject store unicodeKey
  assertEqual "head size" (toInteger (BL.length firstBody)) (oiSize firstHead)
  assertEqual "head metadata" (Just firstIdentity) (oiSha256 firstHead)
  expectRight "overwrite same-size content"
    =<< putObject store unicodeKey (toInteger (BL.length secondBody))
      (Just secondIdentity) secondBody
  secondHead <- expectJust "head overwritten key" =<< headObject store unicodeKey
  assert "same-size content identity differs"
    (oiSha256 firstHead /= oiSha256 secondHead)

  let paged =
        [ (ObjectKey "shared/page/α", "a")
        , (ObjectKey "shared/page/β", "b")
        , (ObjectKey "shared/page/γ", "c")
        ]
  mapM_
    (\(key, body) ->
      expectRight "put paged key"
        =<< putObject store key (toInteger (BL.length body))
          (Just (contentIdentity body)) body)
    paged
  keys <- collectPages store "shared/page/" Nothing []
  assertEqual "paginated Unicode list"
    (sort [unObjectKey key | (key, _) <- paged])
    (sort keys)

  missing <- getObjectBytes store (ObjectKey "shared/missing")
  assertEqual "typed not-found error" (Left S3NotFound) missing
  deniedStore <- newS3 config
    { cfgS3SecretKey =
        secret (revealSecret (cfgS3SecretKey config) <> "-invalid")
    }
  denied <- listObjects deniedStore "" Nothing Nothing 1
  case denied of
    Left (S3ServiceError code) ->
      assert "typed provider rejection" (code >= 400 && code < 500)
    other -> fail ("expected typed provider rejection, got " <> show other)
  expectRight "delete Unicode key" =<< deleteObject store unicodeKey
  absent <- headObject store unicodeKey
  assertEqual "deleted key absent" (Right Nothing) absent

collectPages :: S3 -> Text -> Maybe Text -> [Text] -> IO [Text]
collectPages store prefix marker accumulated = do
  page <- expectRight "list page" =<< listObjects store prefix Nothing marker 1
  let accumulated' = accumulated <> map fst (lpObjects page)
  if lpIsTruncated page
    then case lpContinuationToken page of
      Nothing -> fail "truncated list omitted continuation marker"
      Just next -> collectPages store prefix (Just next) accumulated'
    else pure accumulated'

admit :: Db -> Text -> Text -> Integer -> IO AdmitOutcome
admit database operation key size =
  admitWithWindow database operation key size 3600

admitWithWindow :: Db -> Text -> Text -> Integer -> Integer -> IO AdmitOutcome
admitWithWindow database operation key size window =
  withTransaction database $ \connection ->
    admitMutation
      connection
      AdmissionRequest
        { admissionIdentity =
            AdmissionIdentity (UserId "integration-user") Put
              (OperationId operation)
        , admissionPoolId = sharedPool
        , admissionRelativePath = RelativePath key
        , admissionFingerprint =
            MutationFingerprint
              Put
              (ObjectKey ("shared/" <> key)) size (Just fixedIdentity)
        }
      AdmissionLimits
        { admissionMaxObjectBytes = 1000
        , admissionIdempotencyWindowSec = window
        }

resetAuthority :: Db -> Integer -> IO ()
resetAuthority database limitBytes = do
  _ <- Db.execute_ database
    "TRUNCATE key_locks, mutations, objects_meta, pools CASCADE"
  configurePoolLimit database sharedPool limitBytes

assertAuthority :: Db -> Int -> Integer -> IO ()
assertAuthority database expectedLocks expectedReserved = do
  locks <- Db.query_ database "SELECT count(*) FROM key_locks"
    :: IO [Only Int]
  reservations <- Db.query database
    "SELECT reserved_bytes FROM pools WHERE pool_id = ?"
    (Only (unPoolId sharedPool))
    :: IO [Only Integer]
  assertEqual "active key locks" [Only expectedLocks] locks
  assertEqual "reserved bytes" [Only expectedReserved] reservations

countPrepared :: [AdmitOutcome] -> Int
countPrepared outcomes = length [() | AdmitPrepared _ <- outcomes]

countReplay :: [AdmitOutcome] -> Int
countReplay outcomes = length [() | AdmitReplay MutationPending _ <- outcomes]

countConflict :: [AdmitOutcome] -> Int
countConflict outcomes = length [() | AdmitConflict <- outcomes]

countRejected :: [AdmitOutcome] -> Int
countRejected outcomes = length [() | AdmitRejected _ _ <- outcomes]

preparedRecord :: [AdmitOutcome] -> MutationRecord
preparedRecord outcomes =
  case [record | AdmitPrepared record <- outcomes] of
    [record] -> record
    _ -> error "expected exactly one prepared record"

fixedIdentity :: ContentSha256
fixedIdentity =
  ContentSha256 "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"

contentIdentity :: BL.ByteString -> ContentSha256
contentIdentity =
  ContentSha256
    . TE.decodeUtf8
    . B16.encode
    . BA.convert
    . (hashlazy :: BL.ByteString -> Digest SHA256)

expectRight :: String -> Either S3Error value -> IO value
expectRight label result =
  case result of
    Left err -> fail (label <> ": " <> show err)
    Right value -> pure value

expectJust :: String -> Either S3Error (Maybe value) -> IO value
expectJust label result = do
  value <- expectRight label result
  case value of
    Nothing -> fail (label <> ": object absent")
    Just present -> pure present

assert :: String -> Bool -> IO ()
assert label condition =
  unless condition (fail label)

assertEqual :: (Eq value, Show value) => String -> value -> value -> IO ()
assertEqual label expected actual =
  unless (expected == actual) $
    fail (label <> ": expected " <> show expected <> ", got " <> show actual)

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