{-# LANGUAGE ScopedTypeVariables #-}
-- A Haskell port of ilo.c
-- Build with: ghc -O2 Ilo.hs -o ilo-hs

import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BSI
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BL
import Data.Array.IO (IOArray, newArray, readArray, writeArray, getBounds)
import Data.Bits ((.&.), (.|.), shiftL, shiftR, xor, complement)
import Data.Char (chr)
import Data.Int (Int32)
import Data.Word (Word8, Word32)
import Data.IORef
import Control.Monad (when, forM_, unless, void)
import Control.Exception (IOException, catch, throwIO, Exception)
import System.Environment (getArgs)
import System.IO (withFile, IOMode(ReadMode, ReadWriteMode), hSeek, SeekMode(AbsoluteSeek),
                  stdout, hFlush, hIsEOF, stdin, hSetBuffering, BufferMode(NoBuffering),
                  stderr, hPutStrLn)
import System.Exit (exitSuccess, exitFailure, die)
import Foreign.Ptr (Ptr, castPtr)
import Foreign.Storable (peekByteOff, pokeByteOff)

type Cell = Int32

-- Error types
data VmError
  = StackUnderflow String
  | StackOverflow String
  | MemoryOutOfBounds Int
  | DivisionByZero
  | InvalidBlockNumber Int
  | IOError String
  deriving (Show)

instance Exception VmError

toW32 :: Cell -> Word32
toW32 = fromIntegral

fromW32 :: Word32 -> Cell
fromW32 = fromIntegral

maskShift :: Cell -> Int
maskShift b = fromIntegral (toW32 b .&. 31)

wrapAdd :: Cell -> Cell -> Cell
wrapAdd a b = fromW32 (toW32 a + toW32 b)

wrapSub :: Cell -> Cell -> Cell
wrapSub a b = fromW32 (toW32 a - toW32 b)

wrapMul :: Cell -> Cell -> Cell
wrapMul a b = fromW32 (toW32 a * toW32 b)

wrapAnd :: Cell -> Cell -> Cell
wrapAnd a b = fromW32 (toW32 a .&. toW32 b)

wrapOr :: Cell -> Cell -> Cell
wrapOr a b = fromW32 (toW32 a .|. toW32 b)

wrapXor :: Cell -> Cell -> Cell
wrapXor a b = fromW32 (toW32 a `xor` toW32 b)

wrapShl :: Cell -> Cell -> Cell
wrapShl a b = fromW32 (toW32 a `shiftL` maskShift b)

wrapShr :: Cell -> Cell -> Cell
wrapShr a b = fromIntegral ((a `shiftR` maskShift b) :: Int32)

readLE32 :: Ptr Word8 -> Int -> IO Cell
readLE32 ptr off = do
  b0 <- peekByteOff ptr off       :: IO Word8
  b1 <- peekByteOff ptr (off + 1) :: IO Word8
  b2 <- peekByteOff ptr (off + 2) :: IO Word8
  b3 <- peekByteOff ptr (off + 3) :: IO Word8
  let w = (fromIntegral b0)
        .|. (fromIntegral b1 `shiftL` 8)
        .|. (fromIntegral b2 `shiftL` 16)
        .|. (fromIntegral b3 `shiftL` 24) :: Word32
  pure (fromW32 w)

writeLE32 :: Ptr Word8 -> Int -> Cell -> IO ()
writeLE32 ptr off v = do
  let w = toW32 v
  pokeByteOff ptr off       (fromIntegral ( w          .&. 0xFF) :: Word8)
  pokeByteOff ptr (off + 1) (fromIntegral ((w `shiftR` 8) .&. 0xFF) :: Word8)
  pokeByteOff ptr (off + 2) (fromIntegral ((w `shiftR` 16) .&. 0xFF) :: Word8)
  pokeByteOff ptr (off + 3) (fromIntegral ((w `shiftR` 24) .&. 0xFF) :: Word8)

memSize, dsSize, asSize :: Int
memSize = 65536
dsSize  = 33
asSize  = 257

data Vm = Vm
  { ip         :: IORef Int       -- instruction pointer
  , sp         :: IORef Int       -- data stack pointer
  , rp         :: IORef Int       -- address stack pointer
  , ds         :: IOArray Int Cell
  , as         :: IOArray Int Cell
  , mem        :: IOArray Int Cell
  , blocksPath :: FilePath
  , romPath    :: FilePath
  }

newVm :: FilePath -> FilePath -> IO Vm
newVm blocks rom = do
  ipRef <- newIORef 0
  spRef <- newIORef 0
  rpRef <- newIORef 0
  dataStack <- newArray (0, dsSize - 1) 0
  addrStack <- newArray (0, asSize - 1) 0
  memory    <- newArray (0, memSize - 1) 0
  pure Vm { ip = ipRef, sp = spRef, rp = rpRef
          , ds = dataStack, as = addrStack, mem = memory
          , blocksPath = blocks, romPath = rom
          }

-- Safe stack operations with bounds checking
push :: Vm -> Cell -> IO ()
push vm v = do
  modifyIORef' (sp vm) (+1)
  spVal <- readIORef (sp vm)
  when (spVal >= dsSize) $ throwIO (StackOverflow "Data stack overflow")
  writeArray (ds vm) spVal v

pop :: Vm -> IO Cell
pop vm = do
  spVal <- readIORef (sp vm)
  when (spVal < 0) $ throwIO (StackUnderflow "Data stack underflow")
  val <- readArray (ds vm) spVal
  modifyIORef' (sp vm) (subtract 1)
  pure val

top :: Vm -> IO Cell
top vm = do
  spVal <- readIORef (sp vm)
  when (spVal < 0) $ throwIO (StackUnderflow "Cannot read from empty stack")
  readArray (ds vm) spVal

next :: Vm -> IO Cell
next vm = do
  spVal <- readIORef (sp vm)
  when (spVal < 1) $ throwIO (StackUnderflow "Not enough items on stack")
  readArray (ds vm) (spVal - 1)

setTop :: Vm -> Cell -> IO ()
setTop vm v = do
  spVal <- readIORef (sp vm)
  when (spVal < 0) $ throwIO (StackUnderflow "Cannot write to empty stack")
  writeArray (ds vm) spVal v

setNext :: Vm -> Cell -> IO ()
setNext vm v = do
  spVal <- readIORef (sp vm)
  when (spVal < 1) $ throwIO (StackUnderflow "Not enough items on stack")
  writeArray (ds vm) (spVal - 1) v

addrTop :: Vm -> IO Cell
addrTop vm = do
  rpVal <- readIORef (rp vm)
  when (rpVal < 0) $ throwIO (StackUnderflow "Return stack underflow")
  readArray (as vm) rpVal

setAddrTop :: Vm -> Cell -> IO ()
setAddrTop vm v = do
  rpVal <- readIORef (rp vm)
  when (rpVal < 0) $ throwIO (StackUnderflow "Return stack underflow")
  writeArray (as vm) rpVal v

pushAddr :: Vm -> Cell -> IO ()
pushAddr vm v = do
  modifyIORef' (rp vm) (+1)
  rpVal <- readIORef (rp vm)
  when (rpVal >= asSize) $ throwIO (StackOverflow "Return stack overflow")
  writeArray (as vm) rpVal v

popAddr :: Vm -> IO Cell
popAddr vm = do
  rpVal <- readIORef (rp vm)
  when (rpVal < 0) $ throwIO (StackUnderflow "Return stack underflow")
  val <- readArray (as vm) rpVal
  modifyIORef' (rp vm) (subtract 1)
  pure val

-- Safe memory access
fetchMem :: Vm -> Int -> IO Cell
fetchMem vm addr = do
  when (addr < 0 || addr >= memSize) $ throwIO (MemoryOutOfBounds addr)
  readArray (mem vm) addr

storeMem :: Vm -> Int -> Cell -> IO ()
storeMem vm addr val = do
  when (addr < 0 || addr >= memSize) $ throwIO (MemoryOutOfBounds addr)
  writeArray (mem vm) addr val

loadImage :: Vm -> IO ()
loadImage vm = do
  file <- BS.readFile (romPath vm) `catch` \(e :: IOException) ->
    throwIO (IOError $ "Cannot read ROM file: " ++ show e)
  let cells = min memSize (BS.length file `div` 4)
  forM_ [0 .. cells - 1] $ \i -> do
    let idx = i * 4
        b0 = fromIntegral (BS.index file idx)       :: Word32
        b1 = fromIntegral (BS.index file (idx + 1)) :: Word32
        b2 = fromIntegral (BS.index file (idx + 2)) :: Word32
        b3 = fromIntegral (BS.index file (idx + 3)) :: Word32
        w  = b0 .|. (b1 `shiftL` 8) .|. (b2 `shiftL` 16) .|. (b3 `shiftL` 24)
        c  = fromIntegral (w :: Word32) :: Cell
    writeArray (mem vm) i c
  writeIORef (ip vm) 0
  writeIORef (sp vm) 0
  writeIORef (rp vm) 0

saveImage :: Vm -> IO ()
saveImage vm = do
  builder <- fmap mconcat $
    mapM (\i -> do
            v <- readArray (mem vm) i
            let w = fromIntegral v :: Word32
            pure $ BB.word32LE w) [0 .. memSize - 1]
  BL.writeFile (romPath vm) (BB.toLazyByteString builder) `catch` \(e :: IOException) ->
    throwIO (IOError $ "Cannot write ROM file: " ++ show e)

blockCommon :: Vm -> IO (Int, Int)
blockCommon vm = do
  b <- pop vm -- block buffer (cell offset)
  a <- pop vm -- block number
  let offset = fromIntegral a * 4096
      bufCell = fromIntegral b
  when (a < 0) $ throwIO (InvalidBlockNumber (fromIntegral a))
  when (bufCell < 0 || bufCell + 1024 > memSize) $
    throwIO (MemoryOutOfBounds bufCell)
  pure (offset, bufCell)

readBlock :: Vm -> IO ()
readBlock vm = do
  (offset, bufCell) <- blockCommon vm
  (withFile (blocksPath vm) ReadMode $ \h -> do
    hSeek h AbsoluteSeek (fromIntegral offset)
    bytes <- BS.hGet h 4096
    BS.useAsCString bytes $ \ptr -> do
      let cells = BS.length bytes `div` 4
      forM_ [0 .. cells - 1] $ \i -> do
        v <- readLE32 (castPtr ptr) (i * 4)
        writeArray (mem vm) (bufCell + i) v)
    `catch` \(e :: IOException) ->
      throwIO (IOError $ "Cannot read block file: " ++ show e)

writeBlock :: Vm -> IO ()
writeBlock vm = do
  (offset, bufCell) <- blockCommon vm
  let cells = 4096 `div` 4
  (do chunk <- BSI.create 4096 $ \ptr -> do
        forM_ [0 .. cells - 1] $ \i -> do
          v <- readArray (mem vm) (bufCell + i)
          writeLE32 ptr (i * 4) v
      withFile (blocksPath vm) ReadWriteMode $ \h -> do
        hSeek h AbsoluteSeek (fromIntegral offset)
        BS.hPut h chunk)
    `catch` \(e :: IOException) ->
      throwIO (IOError $ "Cannot write block file: " ++ show e)

saveIp :: Vm -> IO ()
saveIp vm = do
  ipVal <- readIORef (ip vm)
  pushAddr vm (fromIntegral ipVal)

symmetric :: Vm -> Cell -> IO ()
symmetric vm dividend = do
  n <- next vm
  when (dividend >= 0 && n < 0) $ do
    t <- top vm
    setTop vm (t + 1)
    setNext vm (n - dividend)

-- Instruction implementations
instLit, instDup, instDrop, instSwap, instPush, instPop, instJump, instCall, instCCall, instCJump, instRet :: Vm -> IO ()
instEq, instNe, instLt, instGt, instFetch, instStore, instAdd, instSub, instMul, instDiv :: Vm -> IO ()
instAnd, instOr, instXor, instShiftL, instShiftR, instCompare, instCopy :: Vm -> IO ()

instLit vm = do
  modifyIORef' (ip vm) (+1)
  ipVal <- readIORef (ip vm)
  val <- fetchMem vm ipVal
  push vm val

instDup vm = top vm >>= push vm

instDrop vm = do
  spVal <- readIORef (sp vm)
  when (spVal >= 0) $ writeArray (ds vm) spVal 0
  modifyIORef' (sp vm) (subtract 1)

instSwap vm = do
  a <- top vm
  b <- next vm
  setTop vm b
  setNext vm a

instPush vm = do
  t <- pop vm
  pushAddr vm t

instPop vm = do
  t <- popAddr vm
  push vm t

instJump vm = do
  target <- pop vm
  writeIORef (ip vm) (fromIntegral target - 1)

instCall vm = do
  saveIp vm
  target <- pop vm
  writeIORef (ip vm) (fromIntegral target - 1)

instCCall vm = do
  a <- pop vm
  flag <- pop vm
  when (flag /= 0) $ do
    saveIp vm
    writeIORef (ip vm) (fromIntegral a - 1)

instCJump vm = do
  a <- pop vm
  flag <- pop vm
  when (flag /= 0) $ writeIORef (ip vm) (fromIntegral a - 1)

instRet vm = do
  r <- popAddr vm
  writeIORef (ip vm) (fromIntegral r)

instEq vm = do
  a <- next vm
  b <- top vm
  setNext vm (if a == b then -1 else 0)
  void $ pop vm

instNe vm = do
  a <- next vm
  b <- top vm
  setNext vm (if a /= b then -1 else 0)
  void $ pop vm

instLt vm = do
  a <- next vm
  b <- top vm
  setNext vm (if a < b then -1 else 0)
  void $ pop vm

instGt vm = do
  a <- next vm
  b <- top vm
  setNext vm (if a > b then -1 else 0)
  void $ pop vm

instFetch vm = do
  t <- top vm
  val <- fetchMem vm (fromIntegral t)
  setTop vm val

instStore vm = do
  t <- top vm
  n <- next vm
  storeMem vm (fromIntegral t) n
  void $ pop vm
  void $ pop vm

instAdd vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapAdd a b)
  void $ pop vm

