{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE LambdaCase #-}

module Minilang.Lambda.Eval where

import Data.Aeson (FromJSON, ToJSON)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.String (IsString (..))
import Data.Text (Text, pack)
import GHC.Generics (Generic)

data Term
  = Var Text
  | Lam Text Term
  | App Term Term
  | Num Int
  deriving (Show, Read, Eq, Generic, ToJSON, FromJSON)

instance IsString Term where
  fromString = Var . pack

-- | A built-in arithmetic operator value.
-- Carries its name, the number of arguments still needed, and the
-- already-accumulated argument values (in order of application).
data Value
  = Clos EnvValue Text Term
  | VNeut Neutral
  | VNum Int
  | Builtin Text Int [Value]
  deriving (Show, Read, Eq)

data Neutral = NVar Text | NNum Int | NApp Neutral Value
  deriving (Show, Read, Eq)

type Env = [(Text, Term)]

type EnvValue = [(Text, Value)]

newEnv :: Env
newEnv = []

-- | Arity table for every built-in operator.
builtinArity :: Map Text Int
builtinArity = Map.fromList
  [ ("+",   2)
  , ("-",   2)
  , ("*",   2)
  , ("div", 2)
  , ("mod", 2)
  , ("=",   2)
  , ("<",   2)
  , ("neg", 1)
  ]

-- | The set of built-in names, for fast membership testing.
builtinNames :: Map Text Value
builtinNames = Map.mapWithKey (\name ar -> Builtin name ar []) builtinArity

normalise :: Term -> Term
normalise term = readBack (eval builtinEnvValue term)

-- | An 'EnvValue' pre-seeded with all built-in operators.
builtinEnvValue :: EnvValue
builtinEnvValue = Map.toList builtinNames

-- call-by-value evaluator
-- from https://davidchristiansen.dk/tutorials/nbe/#%28part._untyped-eval%29
eval :: EnvValue -> Term -> Value
eval env = \case
  Var x -> case lookup x env of
    Just v  -> v
    Nothing -> VNeut (NVar x)
  Lam x body -> Clos env x body
  App f a ->
    let f' = eval env f
        a' = eval env a
     in f' `seq` a' `seq` apply f' a'
  Num n -> VNum n

apply :: Value -> Value -> Value
apply f a =
  case f of
    Clos closureEnv x body ->
      let env' = (x, a) : closureEnv
       in eval env' body
    VNeut f' -> VNeut (NApp f' a)
    VNum n   -> VNeut (NApp (NNum n) a)
    -- Last argument arrives: dispatch and reduce.
    Builtin name 1 args -> dispatchBuiltin name (args ++ [a])
    -- Still waiting for more arguments: accumulate.
    Builtin name ar args -> Builtin name (ar - 1) (args ++ [a])

-- | Church-encoded True: lam f t -> t
churchTrue :: Value
churchTrue = Clos builtinEnvValue "f" (Lam "t" (Var "t"))

-- | Church-encoded False: lam f t -> f
churchFalse :: Value
churchFalse = Clos builtinEnvValue "f" (Lam "t" (Var "f"))

-- | Reduce a fully-saturated built-in application.
-- If the arguments are not both integers (free variables in function
-- position), fall back to a stuck neutral term.
-- Comparison operators (= and <) return Church booleans rather than integers.
dispatchBuiltin :: Text -> [Value] -> Value
dispatchBuiltin name args = case (name, args) of
  ("+",   [VNum x, VNum y]) -> VNum (x + y)
  ("-",   [VNum x, VNum y]) -> VNum (x - y)
  ("*",   [VNum x, VNum y]) -> VNum (x * y)
  ("div", [VNum x, VNum y])
    | y /= 0   -> VNum (x `div` y)
  ("mod", [VNum x, VNum y])
    | y /= 0   -> VNum (x `mod` y)
  ("=",   [VNum x, VNum y]) -> if x == y then churchTrue else churchFalse
  ("<",   [VNum x, VNum y]) -> if x <  y then churchTrue else churchFalse
  ("neg", [VNum x])         -> VNum (negate x)
  -- Stuck: reconstruct a neutral application chain rooted at the builtin name.
  _ -> foldl (\acc v -> VNeut (NApp (neutralOf acc) v)) (VNeut (NVar name)) args
  where
    neutralOf (VNeut n) = n
    neutralOf v         = NApp (NVar name) v   -- shouldn't happen in practice

readBack :: Value -> Term
readBack = go []
  where
    go names = \case
      Clos env x body ->
        let name  = freshName x names
            names' = name : names
            val'  = eval ((x, VNeut (NVar name)) : env) body
         in Lam name (go names' val')
      VNeut (NVar n)    -> Var n
      VNeut (NNum n)    -> Num n
      VNeut (NApp f a)  -> App (go names (VNeut f)) (go names a)
      VNum n            -> Num n
      -- A partial builtin reads back as a chain of applications rooted at
      -- the builtin name, e.g. Builtin "+" 1 [VNum 3]  →  App (Var "+") (Num 3)
      Builtin bname _ accumulated ->
        foldl (\acc v -> App acc (go names v)) (Var bname) accumulated

freshName :: Text -> [Text] -> Text
freshName x names
  | x `elem` names = freshName (x <> "'") names
  | otherwise = x