{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

module Main where

import Control.Exception (Exception, throw)
import Control.Monad (void)
import Control.Monad.Combinators.NonEmpty qualified as PN
import Data.Aeson qualified as J
import Data.Aeson.Key qualified as JK
import Data.Aeson.KeyMap qualified as JM
import Data.ByteString.Lazy.Char8 qualified as LB
import Data.Char (isSpace)
import Data.List (intercalate, isSuffixOf)
import Data.List.NonEmpty qualified as N
import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
import Data.String (IsString (fromString))
import Data.Text qualified as T
import Data.Vector qualified as V
import Data.Void (Void)
import Field qualified as F
import Record qualified as R
import Store qualified as S
import System.Directory (setCurrentDirectory)
import System.FilePath ((</>))
import Text.Megaparsec qualified as P
import Text.Megaparsec.Char qualified as P
import Text.Megaparsec.Char.Lexer qualified as P
import Text.Printf (printf)

main :: IO ()
main = do
  setCurrentDirectory "./data"

  mapM_
    ( \q -> do
        printf "> %s\n" (show q)
        query' q
    )
    [ "SELECT * FROM c",
      "SELECT * FROM j",
      "SELECT c.id, j.id, is_j FROM c LEFT JOIN j ON j.id == c.j_id",
      "SELECT c.id, j.id FROM c RIGHT JOIN j ON j.id == c.j_id",
      "SELECT c.id, j.id FROM c FULL JOIN j ON j.id == c.j_id",
      "SELECT c.id, j FROM c EMBED j ON j.id == c.j_id"
    ]

data Query
  = Select
      FieldSelector
      Collection
      (JoinClauses FilePath)
      (EmbedClauses FilePath)
      (Maybe WhereClause)

instance Show Query where
  show (Select fs c js es w) =
    intercalate " " $
      catMaybes $
        [ Just "SELECT",
          Just (showFieldSelector fs),
          Just "FROM",
          Just (showCollection c),
          showJoinClauses js,
          showEmbedClauses es,
          showWhereClause w
        ]
    where
      showFieldSelector All = "*"
      showFieldSelector (Only (N.toList -> fs)) =
        intercalate ", " (map showField fs)
      showField = F.toString
      showCollection c = c
      showJoinClauses js = case map showJoinClause js of
        [] -> Nothing
        xs -> Just (intercalate " " xs)
      showJoinClause (JoinClause t c w) =
        intercalate " " $
          catMaybes $
            [ Just (showJoinType t),
              Just (showCollection c),
              Just "ON",
              showWhereClause w
            ]
      showJoinType JoinLeft = "LEFT JOIN"
      showJoinType JoinRight = "RIGHT JOIN"
      showJoinType JoinFull = "FULL JOIN"
      showEmbedClauses js = case map showEmbedClause js of
        [] -> Nothing
        xs -> Just (intercalate " " xs)
      showEmbedClause (EmbedClause c w) =
        intercalate " " $
          catMaybes $
            [ Just "EMBED",
              Just (showCollection c),
              Just "ON",
              showWhereClause w
            ]
      showWhereClause = showWhereClauseWith id
      showWhereClause' = showWhereClauseWith (\x -> "(" <> x <> ")")
      showWhereClauseWith _ Nothing = Nothing
      showWhereClauseWith wrap (Just (And ws)) = Just (wrap (intercalate "AND" (mapMaybe (showWhereClause' . Just) ws)))
      showWhereClauseWith wrap (Just (Or ws)) = Just (wrap (intercalate "OR" (mapMaybe (showWhereClause' . Just) ws)))
      showWhereClauseWith _ (Just (Where p)) = Just (showComparison p)
      showComparison (Eq a b) = intercalate " " [showField a, "==", showField b]

data FieldSelector
  = All
  | Only (N.NonEmpty F.Field)
  deriving (Show)

type Collection = FilePath

type JoinClauses a = [JoinClause a]

data JoinClause a
  = JoinClause JoinType a (Maybe WhereClause)
  deriving (Show)

data JoinType
  = JoinLeft
  | JoinRight
  | JoinFull
  deriving (Show)

type EmbedClauses a = [EmbedClause a]

data EmbedClause a
  = EmbedClause a (Maybe WhereClause)
  deriving (Show)

data WhereClause
  = And [WhereClause]
  | Or [WhereClause]
  | Where Comparison
  deriving (Show)

data Comparison
  = Eq F.Field F.Field
  deriving (Show)

data ParseError = ParseError String
  deriving (Show)

instance Exception ParseError

instance IsString Query where
  fromString =
    either (throw . ParseError . P.errorBundlePretty @String @Void) id
      . P.parse parser ""
    where
      parser = do
        void $ P.many P.space1
        select
        fs <- fieldSelector
        from
        c <- collection
        js <- joinClauses
        es <- embedClauses
        w <- P.optional do
          where_
          whereClause
        P.eof
        pure $ Select fs c js es w

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

      and = void $ lexeme (P.string "AND")
      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")

      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 <- field
        eq
        b <- field
        pure $ Eq a b

      fieldSelector =
        P.choice
          [ do
              void $ lexeme $ P.string "*"
              pure All,
            do
              Only <$> PN.sepBy1 field comma
          ]

      field :: P.Parsec Void String F.Field
      field =
        lexeme . P.choice $
          [ P.try
              do
                F.Qualified
                  <$> (fieldPart <* P.string ".")
                  <*> (N.map T.pack <$> PN.sepBy1 fieldPart (P.string ".")),
              do
                F.Unqualified
                  <$> (N.map T.pack <$> PN.sepBy1 fieldPart (P.string "."))
          ]

      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
          ]

query :: Query -> IO [J.Value]
query (Select fs c js es w) = do
  c' <-
    mapM (fmap (R.fromValue c) . decodeFile . (c </>)) =<< ls c
  js' <-
    mapM
      ( \(JoinClause t c w) ->
          fmap (\j' -> JoinClause t (map (R.fromValue c) j') w) . mapM (decodeFile . (c </>))
            =<< ls c
      )
      js
  es' <-
    mapM
      ( \(EmbedClause c w) ->
          fmap (\e' -> EmbedClause (R.fromValue c e') w) . mapM (decodeFile . (c </>))
            =<< ls c
      )
      es
  pure $ map (select fs) $ where_ w $ embed es' $ combine c' js'
  where
    ls c =
      filter (not . (isSuffixOf "/"))
        <$> S.withStore "." "HEAD" do
          S.listDirectory c

-- TODO use fold
embed ::
  EmbedClauses (R.Record [J.Value]) ->
  [R.Records J.Value] ->
  [R.Records J.Value]
embed es vss = embed' vss es
  where
    embed' vss [] = vss
    embed' vss (EmbedClause (R.Record c es) w : ess) =
      embed'
        ( map
            ( \vs ->
                let es' :: [J.Value]
                    es' = filter (\e -> satisfies w (vs ++ [R.Record c e])) es
                 in vs
                      ++ [ R.Record
                             c
                             ( J.Object
                                 ( JM.singleton
                                     (JK.fromString c)
                                     (J.Array (V.fromList es'))
                                 )
                             )
                         ]
            )
            vss
        )
        ess

-- TODO rename `join`

-- TODO use fold
combine ::
  R.Records J.Value ->
  JoinClauses (R.Records J.Value) ->
  [R.Records J.Value]
combine vs js = combine' (map (: []) vs) js
  where
    combine' vss [] = vss
    combine' vss (JoinClause JoinLeft js w : jss) =
      combine'
        ( concatMap
            ( \vs -> case filter (satisfies w) $ map (\j -> vs ++ [j]) js of
                [] -> [vs]
                vs' -> vs'
            )
            vss
        )
        jss
    combine' vss (JoinClause JoinRight js w : jss) =
      combine'
        ( concatMap
            ( \j -> case filter (satisfies w) $ map (\vs -> vs ++ [j]) vss of
                [] -> [[j]]
                vs' -> vs'
            )
            js
        )
        jss
    combine' vss (JoinClause JoinFull js w : jss) =
      combine'
        ( concatMap
            ( \vs -> case filter (satisfies w) $ map (\j -> vs ++ [j]) js of
                [] -> [vs]
                vs' -> vs'
            )
            vss
            ++ concatMap
              ( \j -> case filter (satisfies w) $ map (\vs -> vs ++ [j]) vss of
                  [] -> [[j]]
                  _ -> []
              )
              js
        )
        jss

data DecodeException = DecodeException
  deriving (Show)

instance Exception DecodeException

decodeFile :: J.FromJSON a => Collection -> IO a
decodeFile fp = S.withStore "." "HEAD" do
  fromMaybe (throw DecodeException) . J.decode <$> S.readFile fp

select :: FieldSelector -> R.Records J.Value -> J.Value
select All vs = R.disjointUnions (map R.toValue vs)
select (Only fs) vs = R.select (N.toList fs) vs

where_ :: Maybe WhereClause -> [R.Records J.Value] -> [R.Records J.Value]
where_ w = filter (satisfies w)

satisfies :: Maybe WhereClause -> R.Records J.Value -> Bool
satisfies Nothing _ = True
satisfies (Just (And ws)) vs = all (\w -> satisfies (Just w) vs) ws
satisfies (Just (Or ws)) vs = any (\w -> satisfies (Just w) vs) ws
satisfies (Just (Where p)) vs = satisfy p vs

satisfy :: Comparison -> R.Records J.Value -> Bool
satisfy (Eq f g) vs = unique f vs == unique g vs

data DuplicateField' = DuplicateField'
  deriving (Show)

instance Exception DuplicateField'

unique :: F.Field -> R.Records J.Value -> J.Value
unique f as = case mapMaybe (R.lookup f) as of
  [v] -> v
  (_ : _) -> throw DuplicateField'

query' :: Query -> IO ()
query' q = mapM_ (LB.putStrLn . J.encode) =<< query q