content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE KindSignatures #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PartialTypeSignatures #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE StandaloneKindSignatures #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeFamilyDependencies #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE NoStarIsType #-}
  
  module Main (main, testServer) where
  
  import Control.Applicative (Alternative (empty, many, (<|>)), optional)
  import Control.Concurrent (threadDelay)
  import Control.Concurrent.Async (mapConcurrently)
  import qualified Control.Concurrent.MSem as MSem
  import Control.Concurrent.STM.TVar (TVar, modifyTVar, newTVar, readTVar, readTVarIO, writeTVar)
  import Control.Exception (throw)
  import Control.Monad (ap, forever, void)
  import Control.Monad.Except (MonadError (throwError), runExceptT)
  import Control.Monad.IO.Class (MonadIO (liftIO))
  import Control.Monad.Reader (runReaderT)
  import Control.Monad.STM (STM, atomically, throwSTM)
  import Control.Monad.State.Strict (MonadState (get), evalStateT, modify)
  import Control.Monad.Trans (MonadTrans (lift))
  import Control.Monad.Trans.Free (FreeT)
  import qualified Control.Monad.Yoctoparsec as Yocto
  import qualified Control.Monad.Yoctoparsec.Class as Yocto
  import qualified Data.Attoparsec.Text as Atto (IResult (..), Parser, Result, feed, parse)
  import qualified Data.ByteString.Char8 as BS8
  import qualified Data.ByteString.Lazy as LBS
  import Data.Foldable (Foldable (foldl'))
  import Data.Function (on)
  import Data.Functor (($>), (<&>))
  import qualified Data.HashMap.Strict as HashMap
  import Data.Kind (Constraint, Type)
  import Data.List (sortBy)
  import Data.Maybe (fromMaybe)
  import Data.Text (Text)
  import qualified Data.Text as Text (length, null, pack, stripPrefix, unpack)
  import qualified Data.Text.Encoding as Text
  import Language.SQL.SpiderSQL.Lexer (lexSpiderSQL)
  import Language.SQL.SpiderSQL.Parse (ParserEnv (..), ParserEnvWithGuards (..), mkParserStateTC, mkParserStateUD, spiderSQL, withGuards)
  import Language.SQL.SpiderSQL.Prelude (caselessString)
  import Language.SQL.SpiderSQL.Syntax (SX (..), SpiderSQLTC, SpiderSQLUD)
  import qualified Network.HTTP.Client as HTTP
  import qualified Network.HTTP.Client.TLS as HTTP
  import qualified Picard.Picard.Client as Picard
  import qualified Picard.Picard.Service as Picard
  import qualified Picard.Types as Picard
  import System.Timeout (timeout)
  import Text.Parser.Char (CharParsing (char, string), spaces)
  import Text.Parser.Combinators (Parsing (eof))
  import qualified Thrift.Api as Thrift
  import qualified Thrift.Channel.HeaderChannel as Thrift
  import qualified Thrift.Protocol.Id as Thrift
  import qualified Thrift.Server.CppServer as Thrift
  import qualified Tokenizers (Tokenizer, createTokenizerFromJSONConfig, decode, freeTokenizer)
  import Util.Control.Exception (catchAll)
  import qualified Util.EventBase as Thrift
  
  trace :: forall a. String -> a -> a
  trace _ = id
  
  type RunParsing :: (Type -> Type) -> Type -> Constraint
  class CharParsing m => RunParsing m a where
    type Result m a = r | r -> m a
    runParser :: m a -> Result m a
    feed :: Result m a -> Text -> Result m a
    finalize :: Result m a -> Result m a
    toFeedResult :: Result m a -> Picard.FeedResult
  
  instance RunParsing Atto.Parser a where
    type Result Atto.Parser a = Atto.Result a
    runParser p = Atto.parse p mempty
    feed = Atto.feed
    finalize r = case Atto.feed r mempty of
      Atto.Done notConsumed _ | not (Text.null notConsumed) -> Atto.Fail mempty mempty "Not consumed: notConsumed"
      r' -> r'
    toFeedResult (Atto.Done notConsumed _) = Picard.FeedResult_feedCompleteSuccess (Picard.FeedCompleteSuccess notConsumed)
    toFeedResult (Atto.Partial _) = Picard.FeedResult_feedPartialSuccess Picard.FeedPartialSuccess
    toFeedResult (Atto.Fail i contexts description) =
      Picard.FeedResult_feedParseFailure
        Picard.FeedParseFailure
          { feedParseFailure_input = i,
            feedParseFailure_contexts = Text.pack <$> contexts,
            feedParseFailure_description = Text.pack description
          }
  
  instance RunParsing (FreeT ((->) Char) []) a where
    type Result (FreeT ((->) Char) []) a = [Yocto.Result [] Char a]
    runParser p = Yocto.runParser p
    feed r s = do
      r' <- r
      Yocto.feed r' . Text.unpack $ s
    finalize =
      foldMap
        ( \case
            (Yocto.Done a []) -> pure $ Yocto.Done a []
            (Yocto.Done _ _) -> mempty
            (Yocto.Partial _) -> mempty
        )
    toFeedResult [] =
      Picard.FeedResult_feedParseFailure
        Picard.FeedParseFailure
          { feedParseFailure_input = mempty,
            feedParseFailure_contexts = mempty,
            feedParseFailure_description = "Nothing remains"
          }
    toFeedResult (Yocto.Done _ notConsumed : _) = Picard.FeedResult_feedCompleteSuccess . Picard.FeedCompleteSuccess $ Text.pack notConsumed
    toFeedResult (Yocto.Partial _ : _) = Picard.FeedResult_feedPartialSuccess Picard.FeedPartialSuccess
  
  data PartialParse m a = PartialParse !Text !(Result m a)
  
  deriving stock instance (Show (Result m a)) => Show (PartialParse m a)
  
  type Detokenize = Picard.InputIds -> IO String
  
  data PicardState = PicardState
    { psCounter :: TVar Int,
      psSQLSchemas :: TVar (HashMap.HashMap Picard.DBId Picard.SQLSchema),
      psTokenizer :: TVar (Maybe Tokenizers.Tokenizer),
      psDetokenize :: TVar (Maybe Detokenize),
      psPartialSpiderSQLParsesWithGuardsAndTypeChecking ::
        TVar
          ( HashMap.HashMap
              Picard.InputIds
              (PartialParse (FreeT ((->) Char) []) SpiderSQLTC)
          ),
      psPartialSpiderSQLParsesWithGuards ::
        TVar
          ( HashMap.HashMap
              Picard.InputIds
              (PartialParse Atto.Parser SpiderSQLUD)
          ),
      psPartialSpiderSQLParsesWithoutGuards ::
        TVar
          ( HashMap.HashMap
              Picard.InputIds
              (PartialParse Atto.Parser SpiderSQLUD)
          ),
      psPartialSpiderSQLLexes ::
        TVar
          ( HashMap.HashMap
              Picard.InputIds
              (PartialParse Atto.Parser [String])
          )
    }
  
  initPicardState :: IO PicardState
  initPicardState =
    atomically $
      PicardState
        <$> newTVar 0
          <*> newTVar mempty
          <*> newTVar Nothing
          <*> newTVar Nothing
          <*> newTVar mempty
          <*> newTVar mempty
          <*> newTVar mempty
          <*> newTVar mempty
  
  mkSchemaParser ::
    forall m.
    CharParsing m =>
    HashMap.HashMap Picard.DBId Picard.SQLSchema ->
    m Picard.SQLSchema
  mkSchemaParser sqlSchemas =
    foldl'
      (\agg (dbId, schema) -> agg <|> caselessString (Text.unpack dbId) $> schema)
      empty
      (sortBy (compare `on` (negate . Text.length . fst)) (HashMap.toList sqlSchemas))
  
  mkParser ::
    forall m a.
    (Monad m, CharParsing m) =>
    m Picard.SQLSchema ->
    (Picard.SQLSchema -> m a) ->
    m a
  mkParser schemaParser mkMainParser = do
    _ <-
      spaces
        *> many
          ( char '<'
              *> (string "pad" <|> string "s" <|> string "/s")
              <* char '>'
          )
          <* spaces
    schema <- schemaParser
    _ <- spaces *> char '|' <* spaces
    mkMainParser schema
      <* optional (spaces <* char ';')
      <* eof
  
  getPartialParse ::
    forall m a.
    (RunParsing m a, Monad m) =>
    HashMap.HashMap Picard.DBId Picard.SQLSchema ->
    (Picard.SQLSchema -> m a) ->
    Text ->
    PartialParse m a
  getPartialParse sqlSchemas mkMainParser =
    let schemaParser = mkSchemaParser sqlSchemas
        m = mkParser schemaParser mkMainParser
     in ap PartialParse $ feed (runParser m)
  
  initializeParserCacheSTM ::
    forall m a.
    (RunParsing m a, Monad m) =>
    (Picard.SQLSchema -> m a) ->
    TVar (HashMap.HashMap Picard.DBId Picard.SQLSchema) ->
    TVar (HashMap.HashMap Picard.InputIds (PartialParse m a)) ->
    STM ()
  initializeParserCacheSTM mainParser sqlSchemas partialParses = do
    nukeParserCache partialParses
    partialParse <-
      getPartialParse
        <$> readTVar sqlSchemas
        <*> pure mainParser
        <*> pure mempty
    modifyTVar partialParses (HashMap.insert mempty partialParse)
  
  nukeParserCache ::
    forall m a.
    TVar (HashMap.HashMap Picard.InputIds (PartialParse m a)) ->
    STM ()
  nukeParserCache partialParses = writeTVar partialParses HashMap.empty
  
  data LookupResult a = Cached !a | Fresh !a
    deriving stock (Show)
  
  lookupResultIO ::
    forall m a.
    (RunParsing m a, Monad m) =>
    HashMap.HashMap Picard.DBId Picard.SQLSchema ->
    (Picard.SQLSchema -> m a) ->
    (Picard.InputIds -> IO String) ->
    HashMap.HashMap Picard.InputIds (PartialParse m a) ->
    Picard.InputIds ->
    IO (LookupResult (PartialParse m a))
  lookupResultIO sqlSchemas mkMainParser decode partialParses inputIds =
    case HashMap.lookup inputIds partialParses of
      Just partialParse ->
        trace ("Server: Found inputIds " <> show inputIds) . pure $ Cached partialParse
      Nothing ->
        trace ("Server: Did not find inputIds " <> show inputIds) $ do
          decodedInputIds <- decode inputIds
          let !partialParse = getPartialParse sqlSchemas mkMainParser (Text.pack decodedInputIds)
          pure $ Fresh partialParse
  
  lookupResultWithTimeoutIO ::
    forall m a n.
    (RunParsing m a, Monad m, MonadState DebugInfo n, MonadError Picard.FeedTimeoutFailure n, MonadIO n) =>
    Int ->
    TVar (HashMap.HashMap Picard.DBId Picard.SQLSchema) ->
    (Picard.SQLSchema -> m a) ->
    (Picard.InputIds -> IO String) ->
    TVar (HashMap.HashMap Picard.InputIds (PartialParse m a)) ->
    Picard.InputIds ->
    n (PartialParse m a)
  lookupResultWithTimeoutIO microSeconds sqlSchemas mkMainParser decode partialParses inputIds =
    resultOrTimeout
      microSeconds
      ( do
          schemas <- readTVarIO sqlSchemas
          parses <- readTVarIO partialParses
          !lr <- lookupResultIO schemas mkMainParser decode parses inputIds
          pure lr
      )
      >>= cache
    where
      cache (Cached partialParse) = pure partialParse
      cache (Fresh partialParse) =
        trace ("Server: Cached inputIds " <> show inputIds) . liftIO . atomically $ do
          modifyTVar partialParses $ HashMap.insert inputIds partialParse
          pure partialParse
  
  decodedTokenFromDifferenceIO ::
    (Picard.InputIds -> IO String) ->
    Picard.InputIds ->
    Picard.Token ->
    Text ->
    IO (Text, Maybe Text)
  decodedTokenFromDifferenceIO decode inputIds token decodedInputIds = do
    decoded <- Text.pack <$> decode (inputIds ++ [token])
    pure (decoded, Text.stripPrefix decodedInputIds decoded)
  
  decodedTokenFromDifferenceM ::
    forall m.
    (MonadState DebugInfo m, MonadIO m) =>
    (Picard.InputIds -> IO String) ->
    Picard.InputIds ->
    Picard.Token ->
    Text ->
    m (Text, Text)
  decodedTokenFromDifferenceM decode inputIds token decodedInputIds = do
    (decoded, maybeDecodedToken) <- liftIO $ decodedTokenFromDifferenceIO decode inputIds token decodedInputIds
    _ <- modify (\debugInfo -> debugInfo {debugDecoded = Just decoded})
    maybe
      ( trace
          ("Server: Prefix error " <> show decodedInputIds <> " " <> show decoded)
          . throw
          . Picard.FeedException
          . Picard.FeedFatalException_tokenizerPrefixException
          . Picard.TokenizerPrefixException
          $ "Prefix error."
      )
      ( \decodedToken -> do
          _ <- modify (\debugInfo -> debugInfo {debugDecodedToken = Just decodedToken})
          pure (decoded, decodedToken)
      )
      maybeDecodedToken
  
  data DebugInfo = DebugInfo
    { debugInputIds :: Maybe Picard.InputIds,
      debugToken :: Maybe Picard.Token,
      debugDecodedInputIds :: Maybe Text,
      debugDecodedToken :: Maybe Text,
      debugDecoded :: Maybe Text
    }
  
  mkDebugInfo :: DebugInfo
  mkDebugInfo = DebugInfo Nothing Nothing Nothing Nothing Nothing
  
  resultOrTimeout ::
    forall r m.
    (MonadIO m, MonadState DebugInfo m, MonadError Picard.FeedTimeoutFailure m) =>
    Int ->
    IO r ->
    m r
  resultOrTimeout microSeconds ior = do
    mr <- liftIO $ timeout microSeconds ior
    case mr of
      Just r -> pure r
      Nothing -> do
        DebugInfo {..} <- get
        trace
          ("Server: Timeout error " <> show debugDecodedInputIds <> " " <> show debugDecoded)
          . throwError
          . Picard.FeedTimeoutFailure
          $ "Timeout error."
  
  feedParserWithTimeoutIO ::
    forall m a n.
    (RunParsing m a, MonadState DebugInfo n, MonadIO n, MonadError Picard.FeedTimeoutFailure n) =>
    Int ->
    Result m a ->
    Text ->
    n (Result m a)
  feedParserWithTimeoutIO microSeconds partialParseResult decodedToken = do
    resultOrTimeout
      microSeconds
      $ let !r = case decodedToken of
              "</s>" -> finalize partialParseResult
              s -> feed partialParseResult s
         in pure r
  
  -- | fix me: need one counter for each hashmap
  nukeParserCacheEverySTM ::
    forall m a.
    Int ->
    TVar Int ->
    TVar (HashMap.HashMap Picard.InputIds (PartialParse m a)) ->
    STM ()
  nukeParserCacheEverySTM n counter partialParses = do
    _ <- modifyTVar counter (+ 1)
    c <- readTVar counter
    case c `mod` n of
      0 -> nukeParserCache partialParses
      _ -> pure ()
  
  toFeedResult' ::
    forall m a n.
    (RunParsing m a, MonadState DebugInfo n, Show (Result m a)) =>
    Result m a ->
    n Picard.FeedResult
  toFeedResult' r = do
    let fr = toFeedResult r
    DebugInfo {..} <- get
    pure . flip trace fr $
      "Server: "
        <> ( case fr of
               Picard.FeedResult_feedParseFailure Picard.FeedParseFailure {} -> "Failure"
               Picard.FeedResult_feedTimeoutFailure (Picard.FeedTimeoutFailure msg) -> "Timeout failure: " <> Text.unpack msg
               Picard.FeedResult_feedPartialSuccess Picard.FeedPartialSuccess -> "Partial"
               Picard.FeedResult_feedCompleteSuccess (Picard.FeedCompleteSuccess _) -> "Success"
               Picard.FeedResult_EMPTY -> "Unknown"
           )
        <> ". Input ids were: "
        <> show debugInputIds
        <> ". Token was: "
        <> show debugToken
        <> ". Decoded input ids were: "
        <> show debugDecodedInputIds
        <> ". Decoded token was: "
        <> show debugDecodedToken
        <> ". Result: "
        <> show r
        <> "."
  
  getDetokenize :: TVar (Maybe Detokenize) -> IO Detokenize
  getDetokenize =
    fmap
      ( fromMaybe
          ( throw
              . Picard.FeedException
              . Picard.FeedFatalException_tokenizerNotRegisteredException
              . Picard.TokenizerNotRegisteredException
              $ "Tokenizer has not been registered."
          )
      )
      . readTVarIO
  
  feedIO ::
    forall m a.
    (RunParsing m a, Monad m, Show (Result m a)) =>
    Int ->
    TVar Int ->
    TVar (HashMap.HashMap Picard.DBId Picard.SQLSchema) ->
    (Picard.SQLSchema -> m a) ->
    Detokenize ->
    TVar (HashMap.HashMap Picard.InputIds (PartialParse m a)) ->
    Picard.InputIds ->
    Picard.Token ->
    IO Picard.FeedResult
  feedIO microSeconds counter sqlSchemas mkMainParser detokenize partialParses inputIds token =
    evalStateT
      ( runExceptT
          ( do
              _ <- liftIO . atomically $ nukeParserCacheEverySTM 10000 counter partialParses
              partialParse <- getPartialParseIO detokenize
              liftIO . atomically . modifyTVar partialParses $ HashMap.insert (inputIds ++ [token]) partialParse
              pure partialParse
          )
          >>= toFeedResultIO
      )
      initialDebugInfo
    where
      initialDebugInfo =
        mkDebugInfo
          { debugInputIds = Just inputIds,
            debugToken = Just token
          }
      getPartialParseIO tokenizer = do
        PartialParse decodedInputIds partialParseResult <-
          lookupResultWithTimeoutIO microSeconds sqlSchemas mkMainParser tokenizer partialParses inputIds
        (decoded, decodedToken) <- decodedTokenFromDifferenceM tokenizer inputIds token decodedInputIds
        modify (\debugInfo -> debugInfo {debugDecodedInputIds = Just decodedInputIds})
        partialParseResult' <- feedParserWithTimeoutIO microSeconds partialParseResult decodedToken
        pure $ PartialParse decoded partialParseResult'
      toFeedResultIO (Left timeoutFailure) = pure $ Picard.FeedResult_feedTimeoutFailure timeoutFailure
      toFeedResultIO (Right (PartialParse _ r)) = toFeedResult' r
  
  batchFeedIO ::
    forall m a.
    (RunParsing m a, Monad m, Show (Result m a)) =>
    Int ->
    TVar Int ->
    TVar (HashMap.HashMap Picard.DBId Picard.SQLSchema) ->
    (Picard.SQLSchema -> m a) ->
    Detokenize ->
    TVar (HashMap.HashMap Picard.InputIds (PartialParse m a)) ->
    [Picard.InputIds] ->
    [[Picard.Token]] ->
    IO [Picard.BatchFeedResult]
  batchFeedIO microSeconds counter sqlSchemas mkMainParser detokenize partialParses inputIds topTokens =
    do
      -- traverse
      -- mapPool (length inputIds)
      mapConcurrently
        ( \(batchId, inputIds', token) ->
            feedIO microSeconds counter sqlSchemas mkMainParser detokenize partialParses inputIds' token
              <&> Picard.BatchFeedResult batchId token
        )
        . concat
        . zipWith3
          (\batchId inputIds' tokens -> (batchId,inputIds',) <$> tokens)
          [0 :: Picard.BatchId ..]
          inputIds
        $ topTokens
  
  mapPool :: forall a b t. Traversable t => Int -> (a -> IO b) -> t a -> IO (t b)
  mapPool size f xs = do
    sem <- MSem.new size
    mapConcurrently (MSem.with sem . f) xs
  
  picardHandler :: forall a. PicardState -> Picard.PicardCommand a -> IO a
  picardHandler PicardState {..} = go
    where
      mkSpiderSQLParserWithGuardsAndTypeChecking :: Picard.SQLSchema -> Yocto.Parser [] Char SpiderSQLTC
      mkSpiderSQLParserWithGuardsAndTypeChecking = runReaderT (spiderSQL STC mkParserStateTC) . ParserEnv (ParserEnvWithGuards (withGuards STC))
      mkSpiderSQLParserWithGuards :: Picard.SQLSchema -> Atto.Parser SpiderSQLUD
      mkSpiderSQLParserWithGuards = runReaderT (spiderSQL SUD mkParserStateUD) . ParserEnv (ParserEnvWithGuards (withGuards SUD))
      mkSpiderSQLParserWithoutGuards :: Picard.SQLSchema -> Atto.Parser SpiderSQLUD
      mkSpiderSQLParserWithoutGuards = runReaderT (spiderSQL SUD mkParserStateUD) . ParserEnv (ParserEnvWithGuards (const id))
      mkSpiderSQLLexer :: Picard.SQLSchema -> Atto.Parser [String]
      mkSpiderSQLLexer = runReaderT lexSpiderSQL
      go (Picard.RegisterSQLSchema dbId sqlSchema) =
        trace ("RegisterSQLSchema " <> show dbId) $
          atomically $ do
            r <- readTVar psSQLSchemas
            case HashMap.lookup dbId r of
              Just _ -> throwSTM $ Picard.RegisterSQLSchemaException dbId "Database schema is already registered"
              Nothing -> do
                modifyTVar psSQLSchemas (HashMap.insert dbId sqlSchema)
                initializeParserCacheSTM mkSpiderSQLParserWithGuardsAndTypeChecking psSQLSchemas psPartialSpiderSQLParsesWithGuardsAndTypeChecking
                initializeParserCacheSTM mkSpiderSQLParserWithGuards psSQLSchemas psPartialSpiderSQLParsesWithGuards
                initializeParserCacheSTM mkSpiderSQLParserWithoutGuards psSQLSchemas psPartialSpiderSQLParsesWithoutGuards
                initializeParserCacheSTM mkSpiderSQLLexer psSQLSchemas psPartialSpiderSQLLexes
      go (Picard.RegisterTokenizer jsonConfig) =
        trace "RegisterTokenizer" $ do
          tok <- Tokenizers.createTokenizerFromJSONConfig . Text.encodeUtf8 $ jsonConfig
          tokSem <- MSem.new (1 :: Int)
          maybeOldTokenizer <- atomically $ do
            maybeOldTokenizer <- readTVar psTokenizer
            writeTVar psTokenizer . Just $ tok
            writeTVar psDetokenize . Just $ \inputIds -> MSem.with tokSem (Tokenizers.decode tok $ fromIntegral <$> inputIds)
            initializeParserCacheSTM mkSpiderSQLParserWithGuardsAndTypeChecking psSQLSchemas psPartialSpiderSQLParsesWithGuardsAndTypeChecking
            initializeParserCacheSTM mkSpiderSQLParserWithGuards psSQLSchemas psPartialSpiderSQLParsesWithGuards
            initializeParserCacheSTM mkSpiderSQLParserWithoutGuards psSQLSchemas psPartialSpiderSQLParsesWithoutGuards
            initializeParserCacheSTM mkSpiderSQLLexer psSQLSchemas psPartialSpiderSQLLexes
            pure maybeOldTokenizer
          case maybeOldTokenizer of
            Just oldTok -> Tokenizers.freeTokenizer oldTok
            Nothing -> pure ()
      go (Picard.Feed inputIds token Picard.Mode_PARSING_WITH_GUARDS_AND_TYPE_CHECKING) =
        trace ("Feed parsing with guards " <> show inputIds <> " " <> show token) $ do
          detokenize <- getDetokenize psDetokenize
          feedIO 100000 psCounter psSQLSchemas mkSpiderSQLParserWithGuardsAndTypeChecking detokenize psPartialSpiderSQLParsesWithGuardsAndTypeChecking inputIds token
      go (Picard.Feed inputIds token Picard.Mode_PARSING_WITH_GUARDS) =
        trace ("Feed parsing with guards " <> show inputIds <> " " <> show token) $ do
          detokenize <- getDetokenize psDetokenize
          feedIO 100000 psCounter psSQLSchemas mkSpiderSQLParserWithGuards detokenize psPartialSpiderSQLParsesWithGuards inputIds token
      go (Picard.Feed inputIds token Picard.Mode_PARSING_WITHOUT_GUARDS) =
        trace ("Feed parsing without guards " <> show inputIds <> " " <> show token) $ do
          detokenize <- getDetokenize psDetokenize
          feedIO 100000 psCounter psSQLSchemas mkSpiderSQLParserWithoutGuards detokenize psPartialSpiderSQLParsesWithoutGuards inputIds token
      go (Picard.Feed inputIds token Picard.Mode_LEXING) =
        trace ("Feed lexing " <> show inputIds <> " " <> show token) $ do
          detokenize <- getDetokenize psDetokenize
          feedIO 100000 psCounter psSQLSchemas mkSpiderSQLLexer detokenize psPartialSpiderSQLLexes inputIds token
      go (Picard.Feed _inputIds _token (Picard.Mode__UNKNOWN n)) =
        throw
          . Picard.FeedException
          . Picard.FeedFatalException_modeException
          . Picard.ModeException
          . Text.pack
          $ "Unknown mode " <> show n
      go (Picard.BatchFeed inputIds topTokens Picard.Mode_PARSING_WITH_GUARDS_AND_TYPE_CHECKING) = do
        detokenize <- getDetokenize psDetokenize
        batchFeedIO 10000000 psCounter psSQLSchemas mkSpiderSQLParserWithGuardsAndTypeChecking detokenize psPartialSpiderSQLParsesWithGuardsAndTypeChecking inputIds topTokens
      go (Picard.BatchFeed inputIds topTokens Picard.Mode_PARSING_WITH_GUARDS) = do
        detokenize <- getDetokenize psDetokenize
        batchFeedIO 10000000 psCounter psSQLSchemas mkSpiderSQLParserWithGuards detokenize psPartialSpiderSQLParsesWithGuards inputIds topTokens
      go (Picard.BatchFeed inputIds topTokens Picard.Mode_PARSING_WITHOUT_GUARDS) = do
        detokenize <- getDetokenize psDetokenize
        batchFeedIO 10000000 psCounter psSQLSchemas mkSpiderSQLParserWithoutGuards detokenize psPartialSpiderSQLParsesWithoutGuards inputIds topTokens
      go (Picard.BatchFeed inputIds topTokens Picard.Mode_LEXING) = do
        detokenize <- getDetokenize psDetokenize
        batchFeedIO 10000000 psCounter psSQLSchemas mkSpiderSQLLexer detokenize psPartialSpiderSQLLexes inputIds topTokens
      go (Picard.BatchFeed _inputIds _token (Picard.Mode__UNKNOWN n)) =
        throw
          . Picard.FeedException
          . Picard.FeedFatalException_modeException
          . Picard.ModeException
          . Text.pack
          $ "Unknown mode " <> show n
  
  withPicardServer :: forall a. Thrift.ServerOptions -> (Int -> IO a) -> IO a
  withPicardServer serverOptions action = do
    st <- initPicardState
    Thrift.withBackgroundServer (picardHandler st) serverOptions $
      \Thrift.Server {..} -> action serverPort
  
  picardServerHost :: BS8.ByteString
  picardServerHost = BS8.pack "127.0.0.1"
  
  mkHeaderConfig :: forall t. Int -> Thrift.ProtocolId -> Thrift.HeaderConfig t
  mkHeaderConfig port protId =
    Thrift.HeaderConfig
      { headerHost = picardServerHost,
        headerPort = port,
        headerProtocolId = protId,
        headerConnTimeout = 5000,
        headerSendTimeout = 5000,
        headerRecvTimeout = 5000
      }
  
  testServer :: IO ()
  testServer = do
    let protId = Thrift.binaryProtocolId
        action :: Thrift.Thrift Picard.Picard ()
        action = do
          Picard.registerSQLSchema "test" $
            let columnNames = HashMap.fromList [("0", "column")]
                columnTypes = HashMap.fromList [("0", Picard.ColumnType_NUMBER)]
                tableNames = HashMap.fromList [("0", "table")]
                columnToTable = HashMap.fromList [("0", "0")]
                tableToColumns = HashMap.fromList [("0", ["0"])]
                foreignKeys = HashMap.fromList []
                primaryKeys = []
             in Picard.SQLSchema {sQLSchema_columnNames = columnNames, sQLSchema_columnTypes = columnTypes, sQLSchema_tableNames = tableNames, sQLSchema_columnToTable = columnToTable, sQLSchema_tableToColumns = tableToColumns, sQLSchema_foreignKeys = foreignKeys, sQLSchema_primaryKeys = primaryKeys}
          Picard.registerSQLSchema "car_1" $
            let columnNames = HashMap.fromList [("1", "ContId"), ("10", "ModelId"), ("11", "Maker"), ("12", "Model"), ("13", "MakeId"), ("14", "Model"), ("15", "Make"), ("16", "Id"), ("17", "MPG"), ("18", "Cylinders"), ("19", "Edispl"), ("2", "Continent"), ("20", "Horsepower"), ("21", "Weight"), ("22", "Accelerate"), ("23", "Year"), ("3", "CountryId"), ("4", "CountryName"), ("5", "Continent"), ("6", "Id"), ("7", "Maker"), ("8", "FullName"), ("9", "Country")]
                columnTypes = HashMap.fromList [("1", Picard.ColumnType_NUMBER), ("10", Picard.ColumnType_NUMBER), ("11", Picard.ColumnType_NUMBER), ("12", Picard.ColumnType_TEXT), ("13", Picard.ColumnType_NUMBER), ("14", Picard.ColumnType_TEXT), ("15", Picard.ColumnType_TEXT), ("16", Picard.ColumnType_NUMBER), ("17", Picard.ColumnType_TEXT), ("18", Picard.ColumnType_NUMBER), ("19", Picard.ColumnType_NUMBER), ("2", Picard.ColumnType_TEXT), ("20", Picard.ColumnType_TEXT), ("21", Picard.ColumnType_NUMBER), ("22", Picard.ColumnType_NUMBER), ("23", Picard.ColumnType_NUMBER), ("3", Picard.ColumnType_NUMBER), ("4", Picard.ColumnType_TEXT), ("5", Picard.ColumnType_NUMBER), ("6", Picard.ColumnType_NUMBER), ("7", Picard.ColumnType_TEXT), ("8", Picard.ColumnType_TEXT), ("9", Picard.ColumnType_TEXT)]
                tableNames = HashMap.fromList [("0", "continents"), ("1", "countries"), ("2", "car_makers"), ("3", "model_list"), ("4", "car_names"), ("5", "cars_data")]
                columnToTable = HashMap.fromList [("1", "0"), ("10", "3"), ("11", "3"), ("12", "3"), ("13", "4"), ("14", "4"), ("15", "4"), ("16", "5"), ("17", "5"), ("18", "5"), ("19", "5"), ("2", "0"), ("20", "5"), ("21", "5"), ("22", "5"), ("23", "5"), ("3", "1"), ("4", "1"), ("5", "1"), ("6", "2"), ("7", "2"), ("8", "2"), ("9", "2")]
                tableToColumns = HashMap.fromList [("0", ["1", "2"]), ("1", ["3", "4", "5"]), ("2", ["6", "7", "8", "9"]), ("3", ["10", "11", "12"]), ("4", ["13", "14", "15"]), ("5", ["16", "17", "18", "19", "20", "21", "22", "23"])]
                foreignKeys = HashMap.fromList [("11", "6"), ("14", "12"), ("16", "13"), ("5", "1"), ("9", "3")]
                primaryKeys = ["1", "3", "6", "10", "13", "16"]
             in Picard.SQLSchema {sQLSchema_columnNames = columnNames, sQLSchema_columnTypes = columnTypes, sQLSchema_tableNames = tableNames, sQLSchema_columnToTable = columnToTable, sQLSchema_tableToColumns = tableToColumns, sQLSchema_foreignKeys = foreignKeys, sQLSchema_primaryKeys = primaryKeys}
          manager <- lift $ HTTP.newTlsManagerWith HTTP.tlsManagerSettings
          request <- lift $ HTTP.parseRequest "https://huggingface.co/t5-base/resolve/main/tokenizer.json"
          response <- lift $ HTTP.httpLbs request manager
          Picard.registerTokenizer . Text.decodeUtf8 . LBS.toStrict $ HTTP.responseBody response
          let tokens = [0 .. 32100]
          mapM_
            ( flip catchAll (const (pure Nothing))
                . (Just <$>)
                . (\token -> Picard.feed [0, 794, 1820, 1738, 953, 5, 3297, 440, 29, 45, 953] token Picard.Mode_PARSING_WITH_GUARDS)
            )
            tokens
          let inputIds = [0, 443, 834, 536, 1820, 1738, 3, 17, 5411, 17529, 23, 26, 6, 3, 17, 5411, 17529, 4350, 45, 1440, 38, 3, 17, 536, 1715, 443, 834, 8910, 38, 3, 17, 357, 30, 3, 17, 5411, 17529, 23, 26, 3274, 3, 17, 4416, 17529, 563, 57, 3, 17, 5411, 17529, 23, 26, 578, 3476, 599, 1935, 61, 2490, 220, 7021, 1738, 3, 17, 5411, 17529, 23, 26, 6, 3, 17, 5411, 17529, 4350, 45, 1440, 38, 3, 17, 536, 1715, 443, 834, 8910, 38, 3, 17, 357, 30, 3, 17, 5411, 17529, 23, 26, 3274, 3, 17, 4416, 17529, 1715, 825, 834, 3350, 38, 3, 17, 519, 30, 3, 17, 4416, 23, 26, 3274, 3, 17, 5787, 8337, 213, 3, 17, 5787, 21770, 3274, 96, 3183, 144, 121, 1]
          void $
            Picard.feed
              (take 127 inputIds)
              (inputIds !! 127)
              Picard.Mode_PARSING_WITH_GUARDS
          void $
            Picard.feed
              (take 128 inputIds)
              (inputIds !! 128)
              Picard.Mode_PARSING_WITH_GUARDS
    withPicardServer Thrift.defaultOptions $
      \port ->
        Thrift.withEventBaseDataplane $ \evb -> do
          let headerConf = mkHeaderConfig port protId
          Thrift.withHeaderChannel evb headerConf action
  
  main :: IO ()
  main = do
    st <- initPicardState
    let serverOptions =
          Thrift.ServerOptions
            { desiredPort = Just 9090,
              customFactoryFn = Nothing,
              customModifyFn = Nothing
            }
        action Thrift.Server {} =
          forever (threadDelay maxBound)
    Thrift.withBackgroundServer (picardHandler st) serverOptions action
  