--- NOTE: IS NOT COMPLETED ---

import Data.Map.Strict as Map
import Data.Bifunctor as Bifunctor
import Data.Maybe as Maybe
import Control.Monad as Monad

type Symb = String
infixl 4 :@:

-------------------------------------------------------------------------

-- %x + 3 := 12
-- addp :: BoxT Int -> Int -> Int -> ???
-- addp :: Out Int -> In Int -> In Int
-- addp :: {x: Box Int} -> {_: Int} -> {_: Int} -> ??? {x: Int}
-- addp :: Box Int -> In Int -> In Int -> Ctx
-- addp :: Box x Int -> Int -> Int -> {x: Int}
-- addp (Box x Int) 3 <- type ??? {_: Int} -> {x: Int}
-- Match ((+') (Box x Int) 3) 12

type CtxT = Map Symb Type

-- TODO: proper types for all nodes
data Type = Boo
          | Type :-> Type
          | ContextT CtxT
          | BoxT Symb Type
  deriving (Read,Show,Eq)

type Ctx = Map Symb Term

data Term = Fls
          | Tru
          | OpenContext {ctx :: Term, term :: Term}
          | GetFromContext {x :: Symb, ctx :: Term}
          | Context Ctx
          | Var Symb
          | Box Symb Type
          | Term :@: Term
          | Match {pattern :: Term, doThen :: Term, doElse :: Term}
  deriving (Read,Show)

assign :: Term -> Term -> Maybe Ctx
assign Fls Fls = Just Map.empty
assign Tru Tru = Just Map.empty
-- assign (OpenContext ctx term) term = Just Map.empty -- TODO ?? no meaning, should have been evaluated
-- assign (GetFromContext x ctx) term = Just Map.empty -- TODO ?? no meaning, should have biie evaluated
-- assign (Context ctx) term = Just Map.empty -- TODO ?? can't assign term, no meaning ??
assign (Var x) (Var y) | x == y = Just Map.empty
assign (Box x t) term = Just $ Map.singleton x term
-- assign (left :@: right) (left' :@: right') = do ctxLeft <- assign left left' -- TODO: ?? probably should hav been evaluated first
--                                                 ctxRight <- assign right right'
--                                                 Just $ Map.union ctxLeft ctxRight 
assign match@(Match {}) term = case whnf (match :@: term) of
                                 Context ctx -> Just ctx
                                 _ -> Nothing
assign _ _ = Nothing

-- TODO: check
-- TODO: choose eval trategy
-- TODO: separate unsuccessful assignment and camputation error during assignment ??
oneStep :: Term -> Maybe Term
oneStep Fls = Nothing
oneStep Tru = Nothing
oneStep (OpenContext {ctx, term}) = error "not implemented yet" -- TODO
oneStep (GetFromContext {x, ctx}) | Just ctx' <- oneStep ctx = Just $ GetFromContext {x, ctx = ctx'}
                                  | Context c <- ctx = Map.lookup x c
oneStep (Context {}) = Nothing
oneStep (Var {}) = Nothing
oneStep (Box {}) = Nothing
oneStep (left :@: right) | Just left' <- oneStep left = Just $ left' :@: right
                         | Just right' <- oneStep right = Just $ left :@: right'
oneStep (Match {pattern, doThen, doElse}) | Just pattern' <- oneStep pattern = Just pattern'
                                          -- | Just doThen' <- oneStep doThen = Just doThen'
                                          -- | Just doElse' <- oneStep doElse = Just doElse'
oneStep (Match {pattern, doThen, doElse} :@: term) | Just c <- assign pattern term = Just $ OpenContext (Context c) doThen
                                                   | otherwise = Just doElse
oneStep _ = Nothing

whnf :: Term -> Term
whnf u = maybe u whnf (oneStep u)

-----------------------

-- TODO: all cases, check proper context usage
-- all functions should return context ??, need to fill variables and use returned expression in some way
infer :: CtxT -> Term -> Maybe Type
infer env Tru = Just Boo
infer env Fls = Just Boo
infer env (OpenContext {ctx, term}) | Just (ContextT ctxT) <- infer env ctx = infer (Map.union env ctxT) term -- TODO: better intersection, properly check duplicates
infer env (GetFromContext {x, ctx}) | Just (ContextT ctxT) <- infer env ctx = Map.lookup x ctxT
infer env (Context ctx) = Just $ ContextT $ Map.fromList $ Maybe.mapMaybe (\(x, term) -> (,) x <$> infer env term) $ Map.toList ctx -- TODO: real context ??
infer env (Var x) = Map.lookup x env
infer env (Box x t) = Just $ BoxT x t -- TODO ??
infer env (left :@: right) = error "TODO" -- let sould be match
infer env (Match {pattern, doThen, doElse}) = let env' = env in -- TODO: context from pattern, infer assignment
                                              let patternT = Boo in -- TODO: real type of pattern 
                                              do  thenT <- infer env' doThen
                                                  elseT <- infer env' doElse
                                                  Monad.guard $ thenT == elseT
                                                  Just $ patternT :-> thenT
infer _ _ = Nothing