instSub vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapSub a b)
  void $ pop vm

instMul vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapMul a b)
  void $ pop vm

instDiv vm = do
  a <- top vm
  when (a == 0) $ throwIO DivisionByZero
  b <- next vm
  let (q, r) = quotRem b a
  setTop vm q
  setNext vm r
  symmetric vm b

instAnd vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapAnd a b)
  void $ pop vm

instOr vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapOr a b)
  void $ pop vm

instXor vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapXor a b)
  void $ pop vm

instShiftL vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapShl a b)
  void $ pop vm

instShiftR vm = do
  a <- next vm
  b <- top vm
  setNext vm (wrapShr a b)
  void $ pop vm
instCompare vm = do
  l <- pop vm
  d <- pop vm
  s <- top vm
  setTop vm (-1)
  forM_ [0 .. fromIntegral l - 1] $ \i -> do
    let srcAddr = fromIntegral s + i
        dstAddr = fromIntegral d + i
    v1 <- fetchMem vm dstAddr
    v2 <- fetchMem vm srcAddr
    when (v1 /= v2) $ setTop vm 0

instCopy vm = do
  l <- pop vm
  d <- pop vm
  s <- pop vm
  forM_ [0 .. fromIntegral l - 1] $ \i -> do
    let srcAddr = fromIntegral s + i
        dstAddr = fromIntegral d + i
    v <- fetchMem vm srcAddr
    storeMem vm dstAddr v
