{-# 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 qualified Data.Text as Text
import GHC.Generics (Generic)

data Term
  = Var Text
  | Lam Text Term
  | App Term Term
  | Num Int
  | Str Text
  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
  | VStr Text
  | Builtin Text Int [Value]
  deriving (Show, Read, Eq)

data Neutral = NVar Text | NNum Int | NStr Text | 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)
  , ("str-concat", 2)
  , ("str-length", 1)
  , ("str-slice",  3)
  , ("str=",       2)
  , ("str<",       2)
  ]

-- | 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
  Str s -> VStr s

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)
    VStr s   -> VNeut (NApp (NStr s) 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
  -- Integer arithmetic
  ("+",   [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)
  -- String operations
  ("str-concat", [VStr x, VStr y])           -> VStr (x <> y)
  ("str-length", [VStr x])                   -> VNum (Text.length x)
  ("str-slice",  [VNum i, VNum j, VStr x])   -> VStr (Text.take (j - i) (Text.drop i x))
  ("str=",       [VStr x, VStr y])           -> if x == y then churchTrue else churchFalse
  ("str<",       [VStr x, VStr y])           -> if x <  y then churchTrue else churchFalse
  -- 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 (NStr s)    -> Str s
      VNeut (NApp f a)  -> App (go names (VNeut f)) (go names a)
      VNum n            -> Num n
      VStr s            -> Str s
      -- 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