import qualified Table
import EnvCont hiding (T)
import qualified EnvCont(T)
import Parser
import Repl

type Symbol = String

data Value =
    VNum Integer
  | VUnd Location
  | VRef Location
  | VCtor Symbol Int [Location]
  | VRec (Table.T Symbol Location)
  | VFun Symbol Environment Expr
  | VMod Environment
  | VCnt ContVS

data Expr =
    EId Symbol
  | ENum Integer
  | EPlus  Expr Expr
  | ETimes Expr Expr
  | EFun Symbol Expr
  | EApp Expr Expr
  | EDecl Decl Expr
  | ERecord [(Symbol, Expr)]
  | EMember Expr Symbol
  | ECase Expr [(Pattern, Expr)]
  | EAssign Expr Expr
  | ESeq Expr Expr
  | EPrint String Expr
  | EUnify Expr Expr
  | EWhile Expr Expr
  | EFor Symbol Expr Expr Expr
  | EReturn Expr
  | EBreak
  | EContinue
  | ELetCC Symbol Expr
  | ETry Expr Symbol Expr
  | EThrow Expr
  | EFail
  | EChoose [Expr]
  deriving Eq

data Decl =
    DEmpty
  | DModule Symbol [Decl]
  | DImport Expr
  | DLet Symbol (Maybe Expr)
  | DTypeVar [(Symbol, Int)]
  deriving Eq

data Pattern = ElsePat | VarPat Symbol | CtorPat Symbol [Symbol] | NumPat Integer
  deriving Eq

type Location       = Int
type Continuation a = a -> [Char]
type CPS a          = Continuation a -> [Char]
type ContVS         = Continuation (Value, Store)
type Environment    = EnvCont.T Location (Value, Store)

instance Eq Value where
  (VNum n1)        == (VNum n2)        = n1 == n2
  (VCtor c1 n1 a1) == (VCtor c2 n2 a2) = c1 == c2 && n1 == n2 && all (\(x,y) -> x == y) (zip a1 a2)
  (VRec f1)        == (VRec f2)        = f1 == f2
  (VFun x1 e1 b1)  == (VFun x2 e2 b2)  = x1 == x2 && e1 == e2 && b1 == b2 -- XXX we should at least use alpha-equivalence here
  (VMod d1)        == (VMod d2)        = d1 == d2
  (VCnt _)         == (VCnt _)         = True
  _                == _                = False

data Store = S Int (Table.T Location Value) String [ContVS] [(Continuation Store, [(Location, Value)])]

empty_store :: Store
empty_store = S 0 Table.empty "" [\_ -> error "uncaught exception"] [((\_ -> error "no checkpoint"), [])]

read_mem :: Store -> Location -> Value
read_mem (S _ t _ _ _) l = Table.lookup t l

set_mem :: Store -> Location -> Value -> Store
set_mem (S n t out es cp) l v =
  record_modification (S n (Table.bind t l v) out es cp) l (Table.lookup t l)

unset_mem :: Store -> (Location, Value) -> Store
unset_mem (S n t out es cp) (l, v) = S n (Table.bind t l v) out es cp

get_value :: Store -> Value -> Value
get_value st (VUnd _) = error "reading uninitialised variable"
get_value st (VRef l) = get_value st (read_mem st l)
get_value st v        = v

new_loc :: Store -> (Location, Store)
new_loc (S n t out es cp) = (n, S (n+1) (Table.bind t n (VUnd n)) out es cp)

add_output :: Store -> String -> Store
add_output (S n t out es cp) str = S n t (out ++ str) es cp

get_output :: Store -> String
get_output (S _ _ out _ _) = out

push_exception :: Store -> ContVS -> Store
push_exception (S n t out es cp) k = S n t out (k:es) cp

raise_exception :: Store -> ContVS
raise_exception (S n t out (k:es) cp) (v,_) = k (v, S n t out es cp)

save_checkpoint :: Store -> Continuation Store -> Store
save_checkpoint (S n t out es cp) k = S n t out es ((k,[]):cp)

