{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE TypeFamilies #-}
module Gyehoek.Scheme.Expand
  (
  ) where

import Gyehoek.Sexp.Syntax
import Gyehoek.Sexp qualified as S
import Gyehoek.Prelude
import Gyehoek.Scheme.Syntax hiding (Prim(..))
import Gyehoek.Scheme.Syntax qualified as Scm
import qualified Data.HashSet as HS
import qualified Data.HashMap.Strict as H
import Data.Foldable
import Data.These
import Data.Zip
import Prelude hiding (filter, zip, mapMaybe)
import qualified Data.List.NonEmpty as NE
import Data.Monoid (Ap(Ap, getAp))
import Gyehoek.Sexp.Grammar.Base ((:-)(..))
import Control.Lens.Extras (is)
import Control.Applicative (Alternative(..))
import Gyehoek.GenSym
import Data.HashSet.Lens (setOf)
import Data.List (sort)
import GHC.Exts (IsList(..))
import Gyehoek.Jalmot
import Data.Maybe (isNothing)
import Data.Kind (Type)
import Data.Traversable (for, mapAccumR)
import Effectful.State.Dynamic
import Witherable
import Data.Semigroup (Arg(..))
import Data.Ord (Down(..))
import qualified Data.Scientific as Sci


data Bind
  = BindLexical { symbol :: Name, identity :: Natural }
  | BindGlobal { symbol :: Name }
  deriving stock (Generic, Eq, Show)
  deriving anyclass (Hashable)

newtype Scope = MkScope { identity :: Natural }
  deriving stock (Show, Generic, Eq, Ord)
  deriving anyclass (Hashable)
  deriving newtype (Gen)

type ScopeSet = HashSet Scope

data Formals
  = FormalsFixed (List Name)
  | FormalsRest (List Name) Name
  deriving (Show, Generic)

data PrimLambda e = MkPrimLambda Formals (List e)
  deriving (Show, Generic, Functor, Foldable, Traversable)
data PrimLet e = MkPrimLet (Maybe Name) (List (Name, e)) (List e)
  deriving (Show, Generic, Functor, Foldable, Traversable)
data PrimIf e = MkPrimIf e e (Maybe e)
  deriving (Show, Generic, Functor, Foldable, Traversable)
data PrimLetSyntax e = MkPrimLetSyntax (List (Name, e)) (List e)
  deriving (Show, Generic, Functor, Foldable, Traversable)

data Prim e
  = PrimLambda (PrimLambda e)
  | PrimLet (PrimLet e)
  | PrimIf (PrimIf e)
  | PrimLetSyntax (PrimLetSyntax e)
  | PrimSyntaxRules Trans
  deriving (Show, Generic)

data Key = MkKey Name (HashSet Scope)
  deriving stock (Show, Generic, Eq)
  deriving anyclass (Hashable)



instance S.DatumIso Scope where
  datumIso = S.with (S.datumIso >>>)

hashSetGrammar
  :: forall a. (Hashable a, Ord a)
  => S.DatumGrammar a -> S.DatumGrammar (HashSet a)
hashSetGrammar g =
  S.list (S.rest g)
  >>> S.iso HS.fromList (sort . HS.toList)

instance S.DatumIso Key where
  datumIso = S.with \g ->
    S.list
      ( S.el (S.symBuiltin "@")
      >>> S.el (S.datumIso @Name)
      >>> S.el (hashSetGrammar S.datumIso)
      )
    >>> g

instance S.DatumIso Formals where
  datumIso = S.match
    $ S.With (fixed >>>)
    $ S.With (rest >>>)
    $ S.End
    where
      fixed :: S.G (Datum :- t) (List Name :- t)
      fixed = S.list $ S.rest $ S.datumIso @Name
      rest :: S.G (Datum :- t) (Name :- List Name :- t)
      rest = S.coproduct
        [ S.dottedList (S.rest $ S.datumIso @Name) (S.datumIso @Name)
        , S.datumIso @Name >>> S.onTail (S.push [] null (const mempty))
        ]

optEl
  :: S.G (S.Datum :- t) (a :- t)
  -> S.G (S.ListContext :- t) (S.ListContext :- Maybe a :- t)
optEl g =
  S.coproduct
    [ S.el g >>> S.onTail (S.partialIso Just \case
      Nothing -> Left mempty
      Just x -> Right x)
    , S.onTail $ S.push Nothing isNothing (const mempty)
    ]

anykw :: Text -> S.G (Datum :- t) t
anykw s = S.Flip $ S.PartialIso
  (\t -> adorn SynBuiltin (Symbol s) :- t)
  \case
    (Symbol _ :- t) -> Right t
    _ -> Left $ S.expected "symbol"

prim_if :: S.DatumIso e => S.DatumGrammar (PrimIf e)
prim_if = S.with \g ->
  S.listWithIndentation (NSpecial 1)
    (S.el (anykw "if")
    >>> S.el S.datumIso
    >>> S.el S.datumIso
    >>> optEl S.datumIso)
  >>> g

prim_lambda :: S.DatumIso e => S.DatumGrammar (PrimLambda e)
prim_lambda = S.with \g ->
  S.lambdaLike (anykw "λ") (S.datumIso @Formals) (S.rest S.datumIso)
  >>> g

prim_let
  :: forall e. S.DatumIso e => S.DatumGrammar (PrimLet e)
prim_let = S.with \g ->
  (S.listWithIndentation (NSpecial 1) $
    S.el (anykw "let")
    >>> optEl (S.datumIso @Name)
    >>> S.el (S.list $ S.rest $ S.datumIso @(Name,e))
    >>> S.rest S.datumIso)
  >>> g

prim_let_syntax :: S.DatumIso e => S.DatumGrammar (PrimLetSyntax e)
prim_let_syntax = S.with \g ->
  (S.listWithIndentation (NSpecial 1) $
    S.el (anykw "let-syntax")
    >>> S.el (S.list $ S.rest $ S.datumIso)
    >>> S.rest S.datumIso)
  >>> g

instance S.DatumIso Bind where
  datumIso = S.match
    $ S.With (\g ->
      S.list (S.el (S.symBuiltin "L") >>> S.el S.datumIso >>> S.el S.datumIso)
      >>> g)
    $ S.With (\g ->
      S.list (S.el (S.symBuiltin "G") >>> S.el S.datumIso)
      >>> g)
    $ S.End



{- |

* Examples

>>> :set -XTemplateHaskellQuotes
>>> pat = S.makeSx [|| S.fromDatumUnsafe @Pat S.datumIso ||]
>>> match [] [pat|(x . y)|] [S.sx|(1 2 . 2)|]
Nothing
>>> match [] [pat|(x . y)|] [S.sx|(1 . 2)|]
Just
...
>>> match ["=>"] [pat|(x => y)|] [S.sx|(1 => 2)|]
Just
...
>>> match ["=>"] [pat|(x => y)|] [S.sx|(1 -> 2)|]
Nothing
-}
match :: Foldable f
  => f Text
  -- ^ Literal keywords
  -> Pat
  -> Datum
  -> Maybe (HashMap Name Datum)
match kws p = getAp . match' p where

  match' :: Pat -> Datum -> Ap Maybe (HashMap Name Datum)
  
  match' PatWildcard _ = pure mempty
  
  match' (PatVar x) e
    | coerce x `elem` kws = case e of
      Symbol x' | coerce x == x' -> pure mempty
      _                          -> empty
    | otherwise = pure $ H.singleton x e
  
  match' (PatList ps Nothing Nothing) (List es) = matches ps es
  
  match' (PatList ps (Just []) Nothing) (List es) = do
    let (ps',p) = ps ^?! _Snoc
    (r,rest) <- fold $ alignWith f ps' es
    rest' <- rest
      & fmap (fmap (fmap (:[])) . match' p)
      & foldr (liftA2 $ H.unionWith (<>)) mempty
      & fmap (fmap List)
    pure $ r <> rest'
    where
      f (These a b) = (,[]) <$> match' a b
      f (This a) = empty
      f (That b) = pure (mempty,[b])
  
  match' (PatList ps Nothing (Just p)) (DotList es e) =
    matches (p:|ps) (NE.cons e es)
  
  match' _ _ = _
  
  matchPrefix
    :: (Semialign f, Foldable f)
    => f Pat -> f Datum -> Ap Maybe (HashMap Name Datum, List Datum)
  matchPrefix ps es = fold $ alignWith f ps es
    where
      f (These a b) = (,[]) <$> match' a b
      f (This a) = empty
      f (That b) = pure (mempty,[b])
  
  matches
    :: (Semialign f, Foldable f)
    => f Pat -> f Datum -> Ap Maybe (HashMap Name Datum)
  matches ps es = fold $ alignWith matchThese ps es
  
  matchThese (These a b) = match' a b
  matchThese _ = empty



-- | this hashmap is "curried" since we often want to traverse the
-- entire set of scopes associated with a given symbol.
newtype SymTable = MkSymTable
  { curried :: HashMap Name (HashMap (HashSet Scope) Bind) }
  deriving (Show, Generic)

type instance Index SymTable = Name
type instance IxValue SymTable = HashMap (HashSet Scope) Bind
instance Ixed SymTable where ix j = #curried . ix j
instance At SymTable where at j = #curried . at j

instance Semigroup SymTable where
  MkSymTable c1 <> MkSymTable c2 = MkSymTable $ H.unionWith (<>) c1 c2

instance Monoid SymTable where mempty = MkSymTable mempty

type Expand es = (State SymTable :> es, GenSym :> es)

runExpand :: SymTable -> Eff (State SymTable : GenSym : es) a -> Eff es (a, SymTable)
runExpand tab = runGenSym . runStateLocal tab

testExpand :: Datum -> IO (CommandOrDef, SymTable)
testExpand = runJalmotIO . runExpand (symTableOfEnv env_scheme_base)
  . expand env_scheme_base mempty

-- | denotations
data Denot
  = DenotVar
  | DenotMacro Trans
  | DenotPrim Name
  | DenotSyntax Datum
  | DenotKeyword Name
  deriving stock (Show, Generic)

newtype Env = MkEnv { names :: HashMap Bind Denot }
  deriving stock (Show, Generic)
  deriving newtype (Semigroup, Monoid)

-- | @(environment '(scheme base))@
env_scheme_base :: Env
env_scheme_base = MkEnv . fromList . fold $
  [ [ (BindGlobal p, DenotPrim p) | p <- prims ]
  , [ (BindGlobal "λ", DenotPrim "lambda")
    ]
  ]
  where
    prims =
      [ "if"
      , "lambda"
      , "let"
      , "let-syntax"
      ]

symTableOfEnv :: Env -> SymTable
symTableOfEnv = ifoldMapOf (#names . itraversed) \cases
  b@(BindGlobal n) d -> binding n mempty b

type instance IxValue Env = Denot
type instance Index Env = Bind

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

err :: Jalmot :> es => Text -> Eff es a
err = throwError . EvalError

run :: S.DatumGrammar a -> Datum -> Maybe a
run g = preview #_Right . runPureEff . runJalmot . S.fromDatum g

parsePrim :: S.DatumIso e => Name -> Datum -> Maybe (Prim e)
parsePrim primname d = case primname of
  "let" -> PrimLet <$> run prim_let d
  "let-syntax" -> PrimLetSyntax <$> run prim_let_syntax d
  "lambda" -> PrimLambda <$> run prim_lambda d
  "if" -> PrimIf <$> run prim_if d

emit :: (Monoid m, State m :> es) => m -> Eff es ()
emit x = modify (<> x)

binding :: Name -> HashSet Scope -> Bind -> SymTable
binding sym scopes = MkSymTable . H.singleton sym . H.singleton scopes

envOfVar :: Bind -> Env
envOfVar = MkEnv . flip H.singleton DenotVar

gensymLexical :: GenSym :> es => Name -> Eff es Bind
gensymLexical symbol = do
  identity <- gensym
  pure $ BindLexical {symbol,identity}

lookupSymbol
  :: forall es. (State SymTable :> es, Jalmot :> es)
  => Env -> HashSet Scope -> Name -> Eff es (Bind, Denot)
lookupSymbol g scopes x = do
  ss <- fmap fold . preuse @SymTable @(Eff es) $ ix x
  case nearest (MkKey x scopes) (H.keys ss) of
    [] -> err [i|심벌 #{x}는 정의되지 않다|]
    (_:_:_) -> err [i|심벌 #{x}는 모호하다|]
    [s] | Just b <- ss ^? ix s
        , Just d <- g ^? ix b -> pure (b,d)
        | otherwise -> error "unreachable"

{- | Given a 'Key' and a collection of 'ScopeSet's, filter that
collection down to the largest subsets of the 'Key'\'s scope set.
This is our analogue of the lexical scoping rule which chooses the
"nearest" binding of a variable when shadowing occurs.

  - If no qualifying subsets are found, the name is not in scope.
  - If one subset is found, we're on the happy path!
  - If more than one subset is found, we're on the rarest and
    saddest path: the reference is ambiguous.

* Examples

> (let ((x {A} 123))
>   (λ (x {A,B})
>     (let ((y {A,B,C} 456))
>       x {A,B,C})))
    
>>> :seti -XOverloadedLists
>>> :{
nearest
  (MkKey "x" [MkScope 0, MkScope 1, MkScope 2])
  [ [MkScope 0]
  , [MkScope 0, MkScope 1] ]
:}
-}
nearest :: (Traversable f, Filterable f) => Key -> f ScopeSet -> List ScopeSet
nearest (MkKey symbol scopes) =
  mapMaybe (\x ->
    if x `HS.isSubsetOf` scopes
    then Just . Down $ Arg (length x) x
    else Nothing)
  -- 나쁨. 안 좋다. 안 좋아하야.
  >>> Data.Foldable.toList >>> sort
  >>> foldr (\cases
    x [] -> [x]
    x acc@(y:_) -> case x `compare` y of
      LT -> acc
      EQ -> x:acc
      GT -> [x])
    []
  >>> fmap (\(Down (Arg _ x)) -> x)

expand
  :: (Expand es, Jalmot :> es)
  => Env -> ScopeSet -> Datum -> Eff es CommandOrDef

expand g scopes datum@(List (Symbol s : es)) = do
  let s' = MkName s
  (_,denot) <- lookupSymbol g scopes s'
  case denot of
    DenotVar -> Command . ExpApply (ExpVar s')
      <$> traverse (expandAsExp g scopes) es
    DenotPrim x -> expandPrim g scopes x datum
    DenotMacro trans -> expandMacro g scopes trans datum

-- 신기하지 않은 경우들
expand g scopes datum = case datum of
  List (x:xs) -> Command <$> (ExpApply <$> go x <*> traverse go xs)
  Symbol s -> pure . Command . ExpVar . MkName $ s
  Boolean b -> lit $ LitBool b
  Number (Sci.floatingOrInteger -> Right n) -> lit $ LitInt n
  where
    go = expandAsExp g scopes
    lit = pure . Command . ExpLit

expandMacro
  :: (Expand es, Jalmot :> es)
  => Env -> ScopeSet -> Trans -> Datum -> Eff es CommandOrDef
expandMacro g scopes trans datum = _

expandPrim 
  :: (Expand es, Jalmot :> es)
  => Env -> ScopeSet -> Name -> Datum -> Eff es CommandOrDef
expandPrim g scopes primName datum
  | Just prim <- parsePrim @Datum primName datum
  = case prim of
    PrimLetSyntax (MkPrimLetSyntax bs [body]) -> do
      scopes' <- flip HS.insert scopes <$> gensym
      (rhss,g') <- (_2 %~ (g<>) . fold) . Prelude.unzip <$> for bs \(x,trans) ->
        expandAsExp g scopes trans >>= \case
          ExpSyntaxRules trans' ->
            (trans',) <$> bindLexical scopes' x (DenotMacro trans')
          e -> err [i|syntax-rules를 원하는데 이것 받는다: #{e}|]
      Command <$> (ExpLetSyntax
        (zip (bs ^.. each . _1) rhss)
        <$> expandAsExp g' scopes' body
        )
    PrimLet (MkPrimLet Nothing bs [body]) -> do
      scopes' <- flip HS.insert scopes <$> gensym
      g' <- (g<>) . fold <$> for (bs ^.. each . _1) \x ->
        bindLexical scopes' x DenotVar
      Command <$> (ExpLet
        <$> traverseOf (each . _2) (expandAsExp g scopes) bs
        <*> expandAsExp g scopes' body)
    PrimLambda (MkPrimLambda (FormalsFixed xs) [body]) -> do
      scopes' <- flip HS.insert scopes <$> gensym
      g' <- (g<>) . fold <$> for xs \x -> bindLexical scopes' x DenotVar
      Command . ExpLambda xs <$> expandAsExp g' scopes' body
    PrimIf (MkPrimIf c t f) ->
      Command <$> (ExpIf
        <$> expandAsExp g scopes c
        <*> expandAsExp g scopes t
        <*> traverse (expandAsExp g scopes) f
        )
  | otherwise = err [i|prim #{primName}에 잘못한 신택스: #{datum}|]

expandAsExp
  :: (Expand es, Jalmot :> es)
  => Env -> ScopeSet -> Datum -> Eff es Exp
expandAsExp g scopes d = expand g scopes d >>= intoExp

bindLexical :: Expand es => ScopeSet -> Name -> Denot -> Eff es Env
bindLexical scopes symbol denot = do
  identity <- gensym
  let bind = BindLexical {symbol,identity}
  modify (<> binding symbol scopes bind)
  pure $ MkEnv (H.singleton bind denot)

intoExp :: Jalmot :> es => CommandOrDef -> Eff es Exp
intoExp (Command e) = pure e
intoExp (Begin es) = traverse intoExp es >>= \case
  [] -> err "begin expression은 빔"
  (x:xs) -> pure . ExpBegin $ x NE.:| xs



trans_when :: Trans
trans_when = MkTrans
  { ellipsis = "..."
  , keywords = []
  , rules =
    [ MkRule
      (PatList
        [ PatVar "when"
        , PatVar "test"
        , PatVar "body" ]
        (Just [])
        Nothing)
      (TemList
        [ El $ TemVar "if"
        , El $ TemVar "test"
        , El $ TemList
          [ El $ TemVar "begin"
          , Ellipsis $ TemVar "body"
          ]
          Nothing
        ]
        Nothing)
    ]
  }

trans_and :: Trans
trans_and = MkTrans
  { ellipsis = "..."
  , keywords = []
  , rules =
    [ MkRule
        (PatList
          [PatVar "and"]
          Nothing
          Nothing)
        (TemLit (LitBool True))
    , MkRule
        (PatList
          [PatVar "and", PatVar "x"]
          Nothing
          Nothing)
        (TemVar "x")
    , MkRule
        (PatList
          [PatVar "and", PatVar "x", PatVar "y"]
          (Just [])
          Nothing)
        (TemList
          [ El (TemVar "if")
          , El (TemVar "x")
          , El (TemList
              [ El (TemVar "and")
              , Ellipsis (TemVar "y")
              ]
              Nothing)
          , El . TemLit . LitBool $ False
          ]
          Nothing)
    ]
  }