{-# LANGUAGE LambdaCase #-}

module Minilang.Lambda.DeBruijn where

import Data.List (elemIndex)
import Data.Text (Text)
import Minilang.Lambda.Eval (Term)
import qualified Minilang.Lambda.Eval as Eval
import Debug.Trace (trace)

data DeBruijn ann
  = FV Text
  | Var ann Int
  | Lam ann (DeBruijn ann)
  | App (DeBruijn ann) (DeBruijn ann)
  deriving (Show, Read, Eq)

type Env ann = [(Text, DeBruijn ann)]

newEnv :: Env ann
newEnv = []

fromDeBruijn :: DeBruijn Text -> Term
fromDeBruijn = \case
  FV x -> Eval.Var x
  Var s _ -> Eval.Var s
  Lam s body -> Eval.Lam s (fromDeBruijn body)
  App f a -> Eval.App (fromDeBruijn f) (fromDeBruijn a)

-- Convert a term with named variables to De Bruijn indices
toDeBruijn :: Term -> DeBruijn Text
toDeBruijn term = go term []
  where
    go (Eval.Var x) env =
      case elemIndex x env of
        Just idx -> Var x idx
        Nothing -> FV x
    go (Eval.Lam x body) env =
      Lam x (go body (x : env))
    go (Eval.App f a) env =
      App (go f env) (go a env)

eval :: DeBruijn ann -> DeBruijn ann
eval (App f a) =
  case eval f of
    Lam _ body -> eval (substitute [] 0 a body)
    f' -> App f' a
eval term = term

eval' :: (Show ann, Eq ann) => Env ann -> DeBruijn ann -> DeBruijn ann
eval' env (FV v) = case lookup v env of
  Just term -> term
  Nothing -> FV v
eval' env (App f a) =
  case eval' env f of
    Lam _ body -> eval' env (substitute env 0 a body)
    f' -> if f' /= f
          then eval' env $ App f' a
          else App f' a
eval' _ term = term

substitute :: Env ann -> Int -> DeBruijn ann -> DeBruijn ann -> DeBruijn ann
substitute env n arg term = case term of
  FV x -> case lookup x env of
    Just term -> term
    Nothing -> FV x
  Var t x
    | x == n -> arg
    | otherwise -> Var t x
  Lam x body -> Lam x (substitute env (n + 1) arg body)
  App f a -> App (substitute env n arg f) (substitute env n arg a)

pretty :: DeBruijn Text -> Text
pretty = \case
  FV x -> x
  Var s _ -> s
  Lam s body -> "(lam (" <> s <> ") " <> pretty body <> ")"
  App f a -> "(" <> pretty f <> " " <> pretty a <> ")"