{-# OPTIONS_GHC -fno-warn-orphans #-}

module Store.Query.Parser () where

import Control.Exception (throw)
import Control.Monad (void)
import Data.Aeson qualified as J
import Data.Aeson.Key qualified as JK
import Data.Aeson.KeyMap qualified as JM
import Data.Char (isSpace)
import Data.Map qualified as M
import Data.String (IsString (fromString))
import Data.Text qualified as T
import Data.Vector qualified as V
import Data.Void (Void)
import Store.Exception (ParseError (ParseError))
import Store.Query.Type
import Text.Megaparsec qualified as P
import Text.Megaparsec.Char qualified as P
import Text.Megaparsec.Char.Lexer qualified as P

instance IsString Query where
  fromString =
    either (throw . ParseError . P.errorBundlePretty @String @Void) id
      . P.parse parser ""
    where
      parser = do
        void $ P.many P.space1
        P.choice
          [ do
              delete
              from
              c <- collection
              w <- P.optional do
                where_
                whereClause
              pure $ Delete c w,
            do
              insert
              vs <- objects
              into
              c <- collection
              pure $ Insert vs c,
            do
              select
              fs <- fieldSelector
              from
              c <- collection
              js <- joinClauses
              es <- embedClauses
              w <- P.optional do
                where_
                whereClause
              l <- P.optional do
                limitClause
              o <- P.optional do
                offsetClause
              pure $ Select fs c js es w l o,
            do
              update
              c <- collection
              set
              v <- object
              w <-
                P.optional do
                  where_
                  whereClause
              pure $ Update c v w
          ]
          <* P.eof

      lexeme :: P.Parsec Void String a -> P.Parsec Void String a
      lexeme = P.lexeme $ void $ P.takeWhileP Nothing isSpace

      set = void $ lexeme (P.string "SET")
      update = void $ lexeme (P.string "UPDATE")
      into = void $ lexeme (P.string "INTO")
      insert = void $ lexeme (P.string "INSERT")
      delete = void $ lexeme (P.string "DELETE")
      and = void $ lexeme (P.string "AND")
      colon = void $ lexeme (P.string ":")
      comma = void $ lexeme (P.string ",")
      embed = void $ lexeme (P.string "EMBED")
      eq = void $ lexeme (P.string "==")
      from = void $ lexeme (P.string "FROM")
      on = void $ lexeme (P.string "ON")
      or = void $ lexeme (P.string "OR")
      select = void $ lexeme (P.string "SELECT")
      where_ = void $ lexeme (P.string "WHERE")
      offset_ = void $ lexeme (P.string "OFFSET")
      limit_ = void $ lexeme (P.string "LIMIT")

      collection = lexeme $ P.takeWhile1P (Just "collection") (not . isSpace)

      joinClauses = P.many joinClause

      joinClause = do
        t <- joinType
        c <- collection
        w <- P.optional do
          on
          whereClause
        pure $ JoinClause t c w

      embedClauses = P.many embedClause

      embedClause = do
        embed
        c <- collection
        w <- P.optional do
          on
          whereClause
        pure $ EmbedClause c w

      whereClause =
        P.choice
          [ P.try (And . map Where <$> P.sepBy1 comparison and),
            P.try (Or . map Where <$> P.sepBy1 comparison or),
            Where <$> comparison
          ]

      comparison = do
        a <- P.choice [Left <$> value, Right <$> field]
        eq
        b <- P.choice [Left <$> value, Right <$> field]
        pure $ Eq a b

      offsetClause = do
        offset_
        Offset <$> lexeme P.decimal

      limitClause = do
        limit_
        Limit <$> lexeme P.decimal

      fieldSelector =
        lexeme $
          P.choice
            [ do
                void $ lexeme $ P.string "{"
                kvs <-
                  P.sepBy
                    ( P.choice
                        [ P.try do
                            k <-
                              lexeme $
                                P.takeWhile1P
                                  Nothing
                                  ( \c ->
                                      not (isSpace c)
                                        && c /= '.'
                                        && c /= ','
                                        && c /= ':'
                                  )
                            lexeme $ colon
                            v <- fieldSelector
                            pure (k, v),
                          do
                            f@(Field c ks) <- field
                            let k
                                  | null ks = c
                                  | otherwise = T.unpack (last ks)
                            pure (k, SelectField f)
                        ]
                    )
                    comma
                void $ lexeme $ P.string "}"
                pure $ SelectObject (M.fromList kvs),
              do
                SelectField <$> field
            ]

      objects = P.sepBy object comma

      object = lexeme do
        void $ lexeme $ P.string "{"
        kvs <-
          P.sepBy
            ( do
                k <-
                  lexeme $ do
                    void $ lexeme $ P.string "\""
                    P.takeWhile1P
                      Nothing
                      (/= '\"')
                      <* (void $ lexeme $ P.string "\"")
                lexeme $ colon
                v <- value
                pure (JK.fromString k, v)
            )
            comma
        void $ lexeme $ P.string "}"
        pure $ J.Object (JM.fromList kvs)

      value =
        P.choice
          [ object,
            array,
            string,
            number,
            bool,
            null_
          ]

      array = lexeme do
        void $ lexeme $ P.string "["
        vs <- P.sepBy value comma
        void $ lexeme $ P.string "]"
        pure $ J.Array (V.fromList vs)

      string = lexeme do
        void $ lexeme $ P.string "\""
        s <- P.takeWhileP Nothing (\c -> c /= '\"')
        void $ lexeme $ P.string "\""
        pure $ J.String (T.pack s)

      number = lexeme do
        J.Number <$> P.scientific

      bool = lexeme do
        J.Bool
          <$> P.choice
            [ const True <$> P.string "true",
              const False <$> P.string "false"
            ]

      null_ = lexeme do
        const J.Null <$> P.string "null"

      field :: P.Parsec Void String Field
      field =
        lexeme do
          Field
            <$> fieldPart
            <*> ( P.choice
                    [ do
                        void $ P.string "."
                        map T.pack <$> P.sepBy fieldPart (P.string "."),
                      do
                        pure []
                    ]
                )

      fieldPart :: P.Parsec Void String String
      fieldPart = P.takeWhile1P Nothing (\c -> not (isSpace c) && c /= ',' && c /= '.')

      joinType :: P.Parsec Void String JoinType
      joinType =
        P.choice
          [ do
              void $ lexeme (P.string "LEFT")
              void $ lexeme (P.string "JOIN")
              pure JoinLeft,
            do
              void $ lexeme (P.string "RIGHT")
              void $ lexeme (P.string "JOIN")
              pure JoinRight,
            do
              void $ lexeme (P.string "FULL")
              void $ lexeme (P.string "JOIN")
              pure JoinFull
          ]