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

module Main where

import Control.Exception (Exception, throw)
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.List (foldl', isSuffixOf)
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Set qualified as S
import Data.Text qualified as T
import Debug.Trace (trace)
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)

debug :: Show a => String -> a -> a
debug s x = trace (printf "%s: %s" s (show x)) x

main :: IO ()
main = do
  setCurrentDirectory "./data"
  putStrLn "> SELECT * FROM c"
  query' $ Select All "c" [] []

  putStrLn "\n> SELECT * FROM j"
  query' $ Select All "j" [] []

  putStrLn "\n> SELECT c.id, j.id, is_j FROM c LEFT JOIN j ON j.id == c.j_id"
  query' $
    Select
      ( Only
          [ Qualified "c" "id",
            Qualified "j" "id",
            Unqualified "is_j"
          ]
      )
      "c"
      [ LeftJoin
          "j"
          [ Eq (Qualified "j" "id") (Qualified "c" "j_id")
          ]
      ]
      []

  putStrLn "\n> SELECT c.id, j.id FROM c RIGHT JOIN j ON j.id == c.j_id"
  query' $
    Select
      ( Only
          [ Qualified "c" "id",
            Qualified "j" "id"
          ]
      )
      "c"
      [ RightJoin
          "j"
          [ Eq (Qualified "j" "id") (Qualified "c" "j_id")
          ]
      ]
      []

  putStrLn "\n> SELECT c.id, j.id FROM c FULL JOIN j ON j.id == c.j_id"
  query' $
    Select
      ( Only
          [ Qualified "c" "id",
            Qualified "j" "id"
          ]
      )
      "c"
      [ FullJoin
          "j"
          [ Eq (Qualified "j" "id") (Qualified "c" "j_id")
          ]
      ]
      []

data Query
  = Select FieldSelector Collection [Join FilePath] Where
  deriving (Show)

data FieldSelector
  = All
  | Only [Field]
  deriving (Show)

data Field
  = Qualified FilePath T.Text
  | Unqualified T.Text
  deriving (Show)

type Collection = FilePath

data Join a
  = LeftJoin a Where
  | RightJoin a Where
  | FullJoin a Where
  deriving (Show)

type Where = [Cmp]

data Cmp
  = Eq Field Field
  deriving (Show)

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

{-
instance IsString Query where
  fromString =
    either throw id (P.parse parser "" s)-}

query :: Query -> IO [J.Value]
query (Select fs c js ws) = do
  c' <-
    mapM (fmap (Record c) . decodeFile . (c </>)) =<< ls c
  js' <-
    mapM
      ( \j ->
          case j of
            LeftJoin c ws ->
              fmap (\j' -> LeftJoin (map (Record c) j') ws) . mapM (decodeFile . (c </>))
                =<< ls c
            RightJoin c ws ->
              fmap (\j' -> RightJoin (map (Record c) j') ws) . mapM (decodeFile . (c </>))
                =<< ls c
            FullJoin c ws ->
              fmap (\j' -> FullJoin (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" (S.listDirectory c)

combine :: [Record J.Value] -> [Join [Record J.Value]] -> [[Record J.Value]]
combine vs js = combine' (map (: []) vs) js
  where
    combine' vss [] = vss
    combine' vss (LeftJoin js ws : jss) =
      combine'
        ( concatMap
            ( \vs -> case filter (satisfies ws) $ map (\j -> vs ++ [j]) js of
                [] -> [vs]
                vs' -> vs'
            )
            vss
        )
        jss
    combine' vss (RightJoin js ws : jss) =
      combine'
        ( concatMap
            ( \j -> case filter (satisfies ws) $ map (\vs -> vs ++ [j]) vss of
                [] -> [[j]]
                vs' -> vs'
            )
            js
        )
        jss
    combine' vss (FullJoin 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" $
    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_ :: Where -> [[Record J.Value]] -> [[Record J.Value]]
where_ ws vss = filter (satisfies ws) vss

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

satisfy :: Cmp -> [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