{-# 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 (foldl', intercalate, isSuffixOf)
import Data.List.NonEmpty qualified as N
import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
import Data.Set qualified as S
import Data.String (IsString (fromString))
import Data.Text qualified as T
import Data.Vector qualified as V
import Data.Void (Void)
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 (Qualified c k) = c <> "." <> T.unpack k
      showField (Unqualified k) = T.unpack k
      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 Field)
  deriving (Show)

data Field
  = Qualified Collection T.Text
  | Unqualified T.Text
  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 Field Field
  deriving (Show)

data Record a
  = Record Collection a
  deriving (Show, Eq)

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 Field
      field =
        lexeme . P.choice $
          [ P.try do
              Qualified
                <$> (P.takeWhile1P Nothing (\c -> not (isSpace c) && c /= ',' && c /= '.'))
                <*> (P.string "." >> T.pack <$> P.takeWhile1P Nothing (\c -> not (isSpace c) && c /= ',')),
            do
              Unqualified
                <$> (T.pack <$> P.takeWhile1P Nothing (\c -> not (isSpace 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 (Record c) . decodeFile . (c </>)) =<< ls c
  js' <-
    mapM
      ( \(JoinClause t c w) ->
          fmap (\j' -> JoinClause t (map (Record c) j') w) . mapM (decodeFile . (c </>))
            =<< ls c
      )
      js
  es' <-
    mapM
      ( \(EmbedClause c w) ->
          fmap (\e' -> EmbedClause (Record 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

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

combine :: [Record J.Value] -> JoinClauses [Record J.Value] -> [[Record 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 -> [Record J.Value] -> J.Value
select All vs =
  join' (map (\(Record _ v) -> v) vs)
select (Only fs) vs =
  mergeUnsafe (join' (map ((\(Record _ v) -> v) . select' fs) vs)) v0
  where
    v0 =
      joinUnsafe $
        mapMaybe
          ( \f -> case f of
              Qualified c k -> Just $ J.Object $ JM.singleton (JK.fromText (T.pack c <> "." <> k)) J.Null
              Unqualified k -> Just $ J.Object $ JM.singleton (JK.fromText k) J.Null
          )
          (N.toList fs)

select' :: N.NonEmpty Field -> Record J.Value -> Record J.Value
select' (N.toList -> fs) (Record c (J.Object kvs)) =
  Record c . J.Object $
    JM.fromList . mapMaybe match . JM.toList $
      kvs
  where
    match (k, v) = case filter (matches (Record c (JK.toText k))) fs of
      (Qualified _ _ : _) -> Just (JK.fromString (c <> "." <> JK.toString k), v)
      (Unqualified _ : _) -> Just (k, v)
      _ -> Nothing

matches :: Record T.Text -> Field -> Bool
matches (Record c k) (Qualified c' k') = c == c' && k == k'
matches (Record _ k) (Unqualified k') = k == k'

join' :: [J.Value] -> J.Value
join' = foldl' merge (J.Object JM.empty)

joinUnsafe :: [J.Value] -> J.Value
joinUnsafe = foldl' mergeUnsafe (J.Object JM.empty)

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

satisfies :: Maybe WhereClause -> [Record 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 -> [Record J.Value] -> Bool
satisfy (Eq f g) vs = unique f vs == unique g vs

data DuplicateField' = DuplicateField'
  deriving (Show)

instance Exception DuplicateField'

unique :: Field -> [Record J.Value] -> J.Value
unique f as = case mapMaybe (get f) as of
  [Record _ v] -> v
  (_ : _) -> throw DuplicateField'

get :: Field -> Record J.Value -> Maybe (Record J.Value)
get (Unqualified k) (Record c (J.Object kvs)) =
  Record c <$> JM.lookup (JK.fromText k) kvs
get (Qualified c' k) (Record c (J.Object kvs))
  | c' == c = Record c <$> JM.lookup (JK.fromText k) kvs
  | otherwise = Nothing

data DuplicateField = DuplicateField
  deriving (Show)

instance Exception DuplicateField

mergeUnsafe :: J.Value -> J.Value -> J.Value
mergeUnsafe (J.Object kvs) (J.Object kvs') =
  J.Object (JM.union kvs kvs')

merge :: J.Value -> J.Value -> J.Value
merge v@(J.Object kvs) v'@(J.Object kvs') =
  case disjoint kvs kvs' of
    True -> mergeUnsafe v v'
    False -> throw DuplicateField

disjoint :: JM.KeyMap v -> JM.KeyMap v -> Bool
disjoint kvs kvs' =
  let ks = S.fromList (JM.keys kvs)
      ks' = S.fromList (JM.keys kvs')
   in S.size ks + S.size ks' == S.size (ks `S.union` ks')

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