content: {-# LANGUAGE DeriveGeneric       #-}
  {-# LANGUAGE DerivingStrategies  #-}
  {-# LANGUAGE FlexibleInstances   #-}
  {-# LANGUAGE KindSignatures      #-}
  {-# LANGUAGE LambdaCase          #-}
  {-# LANGUAGE OverloadedStrings   #-}
  {-# LANGUAGE PolyKinds           #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module Bookstore
    ( Bug(..)
    , prop_bookstore
    , cleanup
    , setup
    )
    where
  
  import           Control.Concurrent
                     (threadDelay)
  import           Control.Monad.Reader
                     (ReaderT, ask, runReaderT)
  import           Data.Int
                     (Int64)
  import           Data.Kind
                     (Type)
  import           Data.List
                     (dropWhileEnd, group, inits, isInfixOf, sort, tails)
  import           Data.Maybe
                     (fromJust)
  import           Database.PostgreSQL.Simple         as PS
                     (ConnectInfo, Connection, Query, close, connect,
                     connectDatabase, connectHost, connectPassword,
                     connectPort, connectUser, defaultConnectInfo,
                     execute, execute_, query)
  import           Database.PostgreSQL.Simple.Errors
                     (ConstraintViolation(..), catchViolation)
  import           Database.PostgreSQL.Simple.FromRow
                     (FromRow(fromRow), field)
  import           GHC.Generics
                     (Generic, Generic1)
  import           Network.Socket                     as Sock
                     (AddrInfoFlag(AI_NUMERICHOST, AI_NUMERICSERV),
                     Socket, SocketType(Stream), addrAddress, addrFamily,
                     addrFlags, addrProtocol, addrSocketType, close,
                     connect, defaultHints, getAddrInfo, socket)
  import           Prelude
  import           System.Process
                     (callProcess, readProcess)
  import           Test.QuickCheck
                     (Arbitrary(arbitrary), Gen, Property,
                     arbitraryPrintableChar, choose, elements, frequency,
                     ioProperty, listOf, oneof, shrink, suchThat,
                     vectorOf, (===))
  import           Test.QuickCheck.Monadic
                     (monadic)
  import           UnliftIO
                     (IOException, bracket, catch, liftIO, onException,
                     throwIO)
  
  import           Test.StateMachine
  import qualified Test.StateMachine.Types.Rank2      as Rank2
  import           Test.StateMachine.Z
                     (codomain, domain)
  
  ------------------------------------------------------------------------
  
  -- Based on Bookstore case study from:
  -- https://propertesting.com/book_case_study_stateful_properties_with_a_bookstore.html
  
  -- System under test: Parametrized SQL statements
  
  data Book = Book {
      isbn   :: String
    , title  :: String
    , author :: String
    , owned  :: Int
    , avail  :: Int
    } deriving stock (Eq, Show, Generic)
  
  instance ToExpr Book where
  
  instance FromRow Book where
    fromRow = Book <$> field <*> field <*> field <*> field <*> field
  
  handleViolations
    :: (Connection -> IO a)
    -> Connection
    -> IO (Maybe a)
  handleViolations fn cn = catchViolation catcher (Just <$> fn cn)
    where
      catcher _ (UniqueViolation _) = return Nothing
      catcher e _                   = throwIO e
  
  setupTable :: Connection -> IO (Maybe Int64)
  setupTable = handleViolations $ \cn ->
    execute_ cn "CREATE TABLE books (\
                \ isbn varchar(20) PRIMARY KEY,\
                \ title varchar(256) NOT NULL,\
                \ author varchar(256) NOT NULL,\
                \ owned smallint DEFAULT 0,\
                \ available smallint DEFAULT 0\
                \ )"
  
  teardown :: Connection -> IO (Maybe Int64)
  teardown = handleViolations $ \cn -> execute_ cn "DROP TABLE books"
  
  addBook
    :: String -> String -> String
    -> Connection
    -> IO (Maybe Int64)
  addBook isbn_ title_ auth = handleViolations $ \cn ->
    execute cn templ (isbn_, title_, auth)
      where templ = "INSERT INTO books (isbn, title, author, owned, available)\
                    \ VALUES ( ?, ?, ?, 0, 0 )"
  
  update :: Query -> String -> Connection -> IO Int64
  update prepStmt param cn = execute cn prepStmt [param]
  
  addCopy :: String -> Connection -> IO (Maybe Int64)
  addCopy isbn_ = handleViolations $ update templ isbn_
    where templ = "UPDATE books SET\
                  \ owned = owned + 1,\
                  \ available = available + 1\
                  \ WHERE isbn = ?"
  
  borrowCopy :: String -> Connection -> IO (Maybe Int64)
  borrowCopy isbn_ = handleViolations (update templ isbn_)
    where templ = "UPDATE books SET available = available - 1 \
                  \ WHERE isbn = ? AND available > 0"
  
  returnCopy :: Bug -> String -> Connection -> IO (Maybe Int64)
  returnCopy bug isbn_ = handleViolations (update templ isbn_)
    where
      templ = if bug == Bug
        then "UPDATE books SET available = available + 1 WHERE isbn = ?;"
        else "UPDATE books SET available = available + 1\
             \ WHERE isbn = ? AND available < owned;"
  
  select :: Query -> String -> Connection -> IO [Book]
  select prepStmt param cn = query cn prepStmt [param]
  
  findByAuthor, findByTitle
    :: Bug -> String -> Connection -> IO (Maybe [Book])
  findByAuthor bug s = case bug of
    Injection -> handleViolations (select templ $ "%"++s++"%")
    _          -> handleViolations (select templ $ "%"++(sanitize s)++"%")
    where templ = "SELECT * FROM books WHERE author LIKE ?"
  
  findByTitle bug s = case bug of
    Injection -> handleViolations (select templ $ "%"++s++"%")
    _          -> handleViolations (select templ $ "%"++(sanitize s)++"%")
    where templ = "SELECT * FROM books WHERE title LIKE ?"
  
  findByIsbn :: String -> Connection -> IO (Maybe [Book])
  findByIsbn s = handleViolations $ select "SELECT * FROM books WHERE isbn = ?" s
  
  -- XXX
  sanitize :: String -> String
  sanitize = concatMap (\c -> if c `elem` ['_', '%', '\\'] then '\\':[c] else [c])
  
  withConnection :: ConnectInfo -> (Connection -> IO a) -> IO a
  withConnection connInfo = bracket (PS.connect connInfo) PS.close
  
  -- Modeling
  
  newtype Model r = Model [(Isbn, Book)]
    deriving stock (Generic, Show)
  
  instance ToExpr (Model Concrete) where
  
  initModel :: Model v
  initModel = Model []
  
  newtype Isbn = Isbn { getString :: String } deriving stock (Eq, Show)
  
  instance ToExpr Isbn where
    toExpr = toExpr . show
  
  instance Arbitrary Isbn where
    arbitrary = oneof [isbn10, isbn13]
      where
        isbn10 = do
          gr <- isbnGroup
          t  <- pubTitle (9 - length gr)
          c  <- elements ('X':['0'..'9'])
          return $ Isbn (gr ++ "-" ++ t ++ "-" ++ [c])
        isbn13 = do
          code <- elements ["978", "979"]
          Isbn x <- isbn10
          return $ Isbn (code ++ "-" ++ x)
  
  isbnGroup :: Gen String
  isbnGroup = elements $ show <$> concat [
      [0..5] :: [Integer]
    , [7]
    , [80..94]
    , [600..621]
    , [950..989]
    , [9926..9989]
    , [99901..99976]
    ]
  
  pubTitle :: Int -> Gen String
  pubTitle size = do
    i  <- choose (1, size - 1)
    v1 <- vectorOf i $ choose ('0', '9')
    v2 <- vectorOf (size - i) $ choose ('0', '9')
    return $ v1 ++ "-" ++ v2
  
  data Tag
    = New
    | Exist
    | Invalid
    | Avail
    | Unavail
    | Taken
    | Full
    deriving stock (Eq, Show)
  
  data Command (r :: Type -> Type)
    = NewBook      Tag Isbn String String
    | AddCopy      Tag Isbn
    | Borrow       Tag Isbn
    | Return       Tag Isbn
    | FindByAuthor Tag String
    | FindByTitle  Tag String
    | FindByIsbn   Tag Isbn
    deriving stock (Show, Generic1)
  
  instance Rank2.Foldable    Command where
  instance Rank2.Functor     Command where
  instance Rank2.Traversable Command where
  
  data Response (r :: Type -> Type)
    = Rows [Book]
    | Updated
    | Inserted
    | NotFound
    | UniqueError
    | OtherError
    deriving stock (Eq, Show, Generic1)
  
  instance Rank2.Foldable Response where
  
  generator :: Model Symbolic -> Maybe (Gen (Command Symbolic))
  generator (Model m) = Just $ oneof $ [
      NewBook      New     <$> newIsbn <*> stringGen <*> stringGen
    , AddCopy      Invalid <$> newIsbn
    , Borrow       Invalid <$> newIsbn
    , Return       Invalid <$> newIsbn
    , FindByIsbn   Invalid <$> newIsbn
    , FindByTitle  Invalid <$> searchStringGen
    , FindByAuthor Invalid <$> searchStringGen
    ] ++ if null m then [] else [
      NewBook      Exist   <$> existIsbn <*> stringGen <*> stringGen
    , AddCopy      Exist   <$> existIsbn
    , Borrow       Avail   <$> existIsbn
    , Borrow       Unavail <$> existIsbn
    , Return       Taken   <$> existIsbn
    , Return       Full    <$> existIsbn
    , FindByIsbn   Exist   <$> existIsbn
    , FindByTitle  Exist   <$> genTitle
    , FindByAuthor Exist   <$> genAuthor
    ]
    where
      newIsbn = arbitrary `suchThat` \x -> x `notElem` domain m
      existIsbn = elements $ domain m
      genTitle  = elements $ (title  <$> codomain m) >>= infixes
      genAuthor = elements $ (author <$> codomain m) >>= infixes
  
  stringGen :: Gen String
  stringGen = listOf arbitraryPrintableChar
  
  searchStringGen :: Gen String
  searchStringGen = listOf $ frequency [ (7, arbitraryPrintableChar)
                                       , (3, elements ['_', '%']) ]
  
  infixes :: Ord a => [a] -> [[a]]
  infixes l = map head . group . sort $ inits l >>= tails
  
  -- How to shrink Commands
  
  shrinker :: Model Symbolic -> Command Symbolic -> [Command Symbolic]
  shrinker _ (NewBook tag key x y) = [ NewBook tag key x y' | y' <- shrink y ] ++
                                     [ NewBook tag key x' y | x' <- shrink x ]
  shrinker _ (FindByAuthor tag a)  = [ FindByAuthor tag a' | a' <- shrink a ]
  shrinker _ (FindByTitle  tag t)  = [ FindByTitle  tag t' | t' <- shrink t ]
  shrinker _ _ = []
  
  
  
  -- Pre-requisites and invariants
  
  preconditions :: Model Symbolic -> Command Symbolic -> Logic
  preconditions (Model m) cmd = case cmd of
    AddCopy    Invalid key   -> Not $ hasKey key
    Borrow     Invalid key   -> Not $ hasKey key
    Return     Invalid key   -> Not $ hasKey key
    FindByIsbn Invalid key   -> Not $ hasKey key
    NewBook    New   key _ _ -> Not $ hasKey key
    NewBook    Exist key _ _ -> hasKey key
    AddCopy    Exist key     -> hasKey key
    FindByIsbn Exist key     -> hasKey key
    Borrow     Avail   key -> keyPred key (\b -> avail b .> 0)
    Borrow     Unavail key -> keyPred key (\b -> avail b .== 0)
    Return     Taken   key -> keyPred key (\b -> owned b .> avail b)
    Return     Full    key -> keyPred key (\b -> owned b .== avail b)
    FindByAuthor Invalid x -> Not $ Exists values (Boolean . isInfixOf x . author)
    FindByTitle Invalid x  -> Not $ Exists values (Boolean . isInfixOf x . title)
    FindByTitle Exist x -> Exists values (Boolean . isInfixOf x . title)
    _ -> Top
    where
      values = codomain m
      hasKey key = key `member` (domain m)
      keyPred key p = maybe Bot p (lookup key m)
  
  postconditions :: Model Concrete -> Command Concrete -> Response Concrete -> Logic
  postconditions (Model m) cmd resp = case cmd of
    NewBook    Exist _ _ _ -> resp .== UniqueError
    NewBook    New   _ _ _ -> resp .== Inserted
    AddCopy      Invalid _ -> resp .== NotFound
    Return       Invalid _ -> resp .== NotFound
    Borrow       Invalid _ -> resp .== NotFound
    Return       Full    _ -> resp .== NotFound
    Borrow       Unavail _ -> resp .== NotFound
    Borrow       Avail   _ -> resp .== Updated
    AddCopy      Exist   _ -> resp .== Updated
    Return       Taken   _ -> resp .== Updated
    FindByIsbn   Invalid _ -> resp .== Rows []
    FindByAuthor Invalid _ -> resp .== Rows []
    FindByTitle  Invalid _ -> resp .== Rows []
    FindByIsbn Exist key -> case lookup key m of
      Just x -> Rows [x] .== resp
      _      -> error "Should not happen"
    FindByAuthor Exist x -> case resp of
      Rows rs -> Forall rs (Boolean . isInfixOf x . author) .&&
                 Forall rs (\b -> Just b .== lookup (Isbn $ isbn b) m)
      _ -> Bot .// "findByAuthor returned " ++ (show resp)
    FindByTitle Exist x -> case resp of
      Rows rs -> Forall rs (Boolean . isInfixOf x . title) .&&
                 Forall rs (\b -> Just b .== lookup (Isbn $ isbn b) m)
      _ -> Bot .// "findByTitle returned " ++ (show resp)
    _ -> Bot
  
  -- Transitions of the state machine that models SUT
  
  transitions :: Model r -> Command r -> Response r -> Model r
  transitions (Model m) cmd _ = Model $ case cmd of
    NewBook New key t a -> (key, Book (getString key) t a 0 0):m
    AddCopy Exist key -> map (applyForKey key $ incOwned . incAvail) m
    Borrow  Avail key -> map (applyForKey key decAvail) m
    Return  Taken key -> map (applyForKey key incAvail) m
    _ -> m
    where
      applyForKey key fn (k, v) = (k, if k == key then fn v else v)
      incOwned row = row { owned = 1 + owned row }
      incAvail row = row { avail = 1 + avail row }
      decAvail row = row { avail = (avail row) - 1 }
  
  -- Semantics
  
  data Bug = Bug | NoBug | Injection deriving stock Eq
  
  semantics :: Bug -> Command Concrete -> ReaderT ConnectInfo IO (Response Concrete)
  semantics bug cmd = do
    connInfo <- ask
    liftIO $ withConnection connInfo $ \cn -> case cmd of
      NewBook _ (Isbn key) t a  -> toResp insertRes <$> addBook key t a cn
      AddCopy _ (Isbn key)      -> toResp updateRes <$> addCopy key cn
      Borrow  _ (Isbn key)      -> toResp updateRes <$> borrowCopy key cn
      Return  _ (Isbn key)      -> toResp updateRes <$> returnCopy bug key cn
      FindByAuthor _ x          -> toResp Rows <$> findByAuthor bug x cn
      FindByTitle  _ x          -> toResp Rows <$> findByTitle bug x cn
      FindByIsbn   _ (Isbn key) -> toResp Rows <$> findByIsbn key cn
    where toResp = maybe UniqueError
          updateRes 0 = NotFound
          updateRes 1 = Updated
          updateRes _ = OtherError
          insertRes 1 = Inserted
          insertRes _ = OtherError
  
  -- Mock is currently not used by the library
  
  mock :: Model Symbolic -> Command Symbolic -> GenSym (Response Symbolic)
  mock (Model m) cmd = return $ case cmd of
    NewBook New _ _ _ -> Inserted
    NewBook Exist _ _ _ -> UniqueError
    AddCopy Invalid _ -> NotFound
    AddCopy Exist _ -> Updated
    Borrow  Invalid _ -> NotFound
    Borrow  Avail _ -> Updated
    Borrow  Unavail _ -> NotFound
    Return  Invalid _ -> NotFound
    Return  Taken _ -> Updated
    Return  Full _ -> NotFound
    FindByTitle Exist x -> Rows $ filter (isInfixOf x . title) (codomain m)
    FindByAuthor Exist x -> Rows $ filter (isInfixOf x . author) (codomain m)
    FindByIsbn Exist key -> Rows [fromJust (lookup key m)]
    FindByTitle Invalid _ -> Rows []
    FindByIsbn Invalid _ -> Rows []
    FindByAuthor Invalid _ -> Rows []
    _ -> error $ (show cmd) ++ " should not happen"
  
  -- Property
  
  sm :: Bug -> StateMachine Model Command (ReaderT ConnectInfo IO) Response
  sm bug = StateMachine initModel transitions preconditions postconditions
    Nothing generator shrinker (semantics bug) mock noCleanup
  
  runner :: IO String -> ReaderT ConnectInfo IO Property -> IO Property
  runner io p = do
    dbIp <- io
    let connInfo = defaultConnectInfo {
      connectUser     = "postgres"
    , connectPassword = "mysecretpassword"
    , connectDatabase = "postgres"
    , connectPort     = 5432
    , connectHost     = dbIp
    }
    bracket (withConnection connInfo setupTable)
            (const (withConnection connInfo teardown))
            (const (runReaderT p connInfo))
  
  prop_bookstore :: Bug -> IO String -> Property
  prop_bookstore bug io =
    forAllCommands sm' Nothing $ \cmds -> monadic (ioProperty . runner io) $ do
      (hist, _, res) <- runCommands sm' cmds
  
      prettyCommands sm' hist $ res === Ok
    where
      sm' = sm bug
  
  -- Setup PostgreSQL db in Docker
  
  setup :: IO (String, String)
  setup = do
    pid <- trim <$> readProcess "docker"
      [ "run"
      , "-d"
      , "-e", "POSTGRES_PASSWORD=mysecretpassword"
      , "postgres:10.2"
      ] ""
    ip <- trim <$> readProcess "docker"
      [ "inspect"
      , pid
      , "--format"
      , "'{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}'"
      ] ""
    healthyDb pid ip `onException` callProcess "docker" [ "rm", "-f", "-v", pid ]
    return (pid, ip)
    where
      trim :: String -> String
      trim = dropWhileEnd isGarbage . dropWhile isGarbage
        where
          isGarbage = flip elem ['\'', '\n']
  
      healthyDb :: String -> String -> IO ()
      healthyDb pid ip = do
        sock <- go 10
        Sock.close sock
        where
          go :: Int -> IO Socket
          go 0 = error "healthyDb: db isn't healthy"
          go n = do
            let hints = defaultHints
                  { addrFlags      = [AI_NUMERICHOST , AI_NUMERICSERV]
                  , addrSocketType = Stream
                  }
            addr : _ <- getAddrInfo (Just hints) (Just ip) (Just "5432")
            sock <- socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
            (Sock.connect sock (addrAddress addr) >>
              readProcess "docker"
                [ "exec"
                , "-u", "postgres"
                , pid
                , "psql", "-U", "postgres", "-d", "postgres", "-c", "SELECT 1 + 1"
                ] "" >> return sock)
              `catch` (\(_ :: IOException) -> do
                          threadDelay 1000000
                          go (n - 1))
  
  cleanup :: (String, String) -> IO ()
  cleanup (pid, _) = callProcess "docker" [ "rm", "-f", "-v", pid ]
  