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

type Symbol = String

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

data Expr =
    EId Symbol
  | ENum Integer
  | EPlus  Expr Expr
  | ETimes Expr Expr
  | EFun Symbol Expr
  | EApp Expr Expr
  | ELet Symbol Expr Expr
  | ERecord [(Symbol, Expr)]
  | EMember Expr Symbol
  | ETypeVar [(Symbol, Int)] Expr
  | ECase Expr [(Pattern, Expr)]
  | EAssign Expr Expr
  | ESeq Expr Expr
  | EPrint String Expr
  | EWhile Expr Expr
  | EFor Symbol Expr Expr Expr
  deriving Eq

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

type Location    = Int
type Environment = EnvSimple.T Location
type Store       = (Int, Table.T Location Value, String)

empty_store :: Store
empty_store = (0, Table.empty, "")

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

set_mem :: Store -> Location -> Value -> Store
set_mem (n, t, out) l v = (n, Table.bind t l v, out)

new_loc :: Store -> (Location, Store)
new_loc (n, t, out) = (n, (n+1, t, out))

add_output :: Store -> String -> Store
add_output (n, t, out) str = (n, t, out ++ str)

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

store_value :: Store -> Value -> (Store, Location)
store_value st v =
  let (l, st2) = new_loc st in
  (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 []

eval :: Store -> Environment -> Expr -> (Value, Store)
eval st env (ENum n)        = (VNum n, st)
eval st env (EId id)        = (read_mem st (lookup_variable env id), st)
eval st env (EFun arg body) = (VFun arg env body, st)
eval st env (EPlus e1 e2)   =
  let (v1, st1) = eval st  env e1 in
  let (v2, st2) = eval st1 env e2 in
  case (v1, v2) of
    (VNum x1, VNum x2) -> (VNum (x1 + x2), st2)
    _                  -> error ("addition of non-numbers" ++ show v1 ++ show v2)
eval st env (ETimes e1 e2)   =
  let (v1, st1) = eval st  env e1 in
  let (v2, st2) = eval st1 env e2 in
  case (v1, v2) of
    (VNum x1, VNum x2) -> (VNum (x1 * x2), st2)
    _                  -> error "multiplication of non-numbers"
eval st env (EApp f a)      =
  let (fv, st1) = eval st  env f in
  let (av, st2) = eval st1 env a in
  case fv of
    VFun p e b         -> let (st3, new_env) = bind_value (st2, e) (p, av) in
                          eval st3 new_env b
    VCtor ctor 0  args -> error "constructor applied to too many arguments"
    VCtor ctor ar args -> let (st3, l) = store_value st2 av in
                          (VCtor ctor (ar-1) (args ++ [l]), st3)
    _                  -> error "application to non-function"
eval st env (ERecord fields) =
  let (new_state, vals) = foldl (\(st, vals) (f,e) -> let (v, st2) = eval st env e     in
                                                      let (st3, l) = store_value st2 v in
                                                      (st3, Table.bind vals f l))
                                (st, Table.empty) fields in
  (VRec vals, new_state)
eval st env (EMember e f) =
  case eval st env e of
    (VRec fields, st2) -> (read_mem st2 (Table.lookup fields f), st2)
    _                  -> error "not a record"
eval st env (ELet x e b) =
  let (l, st2) = new_loc st            in  -- note this is now non-recursive
  let new_env  = bind_variable env x l in
  let (v, st3) = eval st2 new_env e    in
  let st4      = set_mem st3 l v       in
  eval st4 new_env b
eval st env (ETypeVar ctors e) = eval new_st new_env e
  where
    (new_st, new_env) = foldl
                          (\st_env (c,ar) -> bind_value st_env (c, VCtor c ar []))
                          (st, env)
                          ctors
eval st env (ECase e pats) =
    iter pats
    where
      (val, st2) = eval st env e

      iter []                 = error "no matching pattern in case expression"
      iter ((ElsePat,  e):ps) = eval st2 env e
      iter ((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
          Just l  -> case read_mem st2 l of
              (VCtor c 0 []) -> case val of
                                  VCtor d 0 [] | c == d -> eval st2 env e
                                  _                     -> iter ps
      iter ((NumPat n, e):ps) = case val of
          VNum k | n == k -> eval st2 env e
          _               -> iter ps
      iter ((CtorPat c xs,e):ps) = case val of
          VCtor d 0 args | c == d && length xs == length args ->
            eval st2 new_env e
            where
              new_env = foldl (\env (x,l) -> bind_variable env x l)
                          env (zip xs args)
          _ -> iter ps
eval st env (EAssign x e) =
  let (l, st1) = eval_lvalue st  env x in
  let (v, st2) = eval        st1 env e in
  (unit_value, set_mem st2 l v)
eval st env (ESeq e1 e2) =
  let (_, st1) = eval st  env e1 in
  eval st1 env e2
eval st env (EPrint str e) =
  let (v, st1) = eval st env e in
  (unit_value, add_output st1 ("\n" ++ str ++ " " ++ show_value st1 v))
eval st env (EWhile c b)  =
  iter st
  where
    iter st = let (v,st2) = eval st env c in
              case v of
                VCtor "False" 0 [] -> (unit_value, st2)
                VCtor "True"  0 [] -> let (_, st3) = eval st2 env b in
                                      iter st3
                _                  -> error "type error in while loop"
eval st env (EFor x l u b)  = iter st first
  where
    (lv, st2) = eval st  env l
    (uv, st3) = eval st2 env u
    first = case lv of
              VNum n -> n
              _      -> error "type error in for loop"
    last  = case uv of
              VNum n -> n
              _      -> error "type error in for loop"
    iter :: Store -> Integer -> (Value, Store)
    iter st i = if i > last then
                  (unit_value, st)
                else
                  let (st2, new_env) = bind_value (st, env) (x, VNum i) in
                  let (v, st3) = eval st2 new_env b                     in
                  iter st3 (i + 1)

eval_lvalue :: Store -> Environment -> Expr -> (Location, Store)
eval_lvalue st env (EId id)      = (lookup_variable env id, st)
eval_lvalue st env (EMember e f) =
  case eval st env e of
    (VRec fields, st2) -> (Table.lookup fields f, st2)
    _                  -> error "not a record"
eval_lvalue st env (ELet x e b) =
  let (l, st2) = new_loc st            in  -- note this is now non-recursive
  let new_env  = bind_variable env x l in
  let (v, st3) = eval st2 new_env e    in
  let st4      = set_mem st3 l v       in
  eval_lvalue st4 new_env b
eval_lvalue st env (ESeq e1 e2) =
  let (_, st1) = eval st  env e1 in
  eval_lvalue st1 env e2
eval_lvalue st env (ECase e pats) =
    iter pats
    where
      (val, st2) = eval st env e

      iter []                 = error "no matching pattern in case expression"
      iter ((ElsePat,  e):ps) = eval_lvalue st2 env e
      iter ((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_lvalue st3 new_env e
          Just l  -> case read_mem st2 l of
              (VCtor c 0 []) -> case val of
                                  VCtor d 0 [] | c == d -> eval_lvalue st2 env e
                                  _                     -> iter ps
      iter ((NumPat n, e):ps) = case val of
          VNum k | n == k -> eval_lvalue st2 env e
          _               -> iter ps
      iter ((CtorPat c xs,e):ps) = case val of
          VCtor d 0 args | c == d && length xs == length args ->
            eval_lvalue st2 new_env e
            where
              new_env = foldl (\env (x,l) -> bind_variable env x l)
                          env (zip xs args)
          _ -> iter ps
eval_lvalue st env _ = error "expression is not an l-value"

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 (PPrint str e)        = EPrint str (desugar e)
desugar (PDecl (PLet x _ e)         b) = ELet x (desugar e) (desugar b)
desugar (PDecl (PLetFun f args _ b) e) = ELet f (desugar (PFun args b)) (desugar e)
desugar (PDecl (PTypeVar _ _ ctors) e) = ETypeVar
                                           (map (\(c,args) -> (c, length args)) ctors)
                                           (desugar e)
desugar (PDecl (PTypeRec _ _ _)     e) = desugar e
desugar (PDecl (PTypeDef _ _ _)     e) = 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 (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

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 (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_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 ++ " }"

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 (ELet x b e)        = "let " ++ x ++ " = " ++ show b ++ "; " ++ 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 (ETypeVar ctors e)  = foldl (++) "type ?? ="
                               (map
                                 (\(c,ar) -> " | " ++ c ++
                                               show_list "(" ")" id (replicate ar "_"))
                                 ctors)
                             ++ "; " ++ 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 = let (val, st) = eval initial_store initial_env (desugar (parse str)) in
          get_output st ++ "\n" ++ show_value st val

main :: IO ()
main = repl run Nothing