-- I/O device handlers
ioOutput, ioInput, ioReadBlock, ioWriteBlock, ioSaveImage, ioLoadImage, ioHalt, ioStackInfo :: Vm -> IO ()

ioOutput vm = do
  c <- pop vm
  putChar (chr (fromIntegral (c .&. 0xFF)))
  hFlush stdout

ioInput vm = do
  eof <- hIsEOF stdin
  if eof
    then exitSuccess
    else do
      b <- getChar
      push vm (fromIntegral (fromEnum b .&. 0xFF))

ioReadBlock = readBlock
ioWriteBlock = writeBlock
ioSaveImage = saveImage

ioLoadImage vm = do
  loadImage vm
  writeIORef (ip vm) (-1)

ioHalt vm = writeIORef (ip vm) memSize

ioStackInfo vm = do
  spVal <- readIORef (sp vm)
  rpVal <- readIORef (rp vm)
  push vm (fromIntegral spVal)
  push vm (fromIntegral rpVal)

instIO :: Vm -> IO ()
instIO vm = do
  code <- pop vm
  case code of
    0 -> ioOutput vm
    1 -> ioInput vm
    2 -> ioReadBlock vm
    3 -> ioWriteBlock vm
    4 -> ioSaveImage vm
    5 -> ioLoadImage vm
    6 -> ioHalt vm
    7 -> ioStackInfo vm
    _ -> pure ()  -- Ignore unknown I/O codes

