{-# LANGUAGE ViewPatterns, MultilineStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveAnyClass #-}
module Gyehoek.Stack.VM
  ( VM(..)
  , Env(..)
  , eval
  , trace
  , module Gyehoek.Stack.Syntax
  , writeObj
  , traceEval
  ) where

import Gyehoek.Stack.Syntax
import Control.Lens
import qualified Data.HashMap.Strict as H
import Data.List (unfoldr, intersperse, compareLength)
import Gyehoek.Prelude
import qualified Data.List.NonEmpty as NE
import Lucid
import Data.Foldable (traverse_)
import qualified Gyehoek.Sexp as S
import Gyehoek.Jalmot
import Text.Pretty.Simple (pStringNoColor, pShowNoColor)
import Effectful.State.Static.Local (runState, evalState, get)
import Data.Traversable
import Control.Applicative (Alternative(..))
import Gyehoek.Sexp.Print (htmlData)
import Control.DeepSeq (deepseq, ($!!))
import Gyehoek.Sexp.Print (htmlData, htmlDatum)
import Control.DeepSeq (deepseq, ($!!))
import Data.String (fromString)
import Data.Monoid (First)
import GHC.Stack (popCallStack)
import Data.Maybe (fromMaybe)


-- | non-essential information maintained only to aide in debugging.
data DebugVM = MkDebugVM
  { activeRoutine :: Label
  }
  deriving (Show, Generic)

newtype Frame = MkFrame { locals :: List Obj }
  deriving stock (Show, Generic)

-- affine
returnAddress :: Traversal' Frame Obj
returnAddress = #locals . _last

-- affine
activeProcedure :: Traversal' Frame Obj
activeProcedure = #locals . _init . _last

newtype Stack = MkStack { frames :: NonEmpty Frame }
  deriving stock (Show, Generic)

data VM = MkVM
  { stack :: Stack
  , code :: List Instr
  , tail :: Tail
  , registers :: HashMap Reg Obj
  , stdout :: Text
  , result :: Maybe (List Obj)
  , debug :: DebugVM
  }
  deriving (Show, Generic)

type instance Index Frame = Int
type instance IxValue Frame = Obj

instance Ixed Frame where
  ix j = wrappedIso . ix j

instance Cons Frame Frame Obj Obj where
  _Cons = prism'
    (\(x,MkFrame xs) -> MkFrame (x:xs))
    \case
      MkFrame (x:xs) -> Just (x, MkFrame xs)
      MkFrame [] -> Nothing

instance Each Frame Frame Obj Obj where each = wrappedIso . each

instance Each Stack Stack Frame Frame where each = wrappedIso . each

pushes :: Foldable f => f Obj -> Frame -> Frame
pushes = flip $ foldr cons

_NonEmpty :: Iso (NonEmpty a) (NonEmpty b) (a, List a) (b, List b)
_NonEmpty = iso
  (\(x:|xs) -> (x,xs))
  (\(x,xs) -> x:|xs)

pushFrame :: Frame -> Stack -> Stack
pushFrame f (MkStack xs) = MkStack $ NE.cons f xs

activeFrame :: Lens' VM Frame
activeFrame = #stack . #frames . _NonEmpty . _1

data Env = MkEnv
  { labels :: HashMap Label Routine
  }
  deriving (Show, Generic)

step :: Jalmot :> es => Env -> VM -> Eff es VM
step g vm = case vm ^. #code of
  c:cs -> stepI g (vm & #code .~ cs) c
  [] -> stepT g vm vm.tail

vmerror :: (HasCallStack, Jalmot :> es) => Text -> Eff es a
vmerror = throwError . VMError

stepI :: Jalmot :> es => Env -> VM -> Instr -> Eff es VM

stepI e vm (Load r j) = do
  x <- expectOf [i|object at index #{j}|] (activeFrame . ix j) vm
  pure $ vm & #registers . at r ?~ x

stepI e vm (Push v) = traverseOf activeFrame push vm
  where push xs = cons <$> evalVal e vm v <*> pure xs

stepI g vm (Prim p) = stepP g vm p

stepI e vm (Pop r) = case vm ^? activeFrame . _Cons of
  Nothing -> vmerror "empty stack"
  Just (x,xs) -> pure $ vm & #registers . at r ?~ x
               & activeFrame .~ xs

stepI e vm ins = vmerror [i|unimplemented instruction: #{ins}|]

stepT :: Jalmot :> es => Env -> VM -> Tail -> Eff es VM

stepT g vm tc@(Call nargs) = do
  (args,f,ret,frm) <- parseCall nargs (vm ^. activeFrame)
    & expectOf [i|bad call: #{show tc}|] _Just 
  rt <- getRoutine g f
  let newFrame = MkFrame $ args ++ [f,ret]
  pure $ vm
    & jumpToRoutine rt
    & activeFrame .~ frm
    & #stack %~ pushFrame newFrame
    -- it is not essential we clear the registers, but it'll
    -- make bugs more obvious.
    & #registers .~ mempty

stepT g vm tc@(Return nret) = do
  (xs,_) <- splitAtExact nret (vm ^. activeFrame . #locals)
    & expectOf [i|bad return: #{show tc}|] _Just
  expectOf [i|no return addr|] (activeFrame . returnAddress) vm >>= \case
    ObjLabel "halt" -> pure $ vm & #result ?~ xs
    ra -> do
      rt <- getRoutine g ra
      vm & traverseOf #stack (fmap snd . popFrame)
         & mapped . activeFrame %~ pushes xs
         & mapped %~ jumpToRoutine rt
         -- it is not essential we clear the registers, but it'll make
         -- bugs more obvious.
         & mapped . #registers .~ mempty

stepT g vm tc@(TailCall nargs) = do
  (args,f,ra) <- parseTailCall nargs (vm ^. activeFrame)
    & expectOf [i|bad call: #{show tc}|] _Just
  case f of
    ObjLabel "halt" -> pure $ vm & #result ?~ args
    _ -> do
      rt <- getRoutine g f
      let newFrame = MkFrame $ args ++ [f, ra]
      pure $ vm
        & jumpToRoutine rt
        -- replace the active frame; don't push a new one.
        & activeFrame .~ newFrame
        -- it is not essential we clear the registers, but it'll make
        -- bugs more obvious.
        & #registers .~ mempty

stepT g vm (If c t f) = do
  branch <- evalVal g vm c <&> \case
      ObjImm (ImmBool False) -> f
      _ -> t
  pure $ jumpToBlock branch vm

stepP :: Jalmot :> es => Env -> VM -> Prim Val -> Eff es VM
stepP g vm p = traverse (evalVal g vm) p >>= \case
  PrimZeroP x -> case x of
    ObjImm (ImmInt n) -> ret1 . ObjImm . ImmBool $ n == 0
    _ -> vmerror [i|bad arg to zero?: #{x}|]
  PrimAdd x y -> arith_binop (+) x y
  PrimMul x y -> arith_binop (*) x y
  PrimSub x y -> arith_binop (-) x y
  PrimDiv x y -> arith_binop div x y
  PrimMakeClosure f env ->
    case f of
      ObjImm (ImmLabel l) -> ret1 . ObjHob $ HobClosure l env
      _ -> vmerror [i|expected label, got #{f}|]
  PrimEnvCode env ->
    case env of
      ObjHob (HobClosure l _) -> ret1 . ObjImm . ImmLabel $ l
      _ -> vmerror [i|expected closure, got #{env}|]
  PrimEnv -> do
    x <- vm & expectOf "expected closure" (activeFrame . activeProcedure)
    ret1 x
  PrimEnvRef n -> do
    (label,env) <- vm & expectOf "expected closure"
        (activeFrame . activeProcedure . #_ObjHob . #_HobClosure) 
    x <- env & expectOf "expected upval" (ix n)
    ret1 x
  PrimCons x y -> ret1 $ ObjHob $ HobPair x y
  PrimCar x -> case x of
      ObjHob (HobPair car _) -> ret1 car
      _ -> vmerror [i|expected pair, got ${x}|]
  PrimCdr x -> case x of
      ObjHob (HobPair _ cdr) -> ret1 cdr
      _ -> vmerror [i|expected pair, got ${x}|]
  x -> vmerror [i|unimplemented prim: #{p}|]
  where
    ret vs = pure $ vm & activeFrame . #locals <>:~ vs
    ret1 v = ret [v]
    arith_binop op (ObjImm (ImmInt x)) (ObjImm (ImmInt y)) =
      ret1 $ ObjImm (ImmInt (op x y))
    arith_binop _ x y = vmerror [i|bad arith: #{x}, #{y}|]



popFrame :: (HasCallStack, Jalmot :> es) => Stack -> Eff es (Frame, Stack)
popFrame stk = case stk ^. #frames . to NE.uncons of
  (_, Nothing) -> vmerror "no frame to pop"
  (f, Just fs) -> pure (f, stk & #frames .~ fs)

jumpToBlock :: Block -> VM -> VM
jumpToBlock b vm = vm
  & #code .~ b.code
  & #tail .~ b.tail

jumpToRoutine :: Routine -> VM -> VM
jumpToRoutine rt vm = vm
  & jumpToBlock rt.start
  & #debug . #activeRoutine .~ rt.label

getLabel :: Obj -> Maybe Label
getLabel = \case
  ObjHob (HobClosure {label}) -> Just label
  ObjImm (ImmLabel label) -> Just label
  x -> Nothing

getRoutine :: (HasCallStack, Jalmot :> es) => Env -> Obj -> Eff es Routine
getRoutine g f = do
  l <- getLabel f & expectOf [i|no label for #{f}|] _Just
  case g ^. #labels . at l of
    Just rt -> pure rt
    Nothing -> vmerror [i|undefined label #{l}|]

expectOf
  :: (HasCallStack, Jalmot :> es)
  => Text -> Getting (First a) s a -> s -> Eff es a
expectOf msg l = maybe (vmerror msg) pure . preview l

evalToLabel :: Jalmot :> es => Env -> VM -> Val -> Eff es Label
evalToLabel e vm v =
  evalVal e vm v >>= \case
    ObjImm (ImmLabel x) -> pure x
    x -> vmerror [i|not a label: #{x}|]

evalVal :: Jalmot :> es => Env -> VM -> Val -> Eff es Obj
evalVal e vm = \case
  ValImm imm -> pure $ ObjImm imm
  ValReg r -> case vm ^. #registers . at r of
    Just x -> pure x
    Nothing -> vmerror [i|undefined register: #{r}|]

splitAtExact :: Int -> List a -> Maybe (List a, List a)
splitAtExact n xs = case compareLength xs n of
  (EQ;GT) -> Just $ splitAt n xs
  LT -> Nothing

takeExact :: Int -> List a -> Maybe (List a)
takeExact n xs = case compareLength xs n of
  (EQ;GT) -> Just $ take n xs
  LT -> Nothing

parseCall :: Int -> Frame -> Maybe (List Obj, Obj, Obj, Frame)
parseCall nargs frm = do
  (xs,ys) <- splitAtExact (nargs+2) (frm ^. #locals)
  let (xs',[f,ret]) = splitAt nargs xs
  pure (xs',f,ret,MkFrame ys)

parseTailCall :: Int -> Frame -> Maybe (List Obj, Obj, Obj)
parseTailCall nargs frm = do
  (xs,_) <- splitAtExact (nargs+1) (frm ^. #locals)
  let (xs',f) = xs ^?! _Snoc
  pure (xs',f,frm ^?! returnAddress)

initialVM :: VM
initialVM = MkVM
  { stack = MkStack . NE.singleton . MkFrame $
    [ ObjLabel "start"
    , ObjLabel "<nowhere at all>"
    , ObjLabel "halt"
    ]
  , tail = TailCall 0
  , code = []
  , registers = mempty
  , stdout = ""
  , result = Nothing
  , debug = MkDebugVM
    { activeRoutine = "<nowhere>"
    }
  }

initialEnv :: Program -> Env
initialEnv p = MkEnv
  { labels = p.routines
  }

loop :: (a -> Either b a) -> a -> b
loop f a = case f a of
  Right a' -> loop f a'
  Left b -> b

loopM :: Monad m => (a -> m (Either b a)) -> a -> m b
loopM f a = f a >>= \case
  Right a' -> loopM f a'
  Left b -> pure b

eval :: Jalmot :> es => Program -> Eff es (List Obj)
eval p = initialVM & loopM \vm -> case vm ^. #result of
  Nothing -> Right <$> step (initialEnv p) vm
  Just rs -> pure . Left $ rs

data Trace
  = Step { vm :: VM, next :: Trace }
  | StepToSuccess { vm :: VM, result :: List Obj }
  | StepToFailure { vm :: VM, err :: AJalmotCS }
  deriving (Show)

trace :: Program -> Trace
trace p = go (initialEnv p) initialVM
  where
    go g vm =
      case vm.result of
        Just rs -> StepToSuccess vm rs
        Nothing ->
          case runPureEff . runJalmotCS $ step g vm of
            Left err -> StepToFailure vm err
            Right vm' -> Step vm (go g vm')

writeObj :: Obj -> Text
writeObj = runJalmotUnsafe . S.encodeWith' S.datumIso

traceEval :: IOE :> es => Program -> Eff es ()
traceEval p = do
  let t = trace p
  liftIO . renderToFile "trace.html" . ppDoc p $ t

ppDoc :: Program -> Trace -> Html ()
ppDoc p t =
  html_ do
    head_ do
      title_ "stackify trace"
      style_ """
        pre {
          max-width: 95vw;
          overflow: scroll;
        }
        table {
          max-width: 95vw;
        }
        tbody > tr:nth-of-type(even) {
            background-color: rgb(237 238 242);
        }
        .loc {
            font-size: 0.8rem;
        }
        .syn-builtin, .syn-macro {
            color: purple;
            font-style: italic;
            font-weight: bold;
        }
        .syn-constant {
            color: olive;
        }
        .syn-procedure {
            color: teal;
        }
        td pre {
          display: inline
        }
        .syn-paren-0 { color: maroon; }
        .syn-paren-1 { color: olive; }
        .syn-paren-2 { color: green; }
        .syn-paren-3 { color: navy; }
        .syn-paren-4 { color: purple; }
        .stack-frame
        { display: inline-flex
        ; flex-direction: row
        ; column-gap: 0.5em
        }
      """
    body_ do
      details_ do
        summary_ "stack code"
        pre_ $ code_ do
          htmlData . runJalmotUnsafe . S.toData S.dataIso $ p
      ppTrace t

ppTrace :: Trace -> Html ()
ppTrace trace =
  table_ do
    thead_ $ tr_ do
      traverse_ (th_ [scope_ "col"])
        ["routine","next instruction","stack frame"]
    tbody_ do
      go trace
  where
    go :: Trace -> Html ()
    go = \case
      Step vm next -> ppVM vm >> go next
      StepToSuccess vm rs -> do
        tr_ [class_ "trace-result"] do
          td_ do
            details_ do
              summary_ "result"
              pre_ do
                code_ . toHtml . pShowNoColor $ vm
          td_ [colspan_ "2"] do
            sequence_ . intersperse " | " $ code_ . ppDatum <$> rs
      StepToFailure vm err -> do
        ppVM vm
        tr_ [class_ "trace-failure"] do
          td_ [colspan_ "3"] do
            details_ do
              summary_ "error"
              pre_ do
                samp_ do
                  fromString $ displayException err

ppVM :: VM -> Html ()
ppVM vm = do
  tr_ do
    td_ do
      details_ do
        summary_ do
          var_ [class_ "loc"] do
            vm ^. #debug . #activeRoutine . to ppDatum
        pre_ do
          code_ . toHtml . pShowNoColor $ vm
    td_ do
      code_ curi
    td_ do
      ppStack vm.stack
  where
    curi = vm ^?! failing (#code . _head . to ppDatum) (#tail . to ppDatum)

ppStack :: Stack -> Html ()
ppStack stk = do
  span_ [class_ "stack"] do
    stk ^.. each
      & fmap ppFrame
      & intersperse " | "
      & sequence_

ppFrame :: Frame -> Html ()
ppFrame frm = do
  span_ [class_ "stack-frame"] do
    sequence_ $ frm ^.. #locals . each . to ppDatum

ppData :: S.DataIso a => a -> Html ()
ppData = htmlData . runJalmotUnsafe . S.toData S.dataIso

ppDatum :: S.DatumIso a => a -> Html ()
ppDatum = htmlDatum . runJalmotUnsafe . S.toDatum S.datumIso

fac (n :: Int) = [stkP|
(define $start
  (push! $fac)
  (push! #{n})
  (tail-call 1))

(define $fac
  (load %n 0)
  (prim %x0 (zero? %n))
  (if %x0
      (then (push! 1)
            (return 1))
      (else (prim %x1 (- %n 1))
            (push! $fac-c0)
            (push! $fac)
            (push! %x1)
            (call 1))))

(define $fac-c0
  (pop! %x2)
  (pop! %n)
  (prim %x3 (* %n %x2))
  (push! %x3)
  (return 1))
|]