{-# LANGUAGE OverloadedStrings #-}

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')
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Set qualified as S
import Data.String (IsString (fromString))
import Data.Text qualified as T
import Debug.Trace (trace)
import System.Directory (listDirectory, setCurrentDirectory)
import System.FilePath (takeExtension, (</>))
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 = do
  setCurrentDirectory "./data"
  -- query "SELECT . FROM c"
  query' $ Select [Unqualified "."] "c" [] []
  putStrLn ""
  -- query "SELECT id FROM c"
  query' $ Select [Qualified "c" "id"] "c" [] []
  putStrLn ""
  -- query "SELECT c.id, j.id, is_j FROM c JOIN j WHERE j.id == c.j_id"
  query' $
    Select
      [ Qualified "c" "id",
        Qualified "j" "id",
        Unqualified "is_j"
      ]
      "c"
      ["j"]
      [ Eq (Qualified "j" "id") (Qualified "c" "j_id")
      ]

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

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

type Collection = FilePath

type Join = [FilePath]

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 -> mapM (fmap (Record j) . decodeFile . (j </>)) =<< ls j) js
  pure $ map (select fs) $ where_ ws $ combine c' js'

combine c = combine' (map (: []) c)
  where
    combine' cs [] = cs
    combine' cs (js : jss) = combine' [c ++ [j] | c <- cs, j <- js] jss

ls :: FilePath -> IO [FilePath]
ls =
  fmap (filter ((== ".json") . takeExtension)) . listDirectory

data DecodeException = DecodeException
  deriving (Show)

instance Exception DecodeException

decodeFile :: J.FromJSON a => FilePath -> IO a
decodeFile =
  fmap (fromMaybe (throw DecodeException)) . J.decodeFileStrict

select :: [Field] -> [Record J.Value] -> J.Value
select fs vs = join' $ map (select' fs) vs

select' :: [Field] -> Record J.Value -> Record J.Value
select' [Unqualified "."] v = v
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' :: [Record J.Value] -> J.Value
join' vs = foldl' merge (J.Object JM.empty) (map (\(Record _ v) -> v) vs)

where_ :: Where -> [[Record J.Value]] -> [[Record J.Value]]
where_ ws vss = filter (\vs -> all (\w -> satisfy w vs) ws) vss
  where
    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

merge :: J.Value -> J.Value -> J.Value
merge (J.Object kvs) (J.Object kvs') =
  case disjoint kvs kvs' of
    True -> J.Object (JM.union kvs kvs')
    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