module AST where

import Data.List.Safe ((!!))
import Data.Text (Text)
import Prelude hiding ((!!))

type Name = Text

type Index = Int

data Term
  = Var Index
  | Pi Term Term
  | Lam Term
  | Term :$ Term
  | Ann Term Term
  | Star

data Val
  = VLam Term Env
  | VPi Val Term Env
  | VStar
  | Ne Ne

data Ne = NV Index | Ne :$$ Val

type Env = [Val]

type Ctx = [Val]

whnf :: Term -> Env -> Maybe Val
whnf (Var i) env = env !! i
whnf (Pi d r) env = do
  vd <- (whnf d env)
  Just $ VPi vd r env
whnf (Lam t) env = Just $ VLam t env
whnf (f :$ a) env = do
  vf <- whnf f env
  va <- whnf a env
  vf $$ va
whnf (Ann t _) env = whnf t env
whnf Star _ = Just VStar

($$) :: Val -> Val -> Maybe Val
(VLam r env) $$ a = whnf r (a : env)
Ne n $$ a = Just (Ne (n :$$ a))
_ $$ _ = Nothing

wk :: Val -> Val
wk (VLam t env) = VLam t (ewk env)
wk (VPi d r env) = VPi (wk d) r (ewk env)
wk VStar = VStar
wk (Ne ne) = Ne $ nwk ne

ewk :: Env -> Env
ewk = map wk

cwk :: Ctx -> Ctx
cwk = map wk

nwk :: Ne -> Ne
nwk (NV i) = NV (i + 1)
nwk (n :$$ v) = nwk n :$$ wk v

fresh :: Val
fresh = Ne (NV 0)

eq :: Ctx -> Val -> Val -> Val -> Maybe ()
eq _ VStar VStar VStar = return ()
eq ctx VStar (VPi d1 r1 env1) (VPi d2 r2 env2) = do
  eq ctx VStar d1 d2
  whnfr1 <- (whnf r1 ((Ne (NV 0)) : ewk env1))
  whnfr2 <- (whnf r2 ((Ne (NV 0)) : ewk env2))
  eq (fresh : cwk ctx) VStar whnfr1 whnfr2
eq ctx (VPi _ r env) f g = do
  vr <- whnf r ((Ne (NV 0)) : ewk env)
  vf <- wk f $$ (Ne (NV 0))
  vg <- wk g $$ (Ne (NV 0))
  eq (fresh : cwk ctx) vr vf vg
eq ctx _ (Ne n1) (Ne n2) = neq ctx n1 n2 >> return ()
eq _ _ _ _ = Nothing

neq :: Ctx -> Ne -> Ne -> Maybe Val
neq ctx (NV i) (NV j) = if i == j then ctx !! i else Nothing
neq ctx (n1 :$$ v1) (n2 :$$ v2) = do
  VPi d r env <- neq ctx n1 n2
  eq ctx d v1 v2
  whnf r (v1 : env)
neq _ _ _ = Nothing