main.hs

module Main where

import Data.Map (Map)
import qualified Data.Map as Map
import Text.Read (readMaybe)
import Control.Monad.State
import Control.Monad (void)
import System.Environment (getArgs)
import System.IO (hFlush, stdout)

type WordVal = Integer
type Stack   = [WordVal]
-- A definition is now a vector/list of tokens we can index into
type Definition = [String] 

data Env = Env 
  { stack      :: Stack
  , dictionary :: Map String Definition
  } deriving Show

type Forth a = StateT Env IO a

---
-- The Inner Interpreter (Recursive with Index)
---

-- Executes a list of tokens by tracking an index (Instruction Pointer)
runDef :: Definition -> Forth ()
runDef def = step 0
  where
    n = length def
    step ip
      | ip >= n = return ()
      | otherwise = case def !! ip of
            "0BRANCH" -> do
                let offset = read (def !! (ip + 1)) :: Int
                stk <- gets stack
                case stk of
                    []     -> error "Stack Underflow"
                    (0:xs) -> modify (\e -> e { stack = xs }) >> step (ip + offset)
                    (_:xs) -> modify (\e -> e { stack = xs }) >> step (ip + 2) -- Skip offset

            "BRANCH" -> do
                let offset = read (def !! (ip + 1)) :: Int
                step (ip + offset)

            ".\"" -> do
                -- Find the closing quote within the definition list
                let (strParts, rest) = break (\x -> last x == '"') (drop (ip + 1) def)
                let fullStr = unwords (strParts ++ [init (head rest)])
                liftIO $ putStr fullStr
                liftIO $ hFlush stdout
                step (ip + 1 + length strParts + 1)

            token -> do
                execute token
                step (ip + 1)
                
---
-- The Word Executor
---

execute :: String -> Forth ()
execute token = do
  dict <- gets dictionary
  case Map.lookup token dict of
    Just def -> runDef def
    Nothing  -> case readMaybe token of
      Just n  -> modify $ \e -> e { stack = n : stack e }
      Nothing -> primitive token

dup :: Stack -> Stack
dup (x:xs) = x:x:xs
dup []     = []

_drop :: Stack -> Stack
_drop (_:xs) = xs
_drop []     = []

pop :: Stack -> (Maybe WordVal, Stack)
pop (x:xs) = (Just x, xs)
pop []   = (Nothing, [])

mutate :: (Stack -> Stack) -> Forth ()
mutate f = modify $ \e -> e { stack = f (stack e) }

stack_underflow :: String -> a
stack_underflow op = error $ "Stack underflow on " ++ op

comp :: (WordVal -> WordVal -> Bool) -> String -> Forth ()
comp f name = bop (\a b -> if f a b then 1 else 0) name

uop :: (WordVal -> WordVal) -> String -> Forth ()
uop f name = do
  s <- gets stack
  case s of
    (x:xs) -> mutate (const (f x : xs))
    []     -> stack_underflow name

stackOp :: (Stack -> Stack) -> Int -> String -> Forth ()
stackOp f required name = do
  s <- gets stack
  if length s < required
    then error $ "Stack underflow on " ++ name
    else mutate f

-- Define functions for the lambdas to avoid "uni-pattern" warnings
swap' :: Stack -> Stack
swap' (x:y:zs) = y:x:zs
swap' s        = s  -- Fallback (though stackOp prevents this)

over' :: Stack -> Stack
over' (x:y:zs) = y:x:y:zs
over' s        = s


primitive :: String -> Forth ()

primitive "+" = bop (+) "+"
primitive "-" = bop (-) "-"
primitive "*" = bop (*) "*"
primitive "/" = bop div "/"

primitive "<" = comp (<) "<"
primitive "="  = bop (\a b -> if a == b then 1 else 0) "="
primitive ">"  = bop (\a b -> if a > b then 1 else 0) ">"
primitive "<=" = comp (<=) "<="
primitive ">=" = comp (>=) ">="

primitive "0=" = uop (\x -> if x == 0 then 1 else 0) "0="

primitive "AND" = bop (\a b -> if a /= 0 && b /= 0 then 1 else 0) "AND"
primitive "OR"  = bop (\a b -> if a /= 0 || b /= 0 then 1 else 0) "OR"
primitive "DUP" = mutate dup

primitive "1+"   = uop (+ 1) "1+"
primitive "1-"   = uop (\x -> x - 1) "1-"

primitive "SWAP" = stackOp swap' 2 "SWAP"
primitive "OVER" = stackOp over' 2 "OVER"

primitive "MAX"  = bop max "MAX"
primitive "MIN"  = bop min "MIN"

