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
    emptyCatalog <- schemaCatalog database
    resetSchema database
    installLegacySchema database
    migrate database "sql" config
    migrate database "sql" config
    adoptedCatalog <- schemaCatalog database
    assertEqual "legacy and empty migration schema equivalence"
      emptyCatalog adoptedCatalog
    verifySchema database
    verifyAdmission database config
  store <- newS3 config
  verifyStorage config store
  putStrLn "WP-04 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_schema.sql"
    , "002_pool_limit_guard.sql"
    , "003_recovery_settlement.sql"
    ]
    (map fromOnly migrations)

schemaCatalog :: Db -> IO [Text]
schemaCatalog database = do
  rows <- Db.query_ database
    "SELECT description FROM ( \
    \SELECT 'column|' || table_name || '|' || ordinal_position::text || '|' \
    \  || column_name || '|' || data_type || '|' || is_nullable || '|' \
    \  || COALESCE(column_default, '') AS description \
    \FROM information_schema.columns \
    \WHERE table_schema = 'public' \
    \UNION ALL \
    \SELECT 'constraint|' || c.relname || '|' || con.conname || '|' \
    \  || con.contype::text || '|' || pg_get_constraintdef(con.oid) \
    \FROM pg_constraint con \
    \JOIN pg_class c ON c.oid = con.conrelid \
    \JOIN pg_namespace n ON n.oid = c.relnamespace \
    \WHERE n.nspname = 'public' \
    \UNION ALL \
    \SELECT 'index|' || tablename || '|' || indexname || '|' || indexdef \
    \FROM pg_indexes WHERE schemaname = 'public' \
    \) catalog ORDER BY description"
    :: IO [Only Text]
  pure (map fromOnly rows)

resetSchema :: Db -> IO ()
resetSchema database = do
  _ <- Db.execute_ database
    "DROP TABLE IF EXISTS schema_migrations, key_locks, objects_meta, \
    \mutations, pools CASCADE"
  pure ()

installLegacySchema :: Db -> IO ()
installLegacySchema database = do
  mapM_ (Db.execute_ database)
    [ "CREATE TABLE pools ( \
      \pool_id TEXT PRIMARY KEY, \
      \limit_bytes BIGINT NOT NULL CHECK (limit_bytes > 0), \
      \used_bytes BIGINT NOT NULL DEFAULT 0 CHECK (used_bytes >= 0), \
      \reserved_bytes BIGINT NOT NULL DEFAULT 0 CHECK (reserved_bytes >= 0))"
    , "CREATE TABLE mutations ( \
      \user_id TEXT NOT NULL, kind TEXT NOT NULL, operation_id TEXT NOT NULL, \
      \pool_id TEXT NOT NULL REFERENCES pools (pool_id), object_key TEXT NOT NULL, \
      \relative_path TEXT NOT NULL, declared_size BIGINT NOT NULL, \
      \content_sha256 TEXT, prior_size BIGINT NOT NULL, prior_sha256 TEXT, \
      \growth_bytes BIGINT NOT NULL, state TEXT NOT NULL, \
      \external_state TEXT NOT NULL DEFAULT 'idle', \
      \created_at TIMESTAMPTZ NOT NULL DEFAULT now(), \
      \updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), \
      \PRIMARY KEY (user_id, kind, operation_id))"
    , "CREATE TABLE key_locks ( \
      \object_key TEXT PRIMARY KEY, user_id TEXT NOT NULL, kind TEXT NOT NULL, \
      \operation_id TEXT NOT NULL)"
    , "CREATE TABLE objects_meta ( \
      \object_key TEXT PRIMARY KEY, \
      \pool_id TEXT NOT NULL REFERENCES pools (pool_id), \
      \size_bytes BIGINT NOT NULL, content_sha256 TEXT NOT NULL, \
      \updated_at TIMESTAMPTZ NOT NULL DEFAULT now())"
    ]
  pure ()

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 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 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

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
      (UserId "integration-user")
      Put
      (OperationId operation)
      sharedPool
      (ObjectKey ("shared/" <> key))
      (RelativePath key)
      size
      (Just fixedIdentity)
      1000
      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