module Main (main) where

import           Control.Concurrent (threadDelay)
import           Control.Monad (unless)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import           Data.Text (Text)
import qualified Data.Text as T
import           Database.PostgreSQL.Simple
  ( FromRow
  , Only (..)
  , Query
  , ToRow
  )
import           Nuthatch.App.Env (newEnv)
import           Nuthatch.Config
import           Nuthatch.Db
import qualified Nuthatch.Db as Db
import           Nuthatch.Mutation.Service (dispatchAndFinalize)
import           Nuthatch.Protocol
import           Nuthatch.Redis
import           Nuthatch.S3
import           Nuthatch.Types
import           Servant.Types.SourceT
import           System.Environment (getArgs, getEnv)
import           System.IO (BufferMode (LineBuffering), hSetBuffering, stdout)
import           Text.Read (readMaybe)

main :: IO ()
main = do
  hSetBuffering stdout LineBuffering
  cfg <- loadConfig
  withDb (cfgPostgres cfg) $ \database -> do
    migrate database "sql" (cfgProduct cfg)
    withRedis (cfgRedis cfg) $ \redis -> do
      store <- newS3 (cfgS3 cfg)
      getArgs >>= \case
        [] -> getEnv "NUTHATCH_RECOVERY_PHASE" >>=
          runPhase cfg database redis store
        ["wait-phase", operation, expected] ->
          waitForPhase database operation expected
        ["terminal", operation, expected] ->
          verifyTerminal database operation expected
        ["active", operation, expectedReservation] ->
          maybe
            (fail "invalid expected reservation")
            (verifyActive database operation)
            (readMaybe expectedReservation)
        ["object", expectation, key] ->
          verifyObject store expectation (ObjectKey (T.pack key))
        ["cache-invalidated"] ->
          cacheGet redis "nuthatch:list:recovery-user:shared::" >>=
            requireEqual "cache invalidation" Nothing
        ["database-available"] ->
          waitForQuery database "SELECT 1" ((), Only (1 :: Int))
        ["log-exhausted"] -> do
          logBytes <- BS.readFile "/run/nuthatch-recovery/app.log"
          unless ("\"result\":\"exhausted\"" `BS.isInfixOf` logBytes) $
            fail "recovery log omitted exhausted retry result"
        _ -> fail "unknown recovery probe mode"

waitForPhase :: Db -> String -> String -> IO ()
waitForPhase database operation expected =
  waitForQuery database
    "SELECT state || ':' || external_state FROM mutations WHERE operation_id = ?"
    (Only (T.pack operation), Only (T.pack expected))

verifyTerminal :: Db -> String -> String -> IO ()
verifyTerminal database operation expected = do
  waitForQuery database
    "SELECT state FROM mutations WHERE operation_id = ?"
    (Only (T.pack operation), Only (T.pack expected))
  locks <- Db.query database
    "SELECT count(*) FROM key_locks WHERE operation_id = ?"
    (Only (T.pack operation))
    :: IO [Only Int]
  requireEqual "terminal key lock" [Only 0] locks
  reservations <- Db.query database
    "SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
    ()
    :: IO [Only Integer]
  requireEqual "terminal reservation" [Only 0] reservations

verifyActive :: Db -> String -> Integer -> IO ()
verifyActive database operation expectedReservation = do
  locks <- Db.query database
    "SELECT count(*) FROM key_locks WHERE operation_id = ?"
    (Only (T.pack operation))
    :: IO [Only Int]
  requireEqual "active key lock" [Only 1] locks
  reservations <- Db.query database
    "SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
    ()
    :: IO [Only Integer]
  requireEqual "active reservation"
    [Only expectedReservation]
    reservations

verifyObject :: S3 -> String -> ObjectKey -> IO ()
verifyObject store expectation key = do
  observed <- headObject store key
  case (expectation, observed) of
    ("present", Right (Just _)) -> pure ()
    ("absent", Right Nothing) -> pure ()
    _ -> fail
      ("object expectation failed: " <> expectation <> " " <> show observed)

