{-# LANGUAGE GHC2021 #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

module Main where

import Control.Arrow ((>>>))
import Control.Monad (void)
import Data.Bits (shiftR, (.&.))
import Data.ByteArray qualified as BA
import Data.Char (chr, ord)
import Data.Functor (($>))
import Data.Int (Int8)
import Data.Kind (Type)
import Data.Vector qualified as V
import Data.Vector.Storable.Mutable qualified as MV
import Data.Word (Word16, Word8)
import Foreign.Ptr (Ptr, castPtr, minusPtr, plusPtr)
import Foreign.Storable qualified as S
import System.Environment (getArgs, getProgName)
import System.Exit (exitFailure)
import System.IO qualified as IO
import Text.ParserCombinators.ReadP qualified as P

class Interpreter a where
    data Program a :: Type
    parse :: String -> Program a
    interpret :: Memory -> Program a -> IO ()

newtype Memory = Memory {unMemory :: MV.IOVector Int8}
type MemIdx = Int

newMemory :: Int -> IO Memory
newMemory = fmap Memory . MV.new

memorySize :: Memory -> Int
memorySize = MV.length . unMemory

readMemory :: Memory -> MemIdx -> IO Int8
readMemory = MV.unsafeRead . unMemory

writeMemory :: Memory -> MemIdx -> Int8 -> IO ()
writeMemory = MV.unsafeWrite . unMemory

modifyMemory :: Memory -> (Int8 -> Int8) -> MemIdx -> IO ()
modifyMemory = MV.unsafeModify . unMemory

nextMemoryIndex :: Memory -> MemIdx -> MemIdx
nextMemoryIndex memory memIdx = (memIdx + 1) `mod` memorySize memory

prevMemoryIndex :: Memory -> MemIdx -> MemIdx
prevMemoryIndex memory memIdx = (memIdx - 1) `mod` memorySize memory

-- StringInterpreter with char zipper
data StringInterpreter

instance Interpreter StringInterpreter where
  data Program StringInterpreter = ProgramCZ CharZipper
  parse = ProgramCZ . czFromString
  interpret memory (ProgramCZ code) = interpretCharZipper memory code

-- the structure is [...,3,2,1,0] focus [0,1,2,3,...]
data CharZipper = CharZipper
    {czLeft :: String, czFocus :: Maybe Char, czRight :: String}

czFromString :: String -> CharZipper
czFromString = \case
    [] -> CharZipper [] Nothing []
    (x : xs) -> CharZipper [] (Just x) xs

czMoveLeft :: CharZipper -> CharZipper
czMoveLeft = \case
    CharZipper [] (Just focus) right -> CharZipper [] Nothing (focus : right)
    CharZipper (x : xs) (Just focus) right -> CharZipper xs (Just x) (focus : right)
    z -> z

czMoveRight :: CharZipper -> CharZipper
czMoveRight = \case
    CharZipper left (Just focus) [] -> CharZipper (focus : left) Nothing []
    CharZipper left (Just focus) (x : xs) -> CharZipper (focus : left) (Just x) xs
    z -> z

interpretCharZipper :: Memory -> CharZipper -> IO ()
interpretCharZipper memory = go 0
    where
    go !memIdx !program = case czFocus program of
        Nothing -> return ()
        Just c -> case c of
            '+' -> modifyMemory memory (+ 1) memIdx >> goNext
            '-' -> modifyMemory memory (subtract 1) memIdx >> goNext
            '>' -> go (nextMemoryIndex memory memIdx) program'
            '<' -> go (prevMemoryIndex memory memIdx) program'
            ',' -> do
                getChar >>= writeMemory memory memIdx . fromIntegral . ord
                goNext
            '.' -> do
                readMemory memory memIdx >>= putChar . chr . fromIntegral
                goNext
            '[' -> readMemory memory memIdx >>= \case
                0 -> go memIdx $ skipRight 1 program
                _ -> goNext
            ']' -> readMemory memory memIdx >>= \case
                0 -> goNext
                _ -> go memIdx $ skipLeft 1 program
            _ -> goNext
        where
        program' = czMoveRight program
        goNext = go memIdx program'

skipRight :: Int -> CharZipper -> CharZipper
skipRight !depth !program
    | depth == 0 = program'
    | otherwise = case czFocus program' of
        Nothing -> error "No matching [ while skipping the loop forward"
        Just '[' -> skipRight (depth + 1) program'
        Just ']' -> skipRight (depth - 1) program'
        _ -> skipRight depth program'
    where
    program' = czMoveRight program

skipLeft :: Int -> CharZipper -> CharZipper
skipLeft !depth !program
    | depth == 0 = czMoveRight program
    | otherwise = case czFocus program' of
        Nothing -> error "No matching ] while skipping the loop backward"
        Just ']' -> skipLeft (depth + 1) program'
        Just '[' -> skipLeft (depth - 1) program'
        _ -> skipLeft depth program'
    where
    program' = czMoveLeft program

-- AST
data ASTInterpreter
instance Interpreter ASTInterpreter where
    data Program ASTInterpreter = ProgramAST Instructions
    parse = ProgramAST . parseToInstructions
    interpret memory (ProgramAST instrs) = interpretAST memory instrs

type Instructions = V.Vector Instruction
data Instruction
    = Inc                   -- +
    | Dec                   -- -
    | MoveR                 -- >
    | MoveL                 -- <
    | GetC                  -- ,
    | PutC                  -- .
    | Loop Instructions     -- []
    deriving (Show)

parseToInstructions :: String -> Instructions
parseToInstructions code =
    V.fromList $ case P.readP_to_S (P.many instrParser <* P.eof) code of
        [(res, "")] -> res
        out -> error $ "Unexpected output while parsing: " <> show out
    where instrParser = P.choice
            [
                P.char '+' $> Inc,
                P.char '-' $> Dec,
                P.char '>' $> MoveR,
                P.char '<' $> MoveL,
                P.char ',' $> GetC,
                P.char '.' $> PutC,
                Loop . V.fromList <$> P.between (P.char '[') (P.char ']') (P.many instrParser)
            ]

interpretAST :: Memory -> Instructions -> IO ()
interpretAST memory = void . interpretInstrs 0 memory

interpretInstrs :: MemIdx -> Memory -> Instructions -> IO MemIdx
interpretInstrs memIdx !memory !program = go memIdx 0
    where
    go !memIdx !progIdx
        | progIdx == V.length program = return memIdx
        | otherwise = case program V.! progIdx of
            Inc -> modifyMemory memory (+ 1) memIdx >> goNext
            Dec -> modifyMemory memory (subtract 1) memIdx >> goNext
            MoveR -> go (nextMemoryIndex memory memIdx) $ progIdx + 1
            MoveL -> go (prevMemoryIndex memory memIdx) $ progIdx + 1
            GetC -> do
                getChar >>= writeMemory memory memIdx . fromIntegral . ord
                goNext
            PutC -> do
                readMemory memory memIdx >>= putChar . chr . fromIntegral
                goNext
            Loop instrs -> readMemory memory memIdx >>= \case
                0 -> goNext
                _ -> interpretInstrs memIdx memory instrs >>= flip go progIdx
        where
        goNext = go memIdx $ progIdx + 1
    
main :: IO ()
main = do
    IO.hSetBuffering IO.stdin IO.NoBuffering
    IO.hSetBuffering IO.stdout IO.LineBuffering
    progName <- getProgName
    let usage = "Usage: " <> progName <> " -(s|a|b|o) <bf_file>"
    getArgs >>= \case
        [interpreterType, fileName] -> do
            code <- filter (`elem` "+-.,><[]") <$> readFile fileName
            memory <- newMemory 30000
            parseAndInterpret memory code usage interpreterType
        _ -> exitWithMsg usage
    where
        parseAndInterpret memory code usage = \case
            "-s" -> interpret @StringInterpreter memory $ parse code
            "-a" -> interpret @ASTInterpreter memory $ parse code
            -- "-b" -> interpret @BytecodeInterpreter memory $ parse code
            -- "-o" -> interpret @OptimizingBytecodeInterpreter memory $ parse code
            t -> exitWithMsg $ "Invalid interpreter type: " <> t <> "\n" <> usage
        exitWithMsg msg = IO.hPutStrLn IO.stderr msg >> exitFailure