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

module Main where

import Control.Exception (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.ByteString.Lazy.Char8 qualified as LB
import Data.Char (isSpace)
import Data.List (foldl', intercalate, isSuffixOf)
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.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"
    ]

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

instance Show Query where
  show (Select fs c js w) =
    intercalate " " $
      catMaybes $
        [ Just "SELECT",
          Just (showFieldSelector fs),
          Just "FROM",
          Just (showCollection c),
          showJoinClauses js,
          showWhereClause w
        ]
    where
      showFieldSelector All = "*"
      showFieldSelector (Only 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"
      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 [Field]
  deriving (Show)

data Field
  = Qualified FilePath 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)

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

data Comparison
  = Eq Field Field
  deriving (Show)

data Record a
  = Record FilePath 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
        w <- P.optional do
          where_
          whereClause
        P.eof
        pure $ Select fs c js 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 ",")
      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

      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 <$> P.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 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
  pure $ map (select fs) $ where_ w $ combine c' js'
  where
    ls c =
      filter (not . (isSuffixOf "/"))
        <$> S.withStore "." "HEAD" do
          S.listDirectory c

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 => FilePath -> 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
          )
          fs

select' :: [Field] -> Record J.Value -> Record J.Value
select' 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