restore_checkpoint :: Store -> (Store, Continuation Store)
restore_checkpoint (S n t out es ((k,vars):cp)) =
  (foldl unset_mem (S n t out es cp) vars, k)

record_modification :: Store -> Location -> Value -> Store
record_modification (S n t out es ((k,vars):cp)) l v =
  S n t out es ((k, (l,v):vars):cp)

abort :: Store -> [Char]
abort st = let (st2, cnt) = restore_checkpoint st in
           cnt st2

store_value :: Store -> Value -> (Store, Location)
store_value st v =
  let (l, st2) = new_loc st in
  case v of
    VUnd p -> (set_mem st2 l (VRef p), l)
    _      -> (set_mem st2 l v,        l)

bind_value :: (Store, Environment) -> (Symbol, Value) -> (Store, Environment)
bind_value (st, env) (x, v) =
  let (st2, l) = store_value st v in
  (st2, bind_variable env x l)

unit_value = VCtor "()" 0 []

unify :: Store -> Value -> Value -> Maybe Store
unify st v1 v2 = do
  (st,_) <- unify_rec Table.empty st v1 v2
  return st

unify_rec :: Table.T (Value,Value) () -> Store -> Value -> Value -> Maybe (Store, Table.T (Value,Value) ())
unify_rec seen st v1 v2 =
  case Table.maybe_lookup seen (v1, v2) of
    Nothing -> compare (Table.bind seen (v1,v2) ()) st v1 v2
    Just () -> Just (st, seen)
  where
    compare seen st (VRef l) v = compare seen st (read_mem st l) v
    compare seen st v (VRef l) = compare seen st v (read_mem st l)
    compare seen st (VUnd l1) (VUnd l2) =
      if l1 == l2 then
        Just (st, seen)
      else
        Just (set_mem st l1 (VRef l2), seen)
    compare seen st (VUnd l)  v = Just (set_mem st l v, seen)
    compare seen st v  (VUnd l) = Just (set_mem st l v, seen)
    compare seen st (VNum n1) (VNum n2)
      | n1 == n2 = Just (st, seen)
    compare seen st (VCtor c1 n1 args1) (VCtor c2 n2 args2)
      | c1 == c2 && n1 == n2 =
        foldl
          (\mst (a1,a2) -> do
            (st, seen) <- mst
            unify_rec seen st (read_mem st a1) (read_mem st a2))
          (Just (st, seen)) (zip args1 args2)
    compare seen st (VRec fields1) (VRec fields2)
      | Table.same_keys fields1 fields2 =
        Table.fold2
          (\_ f1 f2 mst -> do
            (st, seen) <- mst
            unify_rec seen st (read_mem st f1) (read_mem st f2))
          (Just (st, seen)) fields1 fields2
    compare _ _ u v = Nothing

eval :: Store -> Environment -> Expr -> CPS (Value, Store)
eval st env (ENum n)        k = k (VNum n, st)
eval st env (EId id)        k = k (read_mem st (lookup_variable env id), st)
eval st env (EFun arg body) k = k (VFun arg env body, st)
eval st env (EPlus e1 e2)   k =
  eval st env e1
    (\(v1, st1) -> eval st1 env e2
      (\(v2, st2) -> case (get_value st2 v1, get_value st2 v2) of
        (VNum x1, VNum x2) -> k (VNum (x1 + x2), st2)
        _                  -> error "addition of non-numbers"))
eval st env (ETimes e1 e2) k =
  eval st env e1
    (\(v1, st1) -> eval st1 env e2
      (\(v2, st2) -> case (get_value st2 v1, get_value st2 v2) of
        (VNum x1, VNum x2) -> k (VNum (x1 * x2), st2)
        _                  -> error "multiplication of non-numbers"))
eval st env (EApp f a) k =
--  trace ("### " ++ show (EApp f a)) $
  eval st env f
    (\(fv, st1) -> eval st1 env a
      (\(av, st2) -> case get_value st2 fv of
        VFun p e b         -> let (st3, new_env) = bind_value (st2, e) (p, av) in
                              eval st3 (set_return_cont new_env k) b k
        VCnt cnt           -> cnt (av, st2)
        VCtor ctor 0  args -> error "constructor applied to too many arguments"
        VCtor ctor ar args -> let (st3, l) = store_value st2 av in
                              k (VCtor ctor (ar-1) (args ++ [l]), st3)
        x                  -> error ("application to non-function" ++ show x)))