-- Instruction dispatch
processInstruction :: Vm -> Int -> IO ()
processInstruction vm opcode = case opcode of
  0  -> pure ()              -- nop
  1  -> instLit vm           -- literal
  2  -> instDup vm           -- duplicate
  3  -> instDrop vm          -- drop
  4  -> instSwap vm          -- swap
  5  -> instPush vm          -- push to return stack
  6  -> instPop vm           -- pop from return stack
  7  -> instJump vm          -- jump
  8  -> instCall vm          -- call
  9  -> instCCall vm         -- conditional call
  10 -> instCJump vm         -- conditional jump
  11 -> instRet vm           -- return
  12 -> instEq vm            -- equal
  13 -> instNe vm            -- not equal
  14 -> instLt vm            -- less than
  15 -> instGt vm            -- greater than
  16 -> instFetch vm         -- fetch from memory
  17 -> instStore vm         -- store to memory
  18 -> instAdd vm           -- add
  19 -> instSub vm           -- subtract
  20 -> instMul vm           -- multiply
  21 -> instDiv vm           -- divide/modulo
  22 -> instAnd vm           -- bitwise and
  23 -> instOr vm            -- bitwise or
  24 -> instXor vm           -- bitwise xor
  25 -> instShiftL vm        -- shift left
  26 -> instShiftR vm        -- shift right
  27 -> instCompare vm       -- compare memory regions
  28 -> instCopy vm          -- copy memory region
  29 -> instIO vm            -- I/O operations
  _  -> pure ()              -- unknown opcode (nop)