waitForQuery
  :: (Eq row, Show row, ToRow parameters, FromRow row)
  => Db
  -> Query
  -> (parameters, row)
  -> IO ()
waitForQuery database queryText (parameters, expected) =
  go (120 :: Int)
  where
    go remaining = do
      rows <- Db.query database queryText parameters
      if rows == [expected]
        then pure ()
        else
          if remaining <= 1
            then fail
              ("timed out waiting for database value: " <> show rows)
            else threadDelay 250000 >> go (remaining - 1)

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

runPhase :: Config -> Db -> Redis -> S3 -> String -> IO ()
runPhase cfg database redis store phase =
  case phase of
    "before-dispatch" ->
      preparePut database redis "before-dispatch" 13 phaseIdentity >>= checkpoint
    "during-put" -> do
      rec <- preparePut database redis "during-put" slowSize slowIdentity
      beginDispatch cfg database rec
      result <- putObjectStream
        store
        (mrObjectKey rec)
        slowSize
        (Just slowIdentity)
        65536
        (const (pure ()))
        (slowZeroSource slowSize)
      fail ("during-put completed before process loss: " <> show result)
    "live-timeout" -> do
      rec <- preparePut database redis "live-timeout" slowSize slowIdentity
      result <- dispatchAndFinalize
        (newEnv cfg database redis store)
        recoverySession
        Shared
        rec
        "request timed out"
        (putObjectStream
          store
          (mrObjectKey rec)
          slowSize
          (Just slowIdentity)
          65536
          (const (pure ()))
          (slowZeroSource slowSize)
          >>= pure . either (Left . T.pack . show) Right)
      case result of
        Left "request timed out" -> checkpoint rec
        other -> fail ("live timeout returned unexpected result: " <> show other)
    "after-dispatch" -> do
      rec <- preparePut database redis "after-dispatch" 13 phaseIdentity
      beginDispatch cfg database rec
      expectRight =<< putObject
        store
        (mrObjectKey rec)
        13
        (Just phaseIdentity)
        "recovered-put"
      checkpoint rec
    "after-settlement" -> do
      rec <- preparePut database redis "after-settlement" 13 phaseIdentity
      beginDispatch cfg database rec
      expectRight =<< putObject
        store
        (mrObjectKey rec)
        13
        (Just phaseIdentity)
        "recovered-put"
      settled <- withTransaction database $ \connection ->
        markExternalSettled connection rec
      unless settled (fail "could not record settlement")
      checkpoint rec
    "delete-settled" -> do
      rec <- prepareDelete cfg database redis store
      beginDispatch cfg database rec
      expectRight =<< deleteObject store (mrObjectKey rec)
      settled <- withTransaction database $ \connection ->
        markExternalSettled connection rec
      unless settled (fail "could not record delete settlement")
      checkpoint rec
    _ -> fail ("unknown recovery phase: " <> phase)

preparePut
  :: Db
  -> Redis
  -> Text
  -> Integer
  -> ContentSha256
  -> IO MutationRecord
preparePut database redis name size identity = do
  cacheRecoveryListing redis
  outcome <- withTransaction database $ \connection ->
    admitMutation
      connection
      AdmissionRequest
        { admissionIdentity =
            AdmissionIdentity recoveryUser Put (phaseOperation name)
        , admissionPoolId = sharedPool
        , admissionRelativePath = RelativePath (name <> ".bin")
        , admissionFingerprint =
            MutationFingerprint Put (phaseKey name) size (Just identity)
        }
      (AdmissionLimits (1024 * 1024 * 1024) 3600)
  prepared outcome