eval st env (EDecl d e) k =
  eval_decl st env d
    (\(st2, env2) -> eval st2 env2 e k)
eval st env (ERecord fields) k =
  let add_value f e k (new_state, vals) =
        eval new_state env e
          (\(v, st2) ->
              let (st3, l) = store_value st2 v in
              k (st3, Table.bind vals f l)) in
  (foldr
    (\(f,e) k -> add_value f e k)
    (\(new_state, vals) -> k (VRec vals, new_state))
    fields)
  (st, Table.empty)
eval st env (EMember e f) k =
  eval st env e
    (\(v, st2) -> case get_value st2 v of
      VRec fields -> k (read_mem st2 (Table.lookup    fields f), st2)
      VMod fields -> k (read_mem st2 (lookup_variable fields f), st2)
      _           -> error "value has no fields")
eval st env (ECase e pats) k =
    eval st env e
      (\(val, st2) -> iter val st2 pats)
    where
      iter val st2 []                 = error ("no matching pattern in case expression for value " ++ show val)
      iter val st2 ((ElsePat,  e):ps) = eval st2 env e k
      iter val st2 ((VarPat x, e):ps) = case maybe_lookup_variable env x of
          -- check whether x is a constructor or a variable
          Nothing -> let (st3, new_env) = bind_value (st2, env) (x, val) in
                     eval st3 new_env e k
          Just l  -> case get_value st2 (read_mem st2 l) of
              (VCtor c 0 []) -> case val of
                                  VCtor d 0 [] | c == d -> eval st2 env e k
                                  _                     -> iter val st2 ps
      iter val st2 ((NumPat n, e):ps) = case val of
          VNum m | n == m -> eval st2 env e k
          _               -> iter val st2 ps
      iter val st2 ((CtorPat c xs,e):ps) = case val of
          VCtor d 0 args | c == d && length xs == length args ->
            eval st2 new_env e k
            where
              new_env = foldl (\env (x,l) -> bind_variable env x l)
                          env (zip xs args)
          _ -> iter val st2 ps
eval st env (EAssign x e) k =
  eval_lvalue st env x
    (\(l, st1) ->
      eval st1 env e
        (\(v, st2) -> k (unit_value, set_mem st2 l v)))
eval st env (ESeq e1 e2) k =
  eval st  env e1
    (\(_, st1) -> eval st1 env e2 k)
eval st env (EPrint str e) k =
  eval st env e
    (\(v, st1) ->
      k (unit_value, add_output st1 ("\n" ++ str ++ " " ++ show_value st1 v)))
eval st env (EUnify x y) k =
  eval st env x
    (\(u, st1) ->
      eval st1 env y
        (\(v, st2) -> case unify st2 u v of
                        Nothing  -> abort st
                        Just st3 -> k (unit_value, st3)))
eval st env (EWhile c b) k =
  iter st
  where
    loop_env = set_break_cont env k
    iter st = eval st env c
                (\(v,st2) -> case get_value st2 v of
                  VCtor "False" 0 [] -> k (unit_value, st2)
                  VCtor "True"  0 [] -> let cc (_, st3) = iter st3 in
                                        eval st2 (set_continue_cont loop_env cc) b cc
                  _                  -> error "type error in while loop")
eval st env (EFor x l u b) k =
  eval st env l
    (\(lv, st2) -> eval st2 env u
      (\(uv, st3) ->
        let first = case get_value st3 lv of
                      VNum n -> n
                      _      -> error "type error in for loop" in
        let last  = case get_value st3 uv of
                      VNum n -> n
                      _      -> error "type error in for loop" in
        iter st3 first last))
  where
    loop_env = set_break_cont env k
    iter st i last =
      if i > last then
        k (unit_value, st)
      else
        let (st2, new_env) = bind_value (st, loop_env) (x, VNum i) in
        let cc (_, st3) = iter st3 (i + 1) last in
        eval st2 (set_continue_cont new_env cc) b cc
