{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedLists #-}
module Gyehoek.CPS.Eval
  ( evalProgram
  , module Gyehoek.CPS.Syntax
  , evalExp
  , eGrammar
  ) where

import Gyehoek.CPS.Syntax hiding (Hob(..), Obj(..), cont)
import Gyehoek.Sexp qualified as S
import Control.Lens hiding (assign)
import Data.Maybe (fromMaybe, isJust)
import Text.Show.Functions ()
import qualified Data.HashMap.Strict as H
import Gyehoek.Prelude hiding (assign)
import Debug.Pretty.Simple
import Gyehoek.Jalmot
import Control.Monad.Cont
import Gyehoek.Sexp qualified as S
import GHC.Generics (Generically(..))
import Gyehoek.Sexp ((:-)(..))
import Data.List (nub, mapAccumR, compareLength)
import Data.HashSet.Lens (setOf)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IM
import Data.Monoid
import Control.Monad.State
import Data.Traversable (for)
import Data.Foldable (traverse_, foldrM)


newtype Loc = MkLoc { getLoc :: Int }
  deriving stock (Generic, Data)
  deriving newtype (Show, Eq, Ord, Enum)

data Store = MkStore
  { nextLoc :: Loc
  , heap :: IntMap E
  }
  deriving stock (Show, Generic)

type instance Index Store = Loc
type instance IxValue Store = E

instance Ixed Store where ix (MkLoc j) = #heap . ix j
instance At Store where at (MkLoc j) = #heap . at j

emptyStore :: Store
emptyStore = MkStore
  { nextLoc = MkLoc 0
  , heap = mempty
  }

newtype Env = MkEnv { getEnv :: HashMap Name Loc }
  deriving stock (Show, Generic, Data)
  deriving newtype (Semigroup, Monoid)

emptyEnv :: Env
emptyEnv = mempty

type instance Index Env = Name
type instance IxValue Env = Loc

instance Ixed Env where ix j = #getEnv . ix j
instance At Env where at j = #getEnv . at j

update :: Loc -> E -> Store -> Store
update (MkLoc loc) v = #heap %~ IM.alter f loc
  where
    f (Just _) = Just v
    f Nothing = error "segfault lol"

updates :: Foldable f => f (Loc, E) -> Store -> Store
updates = alaf Endo foldMap (uncurry update)

fetch :: Loc -> M r E
fetch (MkLoc loc) = gets (^?! #heap . ix loc)

new :: M r Loc
new = state \st -> (st.nextLoc, st & #nextLoc %~ succ)

new' :: E -> M r Loc
new' e = state \st ->
  ( st.nextLoc
  , st & #nextLoc %~ succ & at st.nextLoc ?~ e
  )

defines :: Traversable t => t (Name, E) -> M Answer Env
defines = alaf Ap foldMap \(name,e) -> do
  l <- new' e
  pure $ bind name l

var :: HasCallStack => Env -> Name -> M Answer Loc
var g x = case g ^. at x of
  Just l -> pure l
  Nothing -> wrong [i|unbound variable #{x}|]

type CmdCont = Store -> Answer
type ExpCont = List E -> CmdCont

type M r = ContT r (State Store)

data Answer
  = AnswerValues (List E)
  | AnswerError AJalmot
  deriving (Show, Generic)

data Mutability
  = Mut
  | NoMut
  deriving (Show, Generic, Data, Eq)

wrong :: Text -> M Answer a
wrong s = ContT \_ -> pure . AnswerError . EvalError $ s

orWrong
  :: Getting (First a) s a
  -> Text -> s -> M Answer a
orWrong p msg s = case getFirst . getConst $ p (Const . First . Just) s of
  Nothing -> wrong msg
  Just x -> pure x

bind :: Name -> Loc -> Env
bind k = MkEnv . H.singleton k

extends :: Foldable f => f (Name, Loc) -> Env -> Env
extends xs g = g <> foldMap (uncurry bind) xs

assign :: Loc -> E -> M Answer ()
assign l e = do
  use (at l) >>= \case
    Just _ -> at l ?= e
    Nothing -> wrong [i|#{e}에서 #{l}이라는 주소는 없다|]



-- | The denotation of an expressed value.
data E
  = ESymbol Text
  | ECharacter Char
  | EInt Int
  | EBool Bool
  | EUndefined
  | EUnspecified
  | ENull
  | EPair Mutability Loc Loc
  | EVec Mutability (List Loc)
  | EString Mutability (List Loc)
  | EProcedure Procedure
  deriving stock (Show, Generic)

type Procedure = List E -> DynPoints -> M Answer (List E)

eGrammar :: Store -> S.DatumGrammar E
eGrammar st = S.partialOsi (const . Left $ mempty) go
  where
    gofetch x = go $ st ^?! ix x
    go = \case
      ESymbol s -> S.Symbol s
      ECharacter c -> S.Character c
      EInt n -> S.Number (fromIntegral n)
      EBool b -> S.Boolean b
      EUndefined -> S.Unreadable "#<undefined>"
      EUnspecified -> S.Unreadable "#<unspecified>"
      EProcedure _ -> S.Unreadable "#<procedure>"
      ENull -> S.List []
      EPair _mut car cdr -> S.DotList [gofetch car] (gofetch cdr)
      EVec _mut xs -> S.Vector . fmap gofetch $ xs
      EString _mut xs -> S.String _

data DynPoints = MkDynPoints
  deriving (Generic, Data)



truthy :: E -> Bool
truthy (EBool False) = False
truthy _ = True



evalVal :: Env -> Val -> M Answer E

evalVal g (ValVar x) = var g x >>= fetch

evalVal g (ValImm imm) = case imm of
  ImmLabel (MkLabel l) -> var g l >>= fetch
  ImmInt n -> pure $ EInt n
  ImmBool b -> pure $ EBool b
  ImmUndefined -> pure EUndefined

evalKexp :: Env -> Kexp -> M Answer E
evalKexp g (KexpVar x) = var g x >>= fetch
evalKexp g (KexpKappa kap) = evalAbs g (AbsKappa kap)

evalAbs :: Env -> Abs -> M Answer E
evalAbs g (MkAbs formals ktail e) = pure . EProcedure $ \xs dps ->
  let
    formals' = formals ++ foldMap (:[]) ktail
    lformals = length formals'
    lxs = length xs
  in if lformals /= lxs
  then wrong [i|함수는 #{lformals}개의 인자를 필요로 하는데 #{lxs}개 받았다.|]
  else do
    ls <- xs & traverse new'
    let g' = g & extends (zip formals' ls)
    eval g' dps e

eval :: Env -> DynPoints -> Exp -> M Answer (List E)

eval g dps (ExpJump f xs ktail) = do
  f' <- evalVal g f
  xs' <- traverse (evalVal g) xs
  ktail' <- traverse (evalKexp g) (ktail ^.. _Just)
  case f' of
    EProcedure p -> p (xs' ++ ktail') dps
    _ -> wrong "bad procedure"

eval g dps (ExpLetRec bs e) = do
  ls <- for bs . const $ new' EUndefined
  let g' = g & extends (zip (bs ^.. each . _1) ls)
  bs' <- forOf (each . _2) bs (evalAbs g')
  traverse_ (uncurry assign) $ zip ls (bs' ^.. each . _2)
  eval g' dps e

eval g dps (ExpPrim (PrimCallCC withcc) k) = do
  withcc' <- evalVal g withcc >>= orWrong #_EProcedure
    [i|call/cc: 함수가 아닌 것을 받았다|]
  k' <- evalKexp g k
  kproc <- orWrong #_EProcedure [i|call/cc: 몰라...|] k'
  let cc = EProcedure \xs dps -> case unsnoc xs of
        Just (xs',_) -> kproc xs' dps
        Nothing -> wrong [i|call/cc: 잘못하는데!|]
  withcc' [cc,k'] dps

eval g dps (ExpPrim p k) = do
  p' <- evalPrim g dps =<< traverse (evalVal g) p
  evalKexp g k >>= \case
    EProcedure fp -> fp p' dps
    _ -> wrong [i|prim(#{p})의 계속을 나쁘다|]

eval g dps (ExpIf c t f) = do
  c' <- evalVal g c
  let b = if truthy c' then t else f
  var g b >>= fetch >>= \case
    EProcedure fp -> fp [] dps
    _ -> wrong [i|if의 계속을 나쁘다|]

eval g dps e = error [i|unimplemented #{e}|]

evalPrim :: Env -> DynPoints -> Prim E -> M Answer (List E)
evalPrim g dps = \case
  PrimAdd x y -> arith2 (+) x y
  PrimMul x y -> arith2 (*) x y
  PrimSub x y -> arith2 (-) x y
  PrimDiv x y -> arith2 div x y
  PrimZeroP x -> pure1 . EBool . isJust $ x ^? #EInt . only 0
  PrimCons x y -> pcons x y >>= pure1
  PrimCar p -> cr p _2
  PrimCdr p -> cr p _3
  PrimPairP p -> pure1 . EBool . maybe False (const True) $
    p ^? #_EPair
  PrimValues xs -> pure xs
  PrimList xs -> foldrM pcons ENull xs >>= pure1
  p -> wrong [i|prim(#{p})은 벌써 나지 않다|]
  where
    pure1 x = pure [x]
    pcons x y = do
      (x',y') <- traverseOf both new' (x,y)
      pure $ EPair Mut x' y'
    arith2 f (EInt x) (EInt y) = pure [EInt $ f x y]
    arith2 f x y = wrong [i|나쁜 인자: #{x}, #{y}|]
    cr p l = orWrong (#_EPair . l) [i|car/cdr는 pair을 받지 않다|] p
      >>= fmap (:[]) . fetch



evalExp :: Jalmot :> es => Exp -> Eff es _
evalExp e = _

evalProgram :: Jalmot :> es => Program -> Eff es (List S.Datum)
evalProgram (MkProgram lam) = case run (pure . AnswerValues) of
  (AnswerError jm, _) -> throwError jm
  (AnswerValues vs, st) -> traverse (S.toDatum $ eGrammar st) vs
  where
    run f = (`runState` emptyStore) . (`runContT` f) $ do
      g <- setup
      eval g MkDynPoints (ExpLetRec
        [("_start",AbsLambda lam)]
        (ExpApply (ValVar "_start") [] (KexpVar "halt")))



setup :: M Answer Env
setup = defines @List
  [ ("halt", EProcedure prim_halt)
  ]

prim_halt :: Procedure
prim_halt xs _dps = ContT \_ -> pure $ AnswerValues xs