{-# LANGUAGE LambdaCase #-}

module Minilang.Lambda.DeBruijn where

import Data.List (elemIndex)
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import qualified Data.Text as Text
import Minilang.Lambda.Eval (Term)
import qualified Minilang.Lambda.Eval as Eval

data DeBruijn ann
  = FV Text
  | FNum Int
  | FStr Text              -- ^ a string literal (closed value)
  | FBuiltin Text          -- ^ a fully-applied or partial built-in operator
  | 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 = []

-- | An environment pre-seeded with every built-in operator name.
-- Each entry maps the name to an 'FBuiltin' node so that 'eval\'' can
-- recognise and reduce applications of those names.
builtinDeBruijnEnv :: Env ann
builtinDeBruijnEnv =
  map (\name -> (name, FBuiltin name)) (Map.keys Eval.builtinArity)

fromDeBruijn :: DeBruijn Text -> Term
fromDeBruijn = \case
  FV x        -> Eval.Var x
  FNum n      -> Eval.Num n
  FStr s      -> Eval.Str s
  FBuiltin b  -> Eval.Var b        -- reads back as a Var; eval recognises it
  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)
    go (Eval.Num n) _env =
      FNum n
    go (Eval.Str s) _env =
      FStr s

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' :: Env Text -> DeBruijn Text -> DeBruijn Text
eval' env = \case
  FV v -> case lookup v env of
    Just term -> term
    Nothing   -> FV v
  FStr s     -> FStr s
  FBuiltin b -> FBuiltin b
  App f a ->
    let f' = eval' env f
        a' = eval' env a
    in case f' of
      -- Beta-reduction
      Lam _ body -> eval' env (substitute env 0 a' body)
      -- Built-in: collect arguments and reduce when saturated
      FBuiltin b -> evalBuiltin env b [a']
      -- Left-spine is itself a partial builtin application
      App g arg ->
        case collectBuiltin (App g arg) of
          Just (b, accumulated) -> evalBuiltin env b (accumulated ++ [a'])
          Nothing ->
            let applied = App f' a'
            in if applied /= App f a
               then eval' env applied
               else applied
      f'' ->
        let applied = App f'' a'
        in if applied /= App f a
           then eval' env applied
           else applied
  term -> term

-- | Walk a left-spine of 'App' nodes and check whether the root is an
-- 'FBuiltin'.  If so, return the builtin name and the already-accumulated
-- arguments (left-to-right order).
collectBuiltin :: DeBruijn Text -> Maybe (Text, [DeBruijn Text])
collectBuiltin = go []
  where
    go acc (App f arg) = go (arg : acc) f
    go acc (FBuiltin b) = Just (b, acc)
    go _ _ = Nothing

-- | Attempt to reduce a built-in application once all required arguments
-- are present.  Falls back to a neutral 'App' spine when the arity has not
-- been reached yet or when the arguments are not integer literals.
evalBuiltin :: Env Text -> Text -> [DeBruijn Text] -> DeBruijn Text
evalBuiltin env b args =
  case Map.lookup b Eval.builtinArity of
    Nothing -> rebuild (FBuiltin b) args   -- unknown builtin; leave as-is
    Just ar
      | length args < ar ->
          -- Still partial: return the spine as-is for now
          rebuild (FBuiltin b) args
      | length args == ar ->
          case tryReduce b args of
            Just result -> result
            Nothing     -> rebuild (FBuiltin b) args
      | otherwise ->
          -- Over-applied: reduce with exactly `ar` args then apply the rest
          let (sat, extra) = splitAt ar args
          in case tryReduce b sat of
               Just reduced -> foldl (\acc x -> eval' env (App acc x)) reduced extra
               Nothing      -> rebuild (FBuiltin b) args

-- | Church-encoded True as a DeBruijn term: lam f -> lam t -> t
churchTrue :: DeBruijn Text
churchTrue  = Lam "f" (Lam "t" (Var "t" 0))

-- | Church-encoded False as a DeBruijn term: lam f -> lam t -> f
churchFalse :: DeBruijn Text
churchFalse = Lam "f" (Lam "t" (Var "f" 1))

-- | Attempt a numeric reduction.  Returns 'Nothing' when either argument is
-- not an 'FNum' (e.g. it is still a free variable).
-- Comparison operators return Church booleans rather than integer flags.
tryReduce :: Text -> [DeBruijn Text] -> Maybe (DeBruijn Text)
tryReduce name args = case (name, args) of
  -- Integer arithmetic
  ("+",   [FNum x, FNum y]) -> Just $ FNum (x + y)
  ("-",   [FNum x, FNum y]) -> Just $ FNum (x - y)
  ("*",   [FNum x, FNum y]) -> Just $ FNum (x * y)
  ("div", [FNum x, FNum y]) | y /= 0 -> Just $ FNum (x `div` y)
  ("mod", [FNum x, FNum y]) | y /= 0 -> Just $ FNum (x `mod` y)
  ("=",   [FNum x, FNum y]) -> Just $ if x == y then churchTrue else churchFalse
  ("<",   [FNum x, FNum y]) -> Just $ if x <  y then churchTrue else churchFalse
  ("neg", [FNum x])         -> Just $ FNum (negate x)
  -- String operations
  ("str-concat", [FStr x, FStr y])         -> Just $ FStr (x <> y)
  ("str-length", [FStr x])                 -> Just $ FNum (Text.length x)
  ("str-slice",  [FNum i, FNum j, FStr x]) -> Just $ FStr (Text.take (j - i) (Text.drop i x))
  ("str=",       [FStr x, FStr y])         -> Just $ if x == y then churchTrue else churchFalse
  ("str<",       [FStr x, FStr y])         -> Just $ if x <  y then churchTrue else churchFalse
  _                                        -> Nothing

-- | Rebuild a left-spine application from a head and a list of arguments.
rebuild :: DeBruijn Text -> [DeBruijn Text] -> DeBruijn Text
rebuild = foldl App

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 t  -> t
    Nothing -> FV x
  FNum _     -> term
  FStr _     -> term
  FBuiltin _ -> term
  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
  FNum n     -> showNum n
  FStr s     -> "\"" <> escapeStr s <> "\""
  FBuiltin b -> b
  Var s _    -> s
  Lam s body -> "(lam (" <> s <> ") " <> pretty body <> ")"
  App f a    -> "(" <> pretty f <> " " <> pretty a <> ")"

-- | Render an integer, using no space between '-' and digits for negative numbers.
showNum :: Int -> Text
showNum n
  | n < 0     = "-" <> Text.pack (show (abs n))
  | otherwise = Text.pack (show n)

-- | Escape a string for pretty-printing: wrap in double-quotes and escape
-- backslashes and double-quote characters so the output round-trips through
-- the parser.
escapeStr :: Text -> Text
escapeStr = Text.concatMap escapeChar
  where
    escapeChar '"'  = "\\\""
    escapeChar '\\' = "\\\\"
    escapeChar c    = Text.singleton c