eval st env (EReturn e) k =
  eval st env e (get_return_cont env)
eval st env EBreak k =
  (get_break_cont env) (unit_value, st)
eval st env EContinue k =
  (get_continue_cont env) (unit_value, st)
eval st env (ELetCC x e) k =
  let (st2, new_env) = bind_value (st, env) (x, VCnt k) in
  eval st2 new_env e k
eval st env (ETry e x h) k =
  let ec (v,st2) = let (st3, new_env) = bind_value (st2, env) (x, v) in
                   eval st3 new_env h k in
  eval (push_exception st ec) env e k
eval st env (EThrow e) k =
  eval st env e (raise_exception st)
eval st env EFail k =
  let (st2, cnt) = restore_checkpoint st in
  cnt st2
eval st env (EChoose cases) k = case cases of
  []     -> eval st env EFail k
  (e:es) -> let st2 = save_checkpoint st (\st -> eval st env (EChoose es) k) in
            eval st2 env e k

eval_lvalue :: Store -> Environment -> Expr -> CPS (Location, Store)
eval_lvalue st env (EId id)      k = k (lookup_variable env id, st)
eval_lvalue st env (EMember e f) k =
  eval st env e
    (\(v, st2) -> case v of
      VRec fields -> k (Table.lookup fields f, st2)
      VMod fields -> k (lookup_variable fields f, st2)
      _           -> error "value has no fields")
eval_lvalue st env (EDecl d e) k =
  eval_decl st env d
    (\(st2, env2) -> eval_lvalue st2 env2 e k)
eval_lvalue st env (ESeq e1 e2) k =
  eval st env e1
    (\(_, st1) -> eval_lvalue st1 env e2 k)
eval_lvalue st env _ _ = error "expression is not an l-value"

eval_decl :: Store -> Environment -> Decl -> CPS (Store, Environment)
eval_decl st env (DLet x e) k =
  let (l, st2) = new_loc st         in --- note this is now non-recursive
  let new_env  = bind_variable env x l in
  case e of
    Nothing -> k (st2, new_env)
    Just e  -> eval st2 new_env e
                 (\(v, st3) -> k (set_mem st3 l v, new_env))
eval_decl st env (DTypeVar ctors) k =
  k (foldl (\st_env (c,a) -> bind_value st_env (c, VCtor c a []))
           (st, env)
           ctors)
eval_decl st env (DModule name decls) k =
  let add_decl d k (new_st, new_env) =
        eval_decl new_st (merge_declarations env (variable_table new_env)) d k in
  (foldr
    (\d k -> add_decl d k)
    (\(new_st, new_env) -> k (bind_value (new_st, env) (name, VMod new_env)))
    decls)
  (st, empty_env)
eval_decl st env (DImport e) k =
  eval st env e
    (\(v, st2) -> case v of
      VMod menv -> k (st2, merge_declarations env (variable_table menv))
      _         -> error "not a module in import statement")

desugar :: PExpr -> Expr
desugar (PId id)              = EId id
desugar (PNum n)              = ENum n
desugar (PPlus  e1 e2)        = EPlus  (desugar e1) (desugar e2)
desugar (PTimes e1 e2)        = ETimes (desugar e1) (desugar e2)
desugar (PMinus e1 e2)        = EPlus  (desugar e1) (ETimes (ENum (-1)) (desugar e2))
desugar (PEqual e1 e2)        = ECase (desugar (PMinus e1 e2))
                                  [ (NumPat 0, EId "%True"),
                                    (ElsePat,  EId "%False") ]