primitive "." = do
  s <- gets stack
  case s of
    (x:xs) -> do
      liftIO $ putStr (show x ++ " ")
      liftIO $ hFlush stdout
      modify $ \e -> e { stack = xs }
    [] -> error "Stack underflow on ."

primitive "CR" = liftIO (putStrLn "") >> liftIO (hFlush stdout)
primitive token = error $ "Unknown word: '" ++ token ++ "'"

bop :: (WordVal -> WordVal -> WordVal) -> String -> Forth ()
bop f name = do
  s <- gets stack
  case pop s of
    (Just b, s1) -> case pop s1 of
      (Just a, s2) -> mutate (const (f a b : s2))
      (Nothing, _) -> stack_underflow name
    (Nothing, _) -> stack_underflow name

---
-- The Outer Interpreter 
---

filterComments :: [String] -> [String]
filterComments [] = []
filterComments ("(":xs) = filterComments (drop 1 $ dropWhile (/= ")") xs)
filterComments (x:xs) = x : filterComments xs

process :: [String] -> Forth ()
process [] = return ()
process (t:ts) = case t of
  "(" -> do
    let (_, restWithParen) = break (== ")") ts
    process (drop 1 restWithParen)

  ".\"" -> do
    let (strParts, restWithQuote) = break (\x -> last x == '"') ts
    case restWithQuote of
      [] -> error "Missing closing quote for .\""
      (qToken:rs) -> do
        let fullStr = unwords (strParts ++ [init qToken])
        liftIO $ putStr fullStr
        liftIO $ hFlush stdout
        process rs

  ":" -> case ts of
    (name:rest) -> do
      let (rawDef, restWithSemi) = break (== ";") rest
      -- Note: In a real Forth, ." inside a definition is compiled 
      -- as a special word. For now, we'll filter it or handle it simply.
      let compiledDef = compile (filterComments rawDef)
      modify $ \e -> e { dictionary = Map.insert name compiledDef (dictionary e) }
      process (drop 1 restWithSemi)
    [] -> error "Missing name after :"

  token -> do
    execute token
    process ts


replaceAt :: Int -> a -> [a] -> [a]
replaceAt i newVal xs = take i xs ++ [newVal] ++ drop (i + 1) xs

compile :: [String] -> [String]
compile tokens = assemble tokens 0
  where
    -- Pass 1: Create a Map of (Index of IF/ELSE) -> (Relative Offset)
    findJumps :: [String] -> Int -> [Int] -> Map.Map Int Int -> Map.Map Int Int
    findJumps [] _ _ m = m
    findJumps (t:ts) idx stack m = case t of
      "IF"   -> findJumps ts (idx + 2) (idx : stack) m
      "ELSE" -> case stack of
                  (ifIdx:rest) -> 
                    let m' = Map.insert ifIdx (idx + 2 - ifIdx) m
                    in findJumps ts (idx + 2) (idx : rest) m'
                  [] -> error "ELSE without IF"
      "THEN" -> case stack of
                  (prevIdx:rest) -> 
                    let m' = Map.insert prevIdx (idx - prevIdx) m
                    in findJumps ts idx rest m'
                  [] -> error "THEN without IF"
      _      -> findJumps ts (idx + 1) stack m

    -- Pre-calculate the jump map
    jumps = findJumps tokens 0 [] Map.empty

    -- Pass 2: Replace keywords with actual BRANCH/0BRANCH tokens
    assemble :: [String] -> Int -> [String]
    assemble [] _ = []
    assemble (t:ts) idx = case t of
      "IF"   -> "0BRANCH" : show (Map.findWithDefault 0 idx jumps) : assemble ts (idx + 2)
      "ELSE" -> "BRANCH"  : show (Map.findWithDefault 0 idx jumps) : assemble ts (idx + 2)
      "THEN" -> assemble ts idx -- THEN is just a label, it disappears
      _      -> t : assemble ts (idx + 1)

repl :: Env -> IO ()
repl env = do
  putStr "> "
  hFlush stdout
  input <- getLine
  if input == "bye" 
    then putStrLn "Goodbye!"
    else do
      (_, newEnv) <- runStateT (process (words input)) env
      putStrLn " ok"
      repl newEnv

main :: IO ()
main = do
  args <- getArgs
  let initialEnv = Env [] Map.empty
  case args of
    [filename] -> do
      contents <- readFile filename
      void $ runStateT (process (words contents)) initialEnv
      putStrLn "\nok"
    [] -> repl initialEnv -- Start REPL if no file is provided
    _  -> putStrLn "Usage: main <program.txt> or run without args for REPL"