{-# LANGUAGE ViewPatterns #-}
module Gyehoek.Scheme.Expand
  (
  ) where

import Gyehoek.Sexp.Syntax
import Gyehoek.Sexp qualified as S
import Gyehoek.Prelude
import Gyehoek.Scheme.Syntax (Name, Lit(..))
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 (zip)
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(..))


data Scope = MkScope
  { identity :: Natural
  , bindings :: HashSet Name
  }
  deriving (Show, Generic)

data Syntax
  = SyntaxTrans Trans
  | SyntaxPrim Prim
  deriving (Show, Generic)

data Prim
  = PrimLambda
  | PrimLet
  | PrimIf
  deriving (Show, Generic)

data Trans = MkTrans
  { ellipsis :: Name
  , keywords :: List Name
  , rules :: List Rule
  }
  deriving (Show, Generic)

data Rule = MkRule
  { rhs :: Pat
  , lhs :: Tem
  }
  deriving (Show, Generic)

data Pat
  = PatWildcard
  | PatVar Name
  | PatList
    { init :: List Pat
    , ellipsis :: Maybe (List Pat)
    , tail :: Maybe Pat
    }
  -- | PatVec
  --   { init :: List Pat
  --   , ellipsis :: Maybe (List Pat)
  --   }
  deriving (Show, Generic)

data Tem
  -- | @(⟨element⟩ …)@
  -- @(⟨element⟩ ⟨element⟩ … . ⟨element⟩)@
  = TemList
    { init :: List El
    , tail :: Maybe Tem
    }
  -- | @(⟨ellipsis⟩ ⟨template⟩)@
  | TemTrail Tem
  | TemLit Lit
  | TemVar Name
  deriving (Show, Generic)

data El
  -- | @⟨template⟩ ⟨ellipsis⟩@
  = Ellipsis Tem
  -- | @⟨template⟩@
  | El Tem
  deriving (Show, Generic)



instance S.DatumIso Trans where
  datumIso = S.with \g ->
    S.listWithIndentation
      (NSpecial 1)
      ( S.el (S.symBuiltin "syntax-rules")
      >>> S.Iso (\(ctx:-t) -> ctx:-"...":-t) (\(ctx:-_:-t) -> ctx:-t)
      >>> S.el (S.list $ S.rest (S.datumIso @Name))
      >>> S.rest (S.datumIso @Rule)
      )
    >>> g

instance S.DatumIso Rule where
  datumIso = S.with \g ->
    S.list (S.el S.datumIso >>> S.el S.datumIso) >>> g

instance S.DatumIso Pat where
  datumIso = S.match
    $ S.With (S.sym "_" >>>)
    $ S.With (S.datumIso >>>)
    $ S.With (lst >>>)
    $ S.End
    where
      lst :: S.G (Datum :- t) (Maybe Pat :- Maybe (List Pat) :- List Pat :- t)
      lst = S.coproduct
        [ S.list $
            ellipsis
            >>> S.onTail (S.push Nothing (is _Nothing) (const mempty))
        , S.dottedList
            ellipsis
            ( S.datumIso @Pat
              >>> S.partialIso Just (maybe (Left mempty) Right) )
        ]
      ellipsis =
        S.restData split
            >>> S.onTail
              ( S.onHead (S.traversed . S.traversed . S.sealed $
                S.datumIso @Pat)
              >>> S.onTail (S.onHead . S.traversed . S.sealed $
                S.datumIso @Pat)
              )
      split
        :: forall t. S.G (List Datum :- t)
            (Maybe (List Datum) :- List Datum :- t)
      split = S.Iso
        (\(ps0:-t) ->
          let (ps,ell) = splitEllipsis ps0
          in ell :- ps :- t)
        (\(ell:-ps:-t) -> (ps ++ foldMap ([Symbol "..."]++) ell) :- t)

splitEllipsis :: List Datum -> (List Datum, Maybe (List Datum))
splitEllipsis [] = ([], Nothing)
splitEllipsis (Symbol "..." : xs) = ([], Just xs)
splitEllipsis (x:xs) = splitEllipsis xs & _1 %~ (x:)

instance S.DataIso El where
  dataIso = S.match
    $ S.With (ellipsis >>>)
    $ S.With (noellipsis >>>)
    $ S.End
    where
      ellipsis = S.recontextualise $
        S.el (S.datumIso @Tem) >>> S.el (S.sym "...")
      noellipsis = S.recontextualise $ S.el (S.datumIso @Tem)

instance S.DatumIso Tem where
  datumIso = S.match
    $ S.With (lst >>>)
    $ S.With (trail >>>)
    $ S.With (S.datumIso @Lit >>>)
    $ S.With (S.datumIso @Name >>>)
    $ S.End
    where
      trail = S.list $ S.el (S.sym "...") >>> S.el (S.datumIso @Tem)

      lst :: S.G (S.Datum :- t) (Maybe Tem :- List El :- t)
      lst = S.coproduct
        [ S.list els
          >>> (S.push Nothing (is _Nothing) (const mempty))
        , S.dottedList els $
            S.datumIso @Tem
            >>> S.partialIso Just (maybe (Left mempty) Right)
        ]
      els :: S.G (S.ListContext :- t) (S.ListContext :- List El :- t)
      els =
        S.iso
          (\(S.MkListContext ds) -> affixEllipses ds)
          (S.MkListContext . foldMap \(d,b) ->
            d : if b then [Symbol "..."] else [])
        >>> S.onHead (S.traversed . S.sealed $
          S.flipped S.pair
          >>> S.onTail (S.datumIso @Tem)
          >>> S.pair
          >>> S.iso
            (\(t,b) -> if b then Ellipsis t else El t)
            (\case
              Ellipsis t -> (t,True)
              El t -> (t,False)))
        >>> S.push (S.MkListContext [])
          (\(S.MkListContext xs) -> null xs)
          (const mempty)
      affixEllipses :: List Datum -> List (Datum, Bool)
      affixEllipses (x : Symbol "..." : xs) = (x,True) : affixEllipses xs
      affixEllipses (x : xs) = (x,False) : affixEllipses xs
      affixEllipses [] = []



-- |
-- >>> match (PatList [PatVar "x"] Nothing (Just (PatVar "y"))) [S.sx|(1 2 . 2)|]
-- Nothing
-- >>> match (PatList [PatVar "x"] Nothing (Just (PatVar "y"))) [S.sx|(1 . 2)|]
-- Just
-- ...
match :: Pat -> Datum -> Maybe (HashMap Name Datum)
match p = getAp . match' p
  
match' :: Pat -> Datum -> Ap Maybe (HashMap Name Datum)

match' PatWildcard _ = pure mempty

match' (PatVar x) e = 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



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)
    ]
  }