desugar (PFun []           b) = EFun "()" (desugar b)
desugar (PFun [(p,_)]      b) = EFun p (desugar b)
desugar (PFun ((p,_):args) b) = EFun p (desugar (PFun args b))
desugar (PApp f [])           = EApp (desugar f) (EId "()")
desugar (PApp f [a])          = EApp (desugar f) (desugar a)
desugar (PApp f args)         = foldl (\f a -> EApp f a ) (desugar f) (map desugar args)
desugar (PRecord fields)      = ERecord (map (\(f,e) -> (f, desugar e)) fields)
desugar (PMember e f)         = EMember (desugar e) f
desugar (PSeq e1 e2)          = ESeq (desugar e1) (desugar e2)
desugar (PAssign e f)         = EAssign (desugar e) (desugar f)
desugar (PUnify e1 e2)        = EUnify (desugar e1) (desugar e2)
desugar (PPrint str e)        = EPrint str (desugar e)
desugar (PDecl d e)           = EDecl (desugar_decl d) (desugar e)
desugar (PIf c t e)           = ECase (desugar c)
                                  [ (VarPat "%True", desugar t),
                                    (ElsePat,        desugar e) ]
desugar (PCase e pats)        = ECase (desugar e)
                                      (map (\(p,e) -> (desugar_pat p, desugar e)) pats)
desugar (PWhile c b)          = EWhile (desugar c) (desugar b)
desugar (PFor x l u b)        = EFor x (desugar l) (desugar u) (desugar b)
desugar (PReturn e)           = EReturn (desugar e)
desugar PBreak                = EBreak
desugar PContinue             = EContinue
desugar (PLetCC x e)          = ELetCC x (desugar e)
desugar (PTry e x h)          = ETry (desugar e) x (desugar h)
desugar (PThrow e)            = EThrow (desugar e)
desugar PFail                 = EFail
desugar (PChoose cases)       = EChoose (map desugar cases)
desugar (PListLit es t)       = foldr (\x y -> EApp (EApp (EId "%Cons") (desugar x)) y)
                                  (case t of
                                    Nothing -> EId "%Nil"
                                    Just e  -> desugar e)
                                  es
desugar _                     = error "unsupported syntactic construct\n"

desugar_pat :: PPattern -> Pattern
desugar_pat PElsePat        = ElsePat
desugar_pat (PVarPat x)     = VarPat x
desugar_pat (PCtorPat c xs) = CtorPat c xs
desugar_pat (PNumPat n)     = NumPat n

desugar_decl (PModule name [] decls) = DModule name (map desugar_decl decls)
desugar_decl (PImport e)             = DImport (desugar e)
desugar_decl (PLet x _ e)            = DLet x (Just (desugar e))
desugar_decl (PLetFun f args _ b)    = DLet f (Just (desugar (PFun args b)))
desugar_decl (PLetUnIni x _)         = DLet x Nothing
desugar_decl (PTypeVar _ _ ctors)    = DTypeVar
                                         (map (\(c,args) -> (c, length args)) ctors)
desugar_decl (PTypeRec _ _ _)        = DEmpty
desugar_decl (PTypeDef _ _ _)        = DEmpty

show_list :: Show a => String -> String -> (a -> String) -> [a] -> String
show_list l r sh [] = ""
show_list l r sh xs = l ++ foldl1 (\ x y -> x ++ ", " ++ y) (map sh xs) ++ r

instance Show Value where
  show (VUnd l)         = "<undefined " ++ show l ++ ">"
  show (VRef l)         = "<ref " ++ show l ++ ">"
  show (VNum n)         = show n
  show (VCtor c _ args) = c ++ show_list "(" ")" show args
  show (VRec fields)    = show_list "[" "]" (\(f,v) -> f ++ " = " ++ show v)
                            (Table.to_list fields)
  show (VFun a _ b)     = "fun (" ++ a ++ ") { " ++ show b ++ " }"
  show (VMod env)       = "module {" ++
                             show_list "" "" (\(f,v) -> "let " ++ f ++ " = " ++ show v ++ ";\n")
                               (Table.to_list (variable_table env))
                             ++ "}"
  show (VCnt _)         = "<cont>"

show_value st (VUnd l)         = "<undefined " ++ show l ++ ">"
show_value st (VRef l)         = show_value st (read_mem st l)
show_value st (VNum n)         = show n
show_value st (VCtor c _ args) = c ++ show_list "(" ")" (show_value st) (map (read_mem st) args)
show_value st (VRec fields)    = show_list "[" "]"
                                   (\(f,l) -> f ++ " = " ++ show_value st (read_mem st l))
                                   (Table.to_list fields)
show_value st (VFun a _ b)     = "fun (" ++ a ++ ") { " ++ show b ++ " }"
show_value st (VMod env)       = "module {" ++
                                   show_list "" ""
                                     (\(f,l) -> "let " ++ f ++ " = " ++ show_value st (read_mem st l) ++ ";\n")
                                     (Table.to_list (variable_table env))
                                   ++ "}"
show_value st (VCnt _)         = "<cont>"

instance Show Expr where
  show (EId x)        = x
  show (ENum n)       = show n
  show (EPlus  e1 e2) = "(" ++ show e1 ++ " + "  ++ show e2 ++ ")"
  show (ETimes e1 e2) = "(" ++ show e1 ++ " * "  ++ show e2 ++ ")"
  show (EFun a b)     = "fun (" ++ a ++ ") { " ++ show b ++ " }"
  show (EApp f a)     = show fn ++ show_list "(" ")" show args where
      (fn, args) = collect_args f [a]
      collect_args (EApp g b) args = collect_args g (b:args)
      collect_args g          args = (g, args)
  show (ERecord fields)    = show_list "[" "]" (\(f,e) -> f ++ " = " ++ show e) fields
  show (EMember e f)       = show e ++ "." ++ f
  show (ESeq a b)          = show a ++ "; " ++ show b
  show (EAssign e f)       = show e ++ " := " ++ show f
  show (EPrint str e)      = "print " ++ show str ++ " " ++ show e
  show (EDecl d e)         = show d ++ " " ++ show e
  show (ECase e pats)      = foldl (++) ("case " ++ show e)
                               (map (\(p,e) -> " | " ++ show p ++ " => " ++ show e) pats)
  show (EWhile c b)        = "while " ++ show c ++ " " ++ show b
  show (EFor x l u b)      = "for " ++ x ++ " = " ++ show l ++ " .. " ++ show u ++ " "
                               ++ show b
  show (EReturn e)         = "return " ++ show e
  show EBreak              = "break"
  show EContinue           = "continue"
  show (ELetCC k e)        = "letcc " ++ k ++ " => " ++ show e
  show (ETry e x h)        = "try " ++ show e ++ " catch " ++ x ++ " => " ++ show h
  show (EThrow e)          = "throw " ++ show e
  show (EChoose cases)     = foldl (++) "choose"
                               (map (\e -> " | " ++ show e) cases)
  show EFail               = "fail"

instance Show Decl where
  show DEmpty           = ";"
  show (DLet x b)       = "let " ++ x ++ " = " ++ show b ++ ";"
  show (DTypeVar ctors) = foldl (++) "type ?? ="
                            (map
                              (\(c,ar) -> " | " ++ c ++
                                            show_list "(" ")" id (replicate ar "_"))
                                ctors)
                            ++ ";"
  show (DModule name decls) = foldl (++) ("module " ++ name ++ " ") (map show decls)
                                ++ ";"
  show (DImport e)          = "import " ++ show e ++ ";"

instance Show Pattern where
  show ElsePat         = "else"
  show (VarPat x)      = x
  show (NumPat n)      = show n
  show (CtorPat c xs)  = c ++ show_list "(" ")" id xs

(initial_store, initial_env) =
  foldl bind_value
    (empty_store, empty_env)
    [("()",     VCtor "()"    0 []),
     ("True",   VCtor "True"  0 []),
     ("False",  VCtor "False" 0 []),
     ("Pair",   VCtor "Pair"  2 []),
     ("Nil",    VCtor "Nil"   0 []),
     ("Cons",   VCtor "Cons"  2 []),
     ("%True",  VCtor "True"  0 []),
     ("%False", VCtor "False" 0 []),
     ("%Nil",   VCtor "Nil"   0 []),
     ("%Cons",  VCtor "Cons"  2 [])]

run str = eval initial_store initial_env (desugar (parse str))
            (\(val, st) -> get_output st ++ "\n" ++ show_value st val)

main :: IO ()
main = repl run Nothing