-- Process packed opcode bundle (4 instructions per cell)
processBundle :: Vm -> Cell -> IO ()
processBundle vm opcode = do
  processInstruction vm (fromIntegral (opcode .&. 0xFF))
  processInstruction vm (fromIntegral ((opcode `shiftR` 8) .&. 0xFF))
  processInstruction vm (fromIntegral ((opcode `shiftR` 16) .&. 0xFF))
  processInstruction vm (fromIntegral ((opcode `shiftR` 24) .&. 0xFF))

-- Main execution loop
execute :: Vm -> IO ()
execute vm = go
  where
    go = do
      ipVal <- readIORef (ip vm)
      when (ipVal < memSize && ipVal >= 0) $ do
        opcode <- fetchMem vm ipVal
        processBundle vm opcode
        modifyIORef' (ip vm) (+1)
        go

-- Parse command-line arguments safely
parseArgs :: [String] -> (FilePath, FilePath)
parseArgs args = case args of
  (blocks:rom:_) -> (blocks, rom)
  (blocks:_)     -> (blocks, "ilo.rom")
  []             -> ("ilo.blocks", "ilo.rom")

-- Print stack contents on exit
printStack :: Vm -> IO ()
printStack vm = do
  spVal <- readIORef (sp vm)
  vals <- mapM (readArray (ds vm)) [1..spVal]
  unless (null vals) $ do
    putStr $ unwords (map show vals)
    putStrLn ""

main :: IO ()
main = do
  args <- getArgs
  let (blocks, rom) = parseArgs args

  hSetBuffering stdout NoBuffering

  -- Run VM with error handling
  result <- catch (do
    vm <- newVm blocks rom
    loadImage vm
    execute vm
    printStack vm
    pure True) $ \(e :: VmError) -> do
      hPutStrLn stderr $ "VM Error: " ++ show e
      pure False

  unless result exitFailure
