{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE QuasiQuotes #-}

module AutoTypes.Unify
  ( T (..),
    toString,
    fromJson,
    unify1,
  )
where

import Control.Arrow
import Control.Exception (Exception, throw)
import qualified Data.Aeson as A
import qualified Data.Aeson.Key as K
import qualified Data.Aeson.KeyMap as KM
import Data.Aeson.QQ
import Data.List (intercalate, nub)
import Data.Map (Map)
import qualified Data.Map as M
import qualified Data.Map.Merge.Lazy as M
import qualified Data.Vector as V
import Debug.Trace
import Prelude hiding (null)

data T
  = List (Maybe T)
  | Object (Map String T)
  | Option (Maybe T)
  | Scalar String
  | Union [T]
  deriving (Eq, Show)

instance A.ToJSON T where
  toJSON (List t) = A.toJSON t
  toJSON (Object ts) = A.toJSON ts
  toJSON (Option t) = A.toJSON t
  toJSON (Scalar s) = A.toJSON s
  toJSON (Union ts) = A.toJSON ts

toString :: T -> String
toString = intercalate "\n" . map (\(n, s) -> indent n s) . toString_ 0

toString_ :: Int -> T -> [(Int, String)]
toString_ n (Object kvs) =
  concat
    [ [(n, "{")],
      concat . M.elems $
        M.mapWithKey (\s t -> (n + 1, s) : toString_ (n + 2) t) kvs,
      [(n, "}")]
    ]
toString_ n (Scalar s) = [(n, s)]
toString_ n (Union ts) = concatMap (toString_ n) ts
toString_ n (Option Nothing) = [(n, "null")]
toString_ n (Option (Just t)) = map (second (++ "?")) (toString_ n t)

indent n = (++) (replicate (4 * n) ' ')

union :: [T] -> T
union ts =
  case ts of
    [t] -> t
    ts -> Union ts

unify1 :: T -> T -> [T]
unify1 (Scalar n) (Scalar m)
  | n == m = [Scalar n]
  | otherwise = [union [Scalar n, Scalar m]]
unify1 l@(Object ls) r@(Object rs) =
  let os =
        ( map Object . traverse id $
            let f _ (l@(Option _)) = [l]
                f _ l = [Option (Just l)]
             in M.merge
                  (M.mapMissing f)
                  (M.mapMissing f)
                  (M.zipWithMatched (\_ l r -> unify1 l r))
                  ls
                  rs
        )
   in os
        ++ ( if l `subst` r || r `subst` l
               then []
               else [union [l, r]]
           )
unify1 (Option Nothing) (Option Nothing) = [Option Nothing]
unify1 (Option (Just l)) (Option Nothing) = [Option (Just l)]
unify1 (Option Nothing) (Option (Just r)) = [Option (Just r)]
unify1 (Option (Just l)) (Option (Just r)) = map (Option . Just) (unify1 l r)
unify1 (Option Nothing) r = [Option (Just r)]
unify1 (List Nothing) (List Nothing) = [List Nothing]
unify1 (List Nothing) (List (Just t)) = [List (Just t)]
unify1 (List (Just t)) (List Nothing) = [List (Just t)]
unify1 (List (Just l)) (List (Just r)) =
  if l == r || r `subst` l
    then [List (Just l)]
    else
      if l `subst` r
        then [List (Just r)]
        else [List (Just (union [l, r]))]
unify1 l (Option Nothing) = [Option (Just l)]
unify1 l (Option (Just r)) = map (Option . Just) (unify1 l r)
unify1 (Option (Just l)) r = map (Option . Just) (unify1 l r)
unify1 (Union ls) (Union rs) = [union (ls ++ rs)]
unify1 (Union ls) r = [union (ls ++ [r])]
unify1 l (Union rs) = [union ([l] ++ rs)]
unify1 l r = [union [l, r]]

subst :: T -> T -> Bool
subst (Object l) (Object r) =
  and
    ( map
        ( \(k, t) -> case (t, M.lookup k r) of
            (t, Just t') -> t `subst` t'
            _ -> False
        )
        (M.toList l)
    )
subst (Option Nothing) (Option _) = True
subst l (Option Nothing) = True
subst (Option (Just l)) (Option (Just r)) = l `subst` r
subst l (Option (Just r)) = l `subst` r
subst l r = l == r

lims :: [T] -> [T]
lims ts = nub [t | t <- ts, all (\t' -> not (t `subst` t') || t' == t) ts]

unify l r =
  let s =
        unlines
          ( [ toString l ++ "," ++ toString r,
              "---"
            ]
              ++ map toString t
          )
      t = unify1 l r
   in Debug.Trace.trace s t

{-
unify :: T -> T -> Either (T, T) T
unify (Scalar n) (Scalar m)
  | n == m = Right (Scalar n)
  | otherwise = Left (Scalar n, Scalar m)
unify (Object ls) (Object rs) =
  let f _ (l@(Option _)) = Right l
      f _ l = Right (Option (Just l))
   in case M.merge
        (M.mapMissing f)
        (M.mapMissing f)
        (M.zipWithMatched (\_ l r -> unify l r))
        ls
        rs
        & M.partition isRight
        & ( \(rs, ls) ->
              (M.map (fromRight undefined) rs, M.elems (M.map (fromLeft undefined) ls))
          ) of
        (_, e : _) -> Left e
        (lrs, []) -> Right (Object lrs)
unify (Option Nothing) (Option Nothing) = Right (Option Nothing)
unify (Option (Just l)) (Option Nothing) = Right (Option (Just l))
unify (Option Nothing) (Option (Just r)) = Right (Option (Just r))
unify (Option (Just l)) (Option (Just r)) = Option . Just <$> unify l r
-}

object :: Map String T -> T
object = Object

list :: Maybe T -> T
list = List

string, number, bool, null :: T
string = Scalar "string"
number = Scalar "number"
bool = Scalar "bool"
null = Option Nothing

data InferException = InferException [T]
  deriving (Show)

instance Exception InferException

fromJson :: A.Value -> T
fromJson (A.Object kvs) =
  object (M.mapKeys K.toString (M.map fromJson (KM.toMap kvs)))
fromJson t@(A.Array vs) =
  let ts = map fromJson (V.toList vs)
   in case nub ts of
        [] -> list Nothing
        [t] -> list (Just t)
        _ -> throw (InferException ts)
fromJson (A.String _) = string
fromJson (A.Number _) = number
fromJson (A.Bool _) = bool
fromJson A.Null = null

object1 =
  [aesonQQ|{
    "firstName": "firstName",
    "lastName": "lastName"
  }|]

object2 =
  [aesonQQ|{
    "firstName": "firstName",
    "lastName": "lastName",
    "birthDay": null
  }|]

object3 =
  [aesonQQ|{
    "firstName": "firstName",
    "lastName": "lastName",
    "birthDay": "1990-01-01"
  }|]

object4 =
  [aesonQQ|{
  "firstName": "firstName"
  }|]

object5 =
  [aesonQQ|{
  "lastName": 42,
  "birthDay": null
  }|]

main =
  -- fromJson object2 =:= fromJson object3
  -- unify (fromJson object1) (fromJson object2)
  putStrLn
    ( intercalate
        "\n\n"
        ( map
            toString
            ( foldr1
                (\ls rs -> (concat [unify1 l r | l <- ls, r <- rs]))
                ( map
                    ((: []) . fromJson)
                    [ object1,
                      object2,
                      -- object3,
                      -- object4,
                      object5
                    ]
                )
            )
        )
    )

-- >>= unify (fromJson object2)
-- >>= unify (fromJson object4)