prepareDelete :: Config -> Db -> Redis -> S3 -> IO MutationRecord
prepareDelete cfg database redis store = do
  let name = "delete-settled"
      key = phaseKey name
      body = "delete-me"
      size = toInteger (BL.length body)
  cacheRecoveryListing redis
  expectRight =<< putObject store key size (Just deleteIdentity) body
  _ <- Db.execute database
    "UPDATE pools SET used_bytes = ? WHERE pool_id = ?"
    (size, unPoolId sharedPool)
  _ <- Db.execute database
    "INSERT INTO objects_meta \
    \(object_key, pool_id, size_bytes, content_sha256) VALUES (?,?,?,?)"
    ( unObjectKey key
    , unPoolId sharedPool
    , size
    , unContentSha256 deleteIdentity
    )
  outcome <- withTransaction database $ \connection ->
    admitMutation
      connection
      AdmissionRequest
        { admissionIdentity =
            AdmissionIdentity recoveryUser Delete (phaseOperation name)
        , admissionPoolId = sharedPool
        , admissionRelativePath = RelativePath (name <> ".bin")
        , admissionFingerprint =
            MutationFingerprint Delete key 0 Nothing
        }
      AdmissionLimits
        { admissionMaxObjectBytes = cfgMaxObjectBytes (cfgProduct cfg)
        , admissionIdempotencyWindowSec = cfgIdempotencyWindowSec (cfgProduct cfg)
        }
  prepared outcome

beginDispatch :: Config -> Db -> MutationRecord -> IO ()
beginDispatch cfg database rec = do
  started <- withTransaction database $ \connection -> do
    applying <- markApplying connection rec
    pending <- markExternalPending
      connection
      rec
      DispatchTiming
        { dispatchSettlementBoundSec = cfgObjectRequestTimeoutSec (cfgRecovery cfg) * 2
        , dispatchWorkerGraceSec = cfgRecoveryGraceSec (cfgRecovery cfg)
        }
    pure (applying && pending)
  unless started (fail "could not begin dispatch")

cacheRecoveryListing :: Redis -> IO ()
cacheRecoveryListing redis =
  cacheSet redis "nuthatch:list:recovery-user:shared::" 600 "stale"

prepared :: AdmitOutcome -> IO MutationRecord
prepared outcome =
  case outcome of
    AdmitPrepared rec -> pure rec
    other -> fail ("phase admission failed: " <> show other)

checkpoint :: MutationRecord -> IO a
checkpoint rec = do
  putStrLn
    ("recovery checkpoint "
      <> T.unpack (unOperationId (mrOperationId rec)))
  spin

spin :: IO a
spin = threadDelay 1000000 >> spin

slowZeroSource :: Integer -> SourceT IO BS.ByteString
slowZeroSource total =
  SourceT $ \consume -> consume (go total)
  where
    go remaining
      | remaining <= 0 = Stop
      | otherwise = Effect $ do
          threadDelay 200000
          let size = fromInteger (min remaining 65536)
          pure (Yield (BS.replicate size 0) (go (remaining - toInteger size)))

expectRight :: Either S3Error a -> IO a
expectRight result =
  case result of
    Right value -> pure value
    Left err -> fail ("object operation failed: " <> show err)

phaseOperation :: Text -> OperationId
phaseOperation name = OperationId ("recovery-" <> name)

phaseKey :: Text -> ObjectKey
phaseKey name = ObjectKey ("shared/recovery-" <> name <> ".bin")

recoveryUser :: UserId
recoveryUser = UserId "recovery-user"

recoverySession :: SessionUser
recoverySession =
  SessionUser "recovery" "recovery" recoveryUser (Just "Recovery")

slowSize :: Integer
slowSize = 20 * 1024 * 1024

phaseIdentity :: ContentSha256
phaseIdentity =
  ContentSha256
    "1cfe64e6951d3efd9cfe80b53c65cb55d0e87e2a3c72474e6afd8a89fb0e2f69"

deleteIdentity :: ContentSha256
deleteIdentity =
  ContentSha256
    "af77c0879d9cc69681600bb17fa8fedccfe7ae605b65974fee792581ed63a84c"

slowIdentity :: ContentSha256
slowIdentity =
  ContentSha256
    "cd52d81e25f372e6fa4db2c0dfceb59862c1969cab17096da352b34950c973cc"