{-# 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) WhereClauses

instance Show Query where
  show (Select fs c js ws) =
    intercalate " " $
      catMaybes $
        [ Just "SELECT",
          Just (showFieldSelector fs),
          Just "FROM",
          Just (showCollection c),
          showJoinClauses js,
          showWhereClauses ws
        ]
    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 ws) =
        intercalate " " $
          catMaybes $
            [ Just (showJoinType t),
              Just (showCollection c),
              Just "ON",
              showWhereClauses ws
            ]
      showJoinType JoinLeft = "LEFT JOIN"
      showJoinType JoinRight = "RIGHT JOIN"
      showJoinType JoinFull = "FULL JOIN"
      showWhereClauses ws = case map showWhereClause ws of
        [] -> Nothing
        xs -> Just (intercalate " " xs)
      showWhereClause (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 WhereClauses
  deriving (Show)

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

type WhereClauses = [WhereClause]

data WhereClause
  = 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
        ws <-
          fromMaybe [] <$> P.optional do
            where_
            whereClauses1
        P.eof
        pure $ Select fs c js ws

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

      comma, eq, from, on, select, where_ :: P.Parsec Void String ()
      comma = void $ lexeme (P.string ",")
      eq = void $ lexeme (P.string "==")
      from = void $ lexeme (P.string "FROM")
      on = void $ lexeme (P.string "ON")
      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
        on
        ws <- whereClauses
        pure $ JoinClause t c ws

      whereClauses1 = P.some whereClause
      whereClauses = P.many whereClause

      whereClause = 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 ws) = do
  c' <-
    mapM (fmap (Record c) . decodeFile . (c </>)) =<< ls c
  js' <-
    mapM
      ( \(JoinClause t c ws) ->
          fmap (\j' -> JoinClause t (map (Record c) j') ws) . mapM (decodeFile . (c </>))
            =<< ls c
      )
      js
  pure $ map (select fs) $ where_ ws $ 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 ws : jss) =
      combine'
        ( concatMap
            ( \vs -> case filter (satisfies ws) $ map (\j -> vs ++ [j]) js of
                [] -> [vs]
                vs' -> vs'
            )
            vss
        )
        jss
    combine' vss (JoinClause JoinRight js ws : jss) =
      combine'
        ( concatMap
            ( \j -> case filter (satisfies ws) $ map (\vs -> vs ++ [j]) vss of
                [] -> [[j]]
                vs' -> vs'
            )
            js
        )
        jss
    combine' vss (JoinClause JoinFull js ws : jss) =
      combine'
        ( concatMap
            ( \vs -> case filter (satisfies ws) $ map (\j -> vs ++ [j]) js of
                [] -> [vs]
                vs' -> vs'
            )
            vss
            ++ concatMap
              ( \j -> case filter (satisfies ws) $ 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_ :: WhereClauses -> [[Record J.Value]] -> [[Record J.Value]]
where_ ws vss = filter (satisfies ws) vss

satisfies :: WhereClauses -> [Record J.Value] -> Bool
satisfies ws vs = all (\w -> satisfy w vs) ws

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