import Data.Char as Char
import Control.Monad
import qualified Table
import EnvSimple
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 Type =
    TInt
  | TName Symbol [Type]
  | TFun Type Type
  | TVar (Table.T Symbol [Type])
  | TRec (Table.T Symbol Type)
  deriving Eq

data Expr =
    EId Symbol
  | ENum Integer
  | EPlus  Expr Expr
  | ETimes Expr Expr
  | EFun Symbol (Maybe Type) Expr
  | EApp Expr Expr
  | ELet Symbol (Maybe Type) Expr Expr
  | ERecord [(Symbol, Expr)]
  | EMember Expr Symbol
  | ETypeVar Symbol [Symbol] [(Symbol, [Type])] Expr
  | ETypeRec Symbol [Symbol] [(Symbol, Type)]   Expr
  | ETypeDef Symbol [Symbol] Type 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

-- memory

type Location    = Int
type Environment = Env 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 []

-- evaluation

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 of 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,args) -> bind_value st_env (c, VCtor c (length args) []))
                          (st, env)
                          ctors
eval st env (ETypeRec _ _ _ e) = eval st env e
eval st env (ETypeDef _ _ _ e) = eval st env e
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"



-- data structures for type information

data TypeScheme   = TS [Symbol] Type
type TypeDefs     = Table.T Symbol TypeScheme
type Declarations = Table.T Symbol Type
type TypeSubst    = Table.T Symbol Type

data Types = TY TypeDefs TypeSubst Int
  -- type definitions, variable values, number of used variables

new_type_info defs = TY defs Table.empty 0

lookup_type :: Types -> Symbol -> TypeScheme
lookup_type (TY defs _ _) name = Table.lookup defs name

maybe_lookup_type :: Types -> Symbol -> Maybe TypeScheme
maybe_lookup_type (TY defs _ _) name = Table.maybe_lookup defs name

define_type :: Types -> Symbol -> TypeScheme -> Types
define_type (TY defs vars n) name typ = TY (Table.bind defs name typ) vars n

lookup_type_variable :: Types -> Symbol -> Type
lookup_type_variable (TY _ vars _) name = iter (Table.lookup vars name)
  where
    iter (TName n []) = case Table.maybe_lookup vars n of
                          Just t  -> iter t
                          Nothing -> TName n []
    iter typ = typ

maybe_lookup_type_variable :: Types -> Symbol -> Maybe Type
maybe_lookup_type_variable (TY _ vars _) name = iter (Table.maybe_lookup vars name)
  where
    iter Nothing             = Nothing
    iter (Just (TName n [])) = case Table.maybe_lookup vars n of
                                 Just t  -> iter (Just t)
                                 Nothing -> Just (TName n [])
    iter (Just typ) = Just typ

is_variable :: Types -> Symbol -> Bool
is_variable (TY defs _ _) n = case Table.maybe_lookup defs n of
                                Just _  -> False
                                Nothing -> True

set_variable :: Types -> Symbol -> Type -> Types
set_variable (TY defs vars n) name typ = TY defs (Table.bind vars name typ) n

new_variable :: Types -> (Symbol, Types)
new_variable (TY defs vars n) = ("a[" ++ show n ++ "]", TY defs vars (n+1))


-- utility functions for types

uncurry_type :: Type -> ([Type], Type)
uncurry_type (TFun a b) = let (arg_types, ret_type) = uncurry_type b in
                          (a:arg_types, ret_type)
uncurry_type a          = ([], a)


substitute :: Table.T Symbol Type -> Type -> Type
substitute subst (TName n []) =
  case Table.maybe_lookup subst n of
    Nothing -> TName n []
    Just t  -> t

substitute subst (TName n args) =
  TName n (map (substitute subst) args)

substitute _ TInt = TInt

substitute subst (TFun a b) =
  TFun (substitute subst a) (substitute subst b)

substitute subst (TVar vars) =
  TVar
    (Table.map
      (\_ args -> map (substitute subst) args)
      vars)

substitute subst (TRec fields) =
  TRec
    (Table.map
      (\ _ t -> substitute subst t)
      fields)

instantiate :: Types -> [Symbol] -> Type -> (Type, Types)
instantiate types params typ =
  let (subst, new_types) = foldl
                             (\(subst, types) p ->
                               let (var, new_types) = new_variable types in
                               (Table.bind subst p (TName var []), new_types))
                             (Table.empty, types)
                             params in
  (substitute subst typ, new_types)

occurs :: Symbol -> Type -> Bool
occurs x TInt           = False
occurs x (TName n [])   = n == x
occurs x (TName _ args) = any (occurs x) args
occurs x (TFun a b)     = occurs x a || occurs x b
occurs x (TVar vars)    =
  Table.fold
    (\_ args occ -> occ || any (occurs x) args)
    False
    vars
occurs x (TRec fields)  = Table.fold (\_ t occ -> occ || occurs x t) False fields;


extract_parameters :: Types -> Type -> [Symbol]

extract_parameters types TInt = []

extract_parameters types (TName n []) =
  case maybe_lookup_type types n of
    Nothing -> [n]
    Just _  -> []

extract_parameters types (TName n args) =
  foldl (++) [] (map (extract_parameters types) args)

extract_parameters types (TFun a b) =
  extract_parameters types a ++ extract_parameters types b

extract_parameters types (TVar vars) =
  Table.fold
    (\ _ args params -> foldl (++) params (map (extract_parameters types) args))
    []
    vars

extract_parameters types (TRec fields) =
  Table.fold
    (\_ t params -> params ++ extract_parameters types t)
    []
    fields


expand_type_variable :: Types -> Type -> Type

expand_type_variable types (TName n []) =
  case maybe_lookup_type_variable types n of
    Nothing -> TName n []
    Just _  -> expand_type_variable types (lookup_type_variable types n)

expand_type_variable _ typ = typ


expand_type :: Types -> Type -> Type

expand_type types TInt = TInt

expand_type types (TName n []) =
  case maybe_lookup_type_variable types n of
    Nothing -> TName n []
    Just _  -> expand_type types (lookup_type_variable types n)

expand_type types (TName n args) =
  TName n (map (expand_type types) args)

expand_type types (TFun a b) =
  TFun (expand_type types a) (expand_type types b)

expand_type types (TVar vars) =
  TVar
    (Table.map
      (\ _ args -> map (expand_type types) args)
      vars)

expand_type types (TRec fields) =
  TRec
    (Table.map
      (\_ t -> expand_type types t)
      fields)


-- type unification


unify :: Types -> Type -> Type -> Types
unify types t1 t2 =
--   trace ("unifying types " ++ show t1 ++ " [" ++ pp_type types t1 ++ "] and " ++ show t2
--                            ++ " [" ++ pp_type types t2 ++ "]") $
   unify_step types (expand_type_variable types t1) (expand_type_variable types t2)

unify_step types (TName n []) t2 | is_variable types n =
  case maybe_lookup_type_variable types n of
    Just t  -> unify types t t2
    Nothing -> case t2 of
                 TName n2 [] -> if n2 == n then
                                  types
                                else
                                  if occurs n t2 then
                                    error "recursive type"
                                  else
                                    set_variable types n t2
                 _           -> if occurs n t2 then
                                  error "recursive type"
                                else
                                  set_variable types n t2

unify_step types t1 (TName n []) | is_variable types n =
  case maybe_lookup_type_variable types n of
    Nothing -> set_variable types n t1 -- XXX occurs check
    Just t  -> unify types t1 t

unify_step types (TName n1 args1) (TName n2 args2) =
  if n1 == n2 then
    unify_lists types args1 args2
  else
    error ("type error: trying to unify nominal types " ++ n1 ++ " and " ++ n2)

unify_step types TInt TInt = types

unify_step types (TFun a1 b1) (TFun a2 b2) =
  unify (unify types a1 a2) b1 b2

unify_step types (TVar vars1) (TVar vars2) =
  if Table.same_keys vars1 vars2 then
    Table.fold2
      (\_ ct1 ct2 types -> unify_lists types ct1 ct2)
      types
      vars1 vars2
  else
    error "type error"

unify_step types (TRec f1) (TRec f2) =
  if Table.same_keys f1 f2 then
    Table.fold2
      (\_ t1 t2 types -> unify types t1 t2)
      types
      f1 f2
  else
    error "type error"

unify_step _ t1 t2 = error ("type error: trying to unify " ++ show t1 ++ " and " ++ show t2)

unify_lists types lst1 lst2 =
  if length lst1 == length lst2 then
    foldl
      (\types (t1,t2) -> unify types t1 t2)
      types
      (zip lst1 lst2)
  else
    error "type error"



-- type checking

type_check types decl expr =
  type_check_step types decl expr
--  trace ("## type checking " ++ show expr ++ "\n") $
--  let (typ, new_types) = type_check_step types decl expr in
--  (trace ("## type of " ++ show expr ++ " is " ++ show (expand_type types typ)
--                        ++ " [" ++ pp_type types typ ++ "]") $
--  (typ, new_types))

type_check :: Types -> TypeDefs -> Expr -> (Type, Types)

type_check_step types decl (ENum n) = (TInt, types)

type_check_step types decl (EId id) =
  let TS params typ = Table.lookup decl id in
  instantiate types params typ

type_check_step types decl (EFun x t body) =
  let (x_type, new_types) = case t of
                              Just t  -> (t, types)
                              Nothing -> let (n, new_types) = new_variable types in
                                         (TName n [], new_types)
                            in
  let (typ, new_types2)   = type_check new_types (Table.bind decl x (TS [] x_type)) body in
  (TFun x_type typ, new_types2)

type_check_step types decl (EPlus e1 e2)   =
  let (et1, new_types1) = type_check types      decl e1 in
  let (et2, new_types2) = type_check new_types1 decl e2 in
  let new_types3        = unify new_types2 TInt et1 in
  let new_types4        = unify new_types3 TInt et2 in
  (TInt, new_types4)

type_check_step types decl (ETimes e1 e2) =
  let (et1, new_types1) = type_check types      decl e1 in
  let (et2, new_types2) = type_check new_types1 decl e2 in
  let new_types3        = unify new_types2 TInt et1 in
  let new_types4        = unify new_types3 TInt et2 in
  (TInt, new_types4)

type_check_step types decl (EApp f a) =
  let (ft, new_types)   = type_check types     decl f in
  let (at, new_types2)  = type_check new_types decl a in
  let (var, new_types3) = new_variable new_types2     in
  let rt                = TName var []                in
  (rt, unify new_types3 ft (TFun at rt))

type_check_step types decl (ELet x t e b) =
  let (x_type, new_types) = case t of
                              Just t  -> (expand_type types t, types)
                              Nothing -> let (n, new_types) = new_variable types in
                                         (TName n [], new_types)
                            in
  let params           = extract_parameters new_types x_type  in
  let new_decl         = Table.bind decl x (TS params x_type) in
  let (et, new_types2) = type_check new_types new_decl e in
  let new_types3       = unify new_types2 x_type et in
  let x_type2          = expand_type new_types3 x_type in
  let params2          = extract_parameters new_types3 x_type2  in
  let new_decl2        = Table.bind decl x (TS params2 x_type2) in
  type_check new_types3 new_decl2 b

type_check_step types decl (ERecord fields) =
  let (field_types, new_types) = foldl
                                   (\(table, types) (f,e) ->
                                     let (t, new_types) = type_check types decl e in
                                     (Table.bind table f t, new_types))
                                   (Table.empty, types)
                                   fields in
  (TRec field_types, new_types)

type_check_step types decl (EMember e f) =
  let (e_type, new_types) = type_check types decl e in
  case expand_type_variable types e_type of
    TRec fields -> (Table.lookup fields f, new_types)
    TName _ []  -> error "type error: type inference for record types is not supported. Please use an explicit type annotation."
    _           -> error "type error"

type_check_step types decl (ETypeVar name params variants e) =
  let ctors     = foldl
                    (\table (c,args) -> Table.bind table c args)
                    Table.empty
                    variants in
  let typ       = TVar ctors in
  let new_types = define_type types name (TS params typ) in
  let new_decl  = foldl
                    (\decl (c,args) -> Table.bind decl c
                                         (TS params (foldr (\t a -> TFun a t) typ args)))
                    Table.empty
                    variants in
  type_check new_types new_decl e

type_check_step types decl (ETypeRec name params fields e) =
  let field_types = foldl
                     (\table (f,t) -> Table.bind table f t)
                     Table.empty
                     fields in
  let new_types = define_type types name (TS params (TRec field_types)) in
  type_check new_types decl e

type_check_step types decl (ETypeDef name params t e) =
  let new_types = define_type types name (TS params t) in
  type_check new_types decl e

type_check_step types decl (ECase e pats) =
  let (typ,       new_types)  = type_check types decl e in
  let (pat_types, new_types2) = foldl
                                  (\(pt,types) p ->
                                    let (ptyp, new_types) = check_pat types typ p in
                                    (ptyp : pt, new_types))
                                  ([], new_types)
                                  pats in
  case pat_types of
    []     -> error "malformed case statement"
    (p:ps) -> let new_types3 = foldl (\types pat -> unify types p pat) new_types2 ps in
              (p, new_types3)
  where
    check_pat types typ (g, b) =
      let (new_types, new_decl) = check_guard types typ g in
      type_check new_types new_decl b

    check_guard types typ ElsePat     = (types,                decl)
    check_guard types typ (NumPat  n) = (unify types typ TInt, decl)

    check_guard types typ (VarPat  x) = case Table.maybe_lookup decl x of
          -- check whether x is a constructor or a variable
          Nothing             -> (types, Table.bind decl x (TS [] typ))
          Just (TS params ct) -> let (ctor_typ, new_types) = instantiate types params ct in
                                 let (arg_typ,  ret_typ)   = uncurry_type ctor_typ       in
                                 case arg_typ of
                                   [] -> (unify new_types typ ret_typ, decl)
                                   _  -> error "type error"

    check_guard types typ (CtorPat c xs) =
      let TS params ct          = Table.lookup decl c         in
      let (ctor_typ, new_types) = instantiate types params ct in
      let (arg_typ,  ret_typ)   = uncurry_type ctor_typ       in
      let new_types2            = unify new_types typ ret_typ in
      if length xs == length arg_typ then
        let new_decl = foldl (\decl (x,t) -> Table.bind decl x (TS [] t))
                         decl
                         (zip xs arg_typ) in
        (new_types2, new_decl)
      else
        error "type error"

type_check_step types decl (EAssign x v) =
  let (t1, new_types1) = type_check types      decl x in
  let (t2, new_types2) = type_check new_types1 decl v in
  (type_unit, unify new_types2 t1 t2)

type_check_step types decl (ESeq e1 e2) =
  let (t1, new_types1) = type_check types      decl e1 in
  let (t2, new_types2) = type_check new_types1 decl e2 in
  (t2, unify new_types2 t1 type_unit)

type_check_step types decl (EPrint str e) =
  let (_, new_types) = type_check types decl e in
  (type_unit, new_types)

type_check_step types decl (EWhile c e) =
  let (t1, new_types1) = type_check types      decl c in
  let (t2, new_types2) = type_check new_types1 decl e in
  let new_types3       = unify new_types2 t1 type_bool in
  let new_types4       = unify new_types3 t2 type_unit in
  (type_unit, new_types4)

type_check_step types decl (EFor v l u e) =
  let (t1, new_types1) = type_check types      decl l  in
  let (t2, new_types2) = type_check new_types1 decl u  in
  let (t3, new_types3) = type_check new_types2 decl e  in
  let new_types4       = unify new_types3 t1 TInt      in
  let new_types5       = unify new_types4 t2 TInt      in
  let new_types6       = unify new_types5 t3 type_unit in
  (type_unit, new_types6)

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 "()" (Just (TName "unit" [])) (desugar b)
desugar (PFun [(p,t)]      b) = EFun p (liftM desugar_type t) (desugar b)
desugar (PFun ((p,t):args) b) = EFun p (liftM desugar_type t) (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 t e)         b) =
  ELet x (liftM desugar_type t) (desugar e) (desugar b)
desugar (PDecl (PLetFun f args t b) e) =
  ELet f (liftM desugar_type ft) (desugar (PFun args b)) (desugar e)
  where
    ft = foldr arr t (map snd args)
    arr Nothing _ = Nothing
    arr _ Nothing = Nothing
    arr (Just a) (Just b) = Just (PArrow a b)
desugar (PDecl (PTypeVar name params ctors)  e) =
  ETypeVar name params
    (map (\(c,args) -> (c, map desugar_type args)) ctors)
    (desugar e)
desugar (PDecl (PTypeRec name params fields) e) =
  ETypeRec name params
    (map (\(f,t) -> (f, desugar_type t)) fields)
    (desugar e)
desugar (PDecl (PTypeDef t params t2) e) =
  ETypeDef t params (desugar_type t2) (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

desugar_type :: PExpr -> Type
desugar_type (PId "int")         = TInt
desugar_type (PId n)             = TName n []
desugar_type (PApp (PId n) args) = TName n (map desugar_type args)
desugar_type (PArrow a b)        = TFun (desugar_type a) (desugar_type b)
desugar_type _                   = error "invalid type annotation"

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

show_list_acc :: String -> String -> (b -> a -> (String, b)) -> b -> [a] -> (String, b)
show_list_acc l r sh acc []     = ("", acc)
show_list_acc l r sh acc (x:xs) =
  let (str,new_acc) = foldl
                        (\(str,acc) y -> let (s,a) = sh acc y in
                                         (str ++ ", " ++ s, a))
                        (sh acc x)
                        xs in
  (l ++ str ++ r, new_acc)

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

pretty_print_type :: Table.T Symbol String -> Type -> (String, Table.T Symbol String)
pretty_print_type names TInt         = ("int", names)
pretty_print_type names (TName n []) = case n of
  ('a' : '[' : _) -> case Table.maybe_lookup names n of
                       Just str -> (str, names)
                       Nothing  -> let used = Table.size names in
                                   let str  = if used < 26 then
                                                [Char.chr (97 + used)]
                                              else if used < 26 * 26 then
                                                [Char.chr (97 + used `div` 26),
                                                 Char.chr (97 + used `mod` 26)]
                                              else
                                                [Char.chr (97 + used `div` 26),
                                                 Char.chr (97 + used `mod` 26)]
                                                ++ show (used `div` (26 * 26)) in
                                   (str, Table.bind names n str)
  _               -> (n, names)
pretty_print_type names (TName n args) =
  let (str, new_names) = show_list_acc "(" ")" pretty_print_type names args in
  (n ++ str, new_names)
pretty_print_type names (TFun a b) =
  let (str1, names1) = pretty_print_type names  a in
  let (str2, names2) = pretty_print_type names1 b in
  ("(" ++ str1 ++ " -> " ++ str2 ++ ")", names2)
pretty_print_type names (TVar vars) =
  foldl
    (\(str,names) (c,fields) ->
      let (str, new_names) = show_list_acc "(" ")" pretty_print_type names fields in
      (str ++ "| " ++ c, new_names))
    ("", names)
    (Table.to_list vars)
pretty_print_type names (TRec fields) =
  show_list_acc "[" "]"
    (\names (f,t) -> let (str, new_names) = pretty_print_type names t in
                     (f ++ " : " ++ str, new_names))
    names (Table.to_list fields)

instance Show Type where
  show TInt           = "int"
  show (TName n [])   = n
  show (TName n args) = n ++ show_list "(" ")" show args
  show (TFun a b)     = "(" ++ show a ++ " -> " ++ show b ++ ")"
  show (TVar vars)    = foldl
                          (\str (c,fields) ->
                            str ++ "| " ++ c ++ show_list "(" ")" show fields)
                          ""
                          (Table.to_list vars)
  show (TRec fields)  = show_list "[" "]" (\(f,t) -> f ++ " : " ++ show t)
                          (Table.to_list fields)

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 Nothing b)  = "fun (" ++ a ++ ") { " ++ show b ++ " }"
  show (EFun a (Just t) b) = "fun (" ++ a ++ " : " ++ show t ++ ") { " ++ 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 Nothing  b e) = "let " ++ x ++ " = " ++ show b ++ "; " ++ show e
  show (ELet x (Just t) b e) = "let " ++ x ++ " :" ++ show t ++ " = " ++ 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 n ps ctors e)  = foldl (++) ("type " ++ n ++ show_list "(" ")" id ps ++ " =")
                                 (map
                                   (\(c,args) -> " | " ++ c ++ show_list "(" ")" show args)
                                   ctors)
                               ++ "; " ++ show e
  show (ETypeRec n ps fields e) = "type " ++ n ++ show_list "(" ")" id ps ++ " = {"
                                ++ show_list "(" ")"
                                     (\(c,t) -> c ++ " : " ++ show t)
                                     fields
                                ++ "}; " ++ show e
  show (ETypeDef n ps t e)      = "type " ++ n ++ show_list "(" ")" id ps ++ " = "
                                ++ show t ++ "; " ++ 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

-- builtin types and functions

type_unit = TName "unit" []
type_bool = TName "bool" []
type_pair = TName "pair" [TName "a" [], TName "b" []]
type_list = TName "list" [TName "a" []]

type_unit_def = TVar Table.empty
type_bool_def = TVar (Table.bind (Table.bind Table.empty "True" []) "False" [])
type_pair_def = TVar (Table.bind Table.empty "pair" [TName "a" [], TName "b" []])
type_list_def = TVar (Table.bind (Table.bind Table.empty "Nil" [])
                                 "Cons" [TName "a" [], type_list])

builtin_types =
  foldl (\types (n,t) -> Table.bind types n t) Table.empty
    [("unit",    TS []         type_unit_def),
     ("bool",    TS []         type_bool_def),
     ("product", TS ["a", "b"] type_pair_def),
     ("list",    TS ["a"]      type_list_def)]
((initial_store, builtin_env), builtin_decls) =
  foldl (\(st_env, decls) (ctor, val, typ) ->
           (bind_value st_env (ctor, val), Table.bind decls ctor typ))
    ((empty_store, empty_env), Table.empty)
    [("()",     VCtor "()"    0 [], TS [] type_unit),
     ("True",   VCtor "True"  0 [], TS [] type_bool),
     ("False",  VCtor "False" 0 [], TS [] type_bool),
     ("Pair",   VCtor "Pair"  2 [], TS ["a", "b"] (TFun (TName "a" []) (TFun (TName "b" []) type_pair))),
     ("Nil",    VCtor "Nil"   0 [], TS ["a"] type_list),
     ("Cons",   VCtor "Cons"  2 [], TS ["a"] (TFun (TName "a" []) (TFun type_list type_list))),
     ("%True",  VCtor "True"  0 [], TS [] type_bool),
     ("%False", VCtor "False" 0 [], TS [] type_bool),
     ("%Nil",   VCtor "Nil"   0 [], TS ["a"] type_list),
     ("%Cons",  VCtor "Cons"  2 [], TS ["a"] (TFun (TName "a" []) (TFun type_list type_list)))]

run str = let (val, st) = eval initial_store builtin_env (desugar (parse str)) in
          (get_output st, show_value st val)

run_tc str =
  let (typ, types) = type_check
                       (new_type_info builtin_types)
                       builtin_decls
                       (desugar (parse str)) in
  expand_type types typ

main :: IO ()
main = repl run (Just (fst . pretty_print_type Table.empty . run_tc))