module AST where

import qualified Data.Map as M
import qualified Data.Maybe as MB
import qualified Data.List as L

import Types

data Binding = NameBind | VarBind Ty deriving (Show)

type Context = [(String, Binding)]

addBinding :: Context -> String -> Binding -> Context
addBinding ctx x bind = (x, bind) : ctx

type ContextLength = Int

data PairProj = FST | SND deriving (Show, Eq)

getTypeFromContext :: Context -> Term -> Ty
getTypeFromContext ctx t = case t of
  TmVar i l -> case getBinding ctx i of
    VarBind t -> t
    _ -> error "Expected VarBind, got Namebind"
  TmFree s -> case L.lookup s ctx of
    Just (VarBind ty) -> ty
    Nothing -> error "Variable not bound to type"

getBinding :: Context -> Int -> Binding
getBinding ctx ix = snd $ ctx !! ix

data Term =
    TmVar Int ContextLength
  | TmFree String
  | TmAbs String Ty Term
  | TmApp Term Term
  | TmTrue
  | TmFalse
  | TmIf Term Term Term
  | TmPair Term Term
  | TmRecord (M.Map String Term)
  | TmDyn (M.Map String Ty)
  deriving (Show, Eq)

termShift :: Int -> Term -> Term
termShift shift term = walk 0 term
  where walk cutoff term = case term of
          TmVar x cl ->
            if x >= cutoff then TmVar (x + shift) (cl + shift)
            else TmVar x (cl + shift)
          TmAbs s ty t -> TmAbs s ty (walk (cutoff + 1) t)
          TmApp t1 t2 -> TmApp (walk cutoff t1) (walk cutoff t2)
          TmPair t1 t2 -> TmPair (walk cutoff t1) (walk cutoff t2)
          TmRecord fields -> TmRecord (M.map (walk cutoff) fields)
          _ -> term

termSubst :: Int -> Term -> Term -> Term
termSubst ix s term = walk 0 term
  where walk cutoff term = case term of
          TmVar x cl ->
            if x == (ix + cl) then termShift cl s else TmVar x cl
          TmAbs s ty t -> TmAbs s ty (walk (cutoff + 1) t)
          TmApp t1 t2 -> TmApp (walk cutoff t1) (walk cutoff t2)
          TmPair t1 t2 -> TmPair (walk cutoff t1) (walk cutoff t2)
          TmRecord fields -> TmRecord (M.map (walk cutoff) fields)
          _ -> term

termSubstTop :: Term -> Term -> Term
termSubstTop s t = termShift (-1) (termSubst 0 (termShift 1 s) t)

isval :: Context -> Term -> Bool
isval ctx t = case t of
  TmAbs {} -> True
  TmPair {} -> True
  TmRecord {} -> True
  _ -> False

dynAdd :: Term -> String -> Ty -> Term
dynAdd t v ty = case t of
  (TmDyn fields) -> TmDyn (M.insert v ty fields)
  _ -> error "Can only add to dynamic"

eval1 :: Context -> Term -> Maybe Term
eval1 ctx term = case term of
  TmApp (TmAbs s ty t) v ->
    if isval ctx v then Just $ termSubstTop v t
    else Just $ TmApp (TmAbs s ty t) v
  TmApp (TmPair l r) v ->
    if isval ctx v then Just $ termSubstTop (TmPair l r) v
    else eval1 ctx v
  TmApp v t ->
    if isval ctx v then Just $ TmApp v (MB.fromJust $ eval1 ctx t)
    else Just $ TmApp (MB.fromJust $ eval1 ctx v) t
  _ -> error "No rule applies"

eval :: Context -> Term -> Term
eval ctx t = case eval1 ctx t of
  Just t' -> eval ctx t'
  Nothing -> t

typeOf :: Context -> Term -> Ty
typeOf ctx term = case term of
  TmFree s -> getTypeFromContext ctx (TmFree s)
  TmVar i l -> getTypeFromContext ctx (TmVar i l)
  TmAbs s tyS t -> TyArr tyS tyT
    where
      ctx' = addBinding ctx s (VarBind tyS)
      tyT = typeOf ctx' t
  TmApp t1 t2 -> case tyT1 of
    TyArr ty1_1 ty1_2 ->
      if tyT2 == ty1_1 then ty1_2
      else error "Parameter type mismatch"
    _ -> error "Arrow type expected"
    where
      tyT1 = typeOf ctx t1
      tyT2 = typeOf ctx t2
  TmTrue -> TyBool
  TmFalse -> TyBool
  TmIf t1 t2 t3 ->
    if tyT1 == TyBool then
      if tyT2 == tyT3 then tyT2
        else error "Arms of conditional have different types"
      else error "Guard of conditional not a boolean"
    where
      tyT1 = typeOf ctx t1
      tyT2 = typeOf ctx t2
      tyT3 = typeOf ctx t3
  TmRecord fields -> TyRecord typeFields
    where
      typeFields = M.map (typeOf ctx) fields
  TmPair t1 t2 -> TyPair (typeOf ctx t1) (typeOf ctx t2)

subtype :: Context -> Ty -> Ty -> Bool
subtype ctx tyS tyT =
  tyS == tyT || case (tyS, tyT) of
    (_, TyTop) -> True
    (TyArr tyS1 tyS2, TyArr tyT1 tyT2) -> 
      subtype ctx tyT1 tyS1 && subtype ctx tyS2 tyT2
    (TyRecord fS, TyRecord fT) -> fieldMatch && fieldTypeMatch
      where
        -- No keys in fT that are not in fS
        fieldMatch = M.difference fT fS == M.empty
        -- For each key in fT, the type of its match must be the same
        fieldTypeMatch = and $ M.intersectionWith (==) fT fS
    --(TyDyn fS, TyRecord fT) -> 
    (_,_) -> False

