{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-orphans -Wno-unused-matches -Wno-missing-signatures #-}
{- HLINT ignore "Avoid lambda using `infix`" -}
module Gyehoek.ANF.Syntax
  ( Exp(..)
  , toANF
  , lower
  , wrapFunction
  , lowerProgram
  )
  where

import Data.Text (Text)
import Effectful
import Gyehoek.QBE qualified as QBE
import Data.List (List)
import Data.Text.IO qualified as TIO
import Control.Lens
import Data.Generics.Labels
import Data.Vector.Strict (Vector)
import Data.Function (fix)
import Effectful.Writer.Static.Local
import Gyehoek.Scheme.Syntax qualified as Lam
import Gyehoek.Scheme.Syntax (Name, Prim(..), Lit(..))
import Gyehoek.GenSym
import Control.Monad.Cont
import Data.Foldable
import Data.List.NonEmpty (NonEmpty((:|)))
import Data.List.NonEmpty qualified as NE
import Gyehoek.QBE (FuncDef(FuncDef))
import Data.Foldable1
import qualified Data.Text as T
import Data.String (fromString)
import Language.SexpGrammar as Sexp hiding (List, iso, encode, decode, traversed)
import Language.SexpGrammar.Generic
import GHC.Generics (Generic)
import Gyehoek.Sexp
import Control.Category
import Prelude hiding ((.), id)
import Data.InvertibleGrammar.Base qualified as IG
import Data.InvertibleGrammar.Base ((:-)((:-)))
import qualified Gyehoek.Sexp
import Control.Lens.Unsound
import qualified Data.Bits
import qualified GHC.IO.Encoding as T
import qualified Data.Text.Encoding as T
import Data.HashMap.Strict (HashMap)
import Effectful.State.Static.Local
import qualified Data.HashMap.Strict as HM


data Val
  = ValLit Lit
  | ValVar Name
  deriving (Show, Generic)

data Exp
  = ExpLetApply Name Val (List Val) Exp
  | ExpLetPrim Name (Prim Val) Exp
  | ExpBegin (List Exp)
  | ExpVal Val
  deriving (Show, Generic)



expandBindings
  -- | Match constructor. (an affine fold would be preferable to a
  --   prism here)
  :: Prism' e (lhs, rhs, e)
  -> e
  -> (List (lhs, rhs), e)
expandBindings p = go [] where
  go acc e =
    case e ^? p of
      Just (l,r,e') -> go ((l,r):acc) e'
      Nothing -> (acc, e)

collapseBindings
  :: Foldable f => AReview e (lhs, rhs, e) -> f (lhs, rhs)
  -> e -> e
collapseBindings p bs e = foldr (\(l,r) e' -> p # (l,r,e')) e bs

-- | Technically unlawful.
bindingTelescope
  :: Prism' e (lhs, rhs, e)
  -> Iso' e (List (lhs, rhs), e)
bindingTelescope p = iso
  (expandBindings p)
  (uncurry $ collapseBindings p)

foldLet
  :: Prism' Exp (lhs, rhs, Exp)
  -> Grammar
       Position
       (Exp :- NonEmpty (lhs, rhs) :- t)
       (Exp :- rhs :- lhs :- t)
foldLet p =
  IG.Iso
    (\(e :- ((l1,r1):|bs) :- t) ->
       collapseBindings p bs e :- r1 :- l1 :- t)
    (\(e :- r :- l :- t) ->
       let (bs,e') = expandBindings p e
       in e' :- ((l,r) :| bs) :- t)

instance SexpIso Val where
  sexpIso = match
    $ With (. sexpIso)
    $ With (. symbol)
    $ End

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

-- nonEmptyGrammar :: Grammar p (NonEmpty x :- t) (List x :- x :- t)
-- nonEmptyGrammar = IG.Iso
--   (\((x:|xs) :- t) -> xs :- x :- t)
--   (\(xs :- x :- t) -> (x:|xs) :- t)

instance SexpIso Exp where
  sexpIso = match
    $ With (. letapp)
    $ With (. letprim)
    $ With (\bgn -> bgn . list (el (sym "begin") >>> rest sexpIso))
    $ With (. sexpIso)
    $ End
    where
      letprim
        :: Grammar Position (Sexp :- t) (Exp :- (Prim Val :- (Text :- t)))
      letprim =
        Gyehoek.Sexp.let_ symbol (sexpIso @(Prim Val)) (sexpIso @Exp)
        >>> foldLet #ExpLetPrim
      letapp :: Grammar
        Position (Sexp :- t) (Exp :- List Val :- Val :- Text :- t)
      letapp =
        Gyehoek.Sexp.let_ symbol (sexpIso @(NonEmpty Val)) (sexpIso @Exp)
         >>> foldLet (#ExpLetApply
                      . iso (\(rhs,f,xs,e) -> (rhs, f:|xs, e))
                            (\(rhs,f:|xs,e) -> (rhs,f,xs,e)))
         >>> onTail nonEmptyGrammar



-- 뻘짓이어라
telescope :: Traversable t => t ((a -> r) -> r) -> (t a -> r) -> r
telescope = runCont . traverse cont

toANF'
  :: forall es. GenSym :> es
  => Lam.Exp
  -> (Val -> Eff es Exp)
  -> Eff es Exp

toANF' (Lam.ExpLit v) k = k . ValLit $ v

toANF' (Lam.ExpPrim p) k =
  telescope (toANF' <$> p) \p' -> do
    r <- gensym
    ExpLetPrim r p' <$> k (ValVar r)

toANF' (Lam.ExpApply f xs) k =
  telescope (toANF' <$> (f:|xs)) \(f':|xs') -> do
    r <- gensym
    ExpLetApply r f' xs' <$> k (ValVar r)

toANF' (Lam.ExpBegin xs) k = ExpBegin <$> traverse anf xs
  where
    anf x = toANF' x (pure . ExpVal)

toANF' (Lam.ExpLet xs e) k = _

toANF' e k = _

toANF e = toANF' e (pure . ExpVal)



expr =
  Lam.ExpPrim
    (PrimAdd
     (Lam.ExpPrim
      (PrimMul
       (Lam.ExpLit (LitInt 2))
       (Lam.ExpLit (LitInt 3))))
    (Lam.ExpLit (LitInt 4)))

expr2 =
  Lam.ExpBegin
    [ Lam.ExpPrim
      (PrimWrite
       (Lam.ExpPrim
         (PrimCons
          (Lam.ExpLit (LitInt 2))
          (Lam.ExpLit (LitInt 3)))))
    , Lam.ExpPrim
      (PrimWrite
       (Lam.ExpPrim
        (PrimMul
         (Lam.ExpLit (LitInt 5))
         (Lam.ExpLit (LitInt 4)))))
    ]



instance Semigroup QBE.Program where
  QBE.Program ts ds fs <> QBE.Program ts' ds' fs' =
    QBE.Program (ts <> ts') (ds <> ds') (fs <> fs')

instance Monoid QBE.Program where
  mempty :: QBE.Program
  mempty = QBE.Program mempty mempty mempty

funcdef
  :: QBE.Ident QBE.Global
  -> List QBE.Param -> NonEmpty QBE.Block -> FuncDef
funcdef name ps =
  QBE.FuncDef
  mempty
  (Just (QBE.AbiBaseTy QBE.Long))
  name Nothing ps QBE.NoVariadic

prims :: QBE.Program
prims = QBE.Program primtys mempty primfns where
  primtys =
    [ QBE.TypeDef "scm" Nothing
      [ (QBE.SubExtTy (QBE.BaseTy QBE.Long), Just 2) ]
    ]
  primfns = [ -- write
            -- , mkArith "plus" QBE.Add
            -- , mkArith "star" QBE.Mul
            -- , mkArith "_" QBE.Sub
            -- , mkArith "slash" (QBE.Div QBE.Signed)
            ]
  mkArith name bop =
    funcdef name
    [ QBE.Param (QBE.AbiBaseTy QBE.Long) "x"
    , QBE.Param (QBE.AbiBaseTy QBE.Long) "y"
    ]
    [ QBE.Block "start" []
      [ QBE.BinaryOp ("r" QBE.:= QBE.Long) bop
        (QBE.ValTemporary "x") (QBE.ValTemporary "y")
      ]
      (QBE.Ret (Just (QBE.ValTemporary "r")))
    ]

data BlockBuilder
  = Emit (Vector QBE.Inst) !BlockBuilder
  | Exit QBE.Jump
  deriving (Show)

instance Semigroup BlockBuilder where
  Emit a as <> bs = Emit a (as <> bs)
  Exit _ <> bs = bs

instance Each BlockBuilder BlockBuilder QBE.Inst QBE.Inst where
  each k (Emit is bb) = Emit <$> traverse k is <*> each k bb
  each k (Exit j) = pure (Exit j)

evalBlockBuilder :: BlockBuilder -> (Vector QBE.Inst, QBE.Jump)
evalBlockBuilder (Emit is bb) = evalBlockBuilder bb & _1 <>:~ is
evalBlockBuilder (Exit j)     = ([],j)

buildBlock :: QBE.Ident QBE.Label -> BlockBuilder -> QBE.Block
buildBlock n bb = QBE.Block n [] (is ^.. each) j
  where (is,j) = evalBlockBuilder bb

lowerName :: Name -> QBE.Ident t
lowerName = fromString . T.unpack

lowerInt' = QBE.ValConst . QBE.CInt . fromIntegral

lowerInt = QBE.ValConst . QBE.CInt
  . (Data.Bits..|. 2)
  . (Data.Bits..<<. 2)
  . fromIntegral

lowerString
  :: forall es. (GenSym :> es, State StringLiterals :> es)
  => Text -> (QBE.Val -> Eff es BlockBuilder) -> Eff es BlockBuilder
lowerString s k = do
  let len = lengthOf each $ T.encodeUtf8 s
  rawString <- getRawString
  r <- gensym
  Emit (alloc r rawString len) <$> k (QBE.ValTemporary r)
  where
    -- getRawString
    --   :: forall es. (GenSym :> es, State StringLiterals :> es)
    --   => Eff es _
    getRawString = do
      x <- get
      case x ^. at s of
        Just s' -> pure s'
        Nothing -> do r <- gensym
                      state \lits -> (r, HM.insert s r lits)
    alloc r rs len =
      [ QBE.Call
        (Just (r, QBE.AbiBaseTy QBE.Long))
        (QBE.ValGlobal "scm_from_utf8_string")
        Nothing
        [ QBE.Arg (QBE.AbiBaseTy QBE.Long) (QBE.ValGlobal rs)
        -- N.b. The C function declares this argument as size_t, which
        -- /is/ long on my system.
        , QBE.Arg (QBE.AbiBaseTy QBE.Long) (lowerInt' len)
        ]
        []
      ]

type StringLiterals = HashMap Text (QBE.Ident QBE.Global)

lowerVal
  :: forall es. (GenSym :> es, State StringLiterals :> es)
  => Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder

lowerVal (ValLit (LitInt n)) k = k . lowerInt $ n

lowerVal (ValLit (LitQuote (Lam.SexpSymbol s))) k =
  lowerString s \s' -> do
    r <- gensym
    Emit (intern r s') <$> k (QBE.ValTemporary r)
  where
    intern r s' =
      [ QBE.Call
        (Just (r, QBE.AbiBaseTy QBE.Long))
        (QBE.ValGlobal "scm_string_to_symbol")
        Nothing
        [ QBE.Arg (QBE.AbiBaseTy QBE.Long) s'
        ]
        []
      ]

lowerVal (ValLit (LitString s)) k = lowerString s k

lowerVal (ValLit _) k = error "todo"
lowerVal (ValVar x) k = k . QBE.ValTemporary . lowerName $ x

binaryPrim :: Prism' (Prim a) (QBE.BinaryOp, a, a)
binaryPrim = prism' up down where
  up (bop,a,b) = case bop of
    QBE.Add -> _
    QBE.Mul -> _
    _ -> _
  down = \case
    PrimAdd a b -> Just (QBE.Add,a,b)
    PrimMul a b -> Just (QBE.Mul,a,b)
    _ -> Nothing

lowerArithmetic :: QBE.Assignment -> Prim QBE.Val -> QBE.Inst
lowerArithmetic r p = QBE.BinaryOp r bop x y
  where
    (bop,x,y) = case p of
      PrimAdd a b -> (QBE.Add,a,b)
      PrimMul a b -> (QBE.Mul,a,b)
      _ -> _

sizeofScm :: Integral a => a
sizeofScm = 8

lowerCons
  :: (GenSym :> es, State StringLiterals :> es)
  => Name -> QBE.Val -> QBE.Val -> Exp
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
lowerCons r car cdr e k = do
  r1 <- gensym
  Emit (alloc <> initialise r1) <$> lower' e k
  where
    alloc = [ QBE.Call
              (Just (lowerName r, QBE.AbiBaseTy QBE.Long))
              (QBE.ValGlobal "GC_malloc")
              Nothing
              [ QBE.Arg
                (QBE.AbiBaseTy QBE.Long)
                (QBE.ValConst (QBE.CInt (sizeofScm * 2))) ]
              []
            ]
    initialise r1 =
      [ QBE.BinaryOp (r1 QBE.:= QBE.Long) QBE.Add
        (QBE.ValTemporary (lowerName r)) (QBE.ValConst (QBE.CInt 8))
      , QBE.Store (QBE.BaseTy QBE.Long) car (QBE.ValTemporary (lowerName r))
      , QBE.Store (QBE.BaseTy QBE.Long) cdr (QBE.ValTemporary r1)
      ]

smallIntHelper'
  :: GenSym :> es
  => QBE.Ident 'QBE.Temporary
  -> QBE.BinaryOp
  -> QBE.Val -> QBE.Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
smallIntHelper' r bop v1 v2 k = do
  Emit [ QBE.BinaryOp (r QBE.:= QBE.Long)
         bop v1 v2 ]
    <$> k (QBE.ValTemporary r)

smallIntHelper
  :: GenSym :> es
  => QBE.BinaryOp
  -> QBE.Val -> QBE.Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
smallIntHelper bop a b k = do
  r <- gensym
  smallIntHelper' r bop a b k

makeSmallInt'
  :: forall es. (GenSym :> es)
  => QBE.Ident 'QBE.Temporary
  -> QBE.Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
makeSmallInt' r n k =
  smallIntHelper QBE.Shl n (lowerInt' 2) \n' ->
    smallIntHelper' r QBE.Add n' (lowerInt' 2) k

makeSmallInt
  :: forall es. (GenSym :> es)
  => QBE.Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
makeSmallInt n k = do
  r <- gensym
  makeSmallInt' r n k

getSmallInt
  :: forall es. (GenSym :> es)
  => QBE.Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
getSmallInt n = smallIntHelper QBE.Shr n (lowerInt' 2)

lowerWrite 
  :: forall es. (GenSym :> es)
  => Name -> QBE.Val -> Exp
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
lowerWrite r x e k =
  Emit [ QBE.Call (Just (lowerName r, QBE.AbiBaseTy QBE.Long))
         (QBE.ValGlobal "scm_write") Nothing
         [QBE.Arg (QBE.AbiBaseTy QBE.Long) x]
         []
       ]
    <$> k (QBE.ValTemporary (lowerName r))

smallIntMask :: Integer
smallIntMask = 2 ^ (sizeofScm * 8) - 2

lowerCar
  :: (GenSym :> es, State StringLiterals :> es)
  => Name -> QBE.Val -> _
  -> (QBE.Val -> Eff es BlockBuilder) -> Eff es BlockBuilder
lowerCar r x e k = do
  Emit [ QBE.Load (lowerName r QBE.:= QBE.Long) QBE.Long x
       ]
    <$> lower' e k

lowerCdr
  :: (GenSym :> es, State StringLiterals :> es)
  => Name -> QBE.Val -> Exp
  -> (QBE.Val -> Eff es BlockBuilder) -> Eff es BlockBuilder
lowerCdr r x e k = do
  x1 <- gensym
  Emit [ QBE.BinaryOp (x1 QBE.:= QBE.Long)
         QBE.Add x (lowerInt' sizeofScm)
       , QBE.Load (lowerName r QBE.:= QBE.Long) QBE.Long
         (QBE.ValTemporary x1)
       ]
    <$> lower' e k

lowerNewline r k =
  Emit [ QBE.Call (Just (lowerName r, QBE.AbiBaseTy QBE.Long))
         (QBE.ValGlobal "scm_newline") Nothing
         []
         []
       ]
    <$> k (QBE.ValTemporary (lowerName r))

lowerPrim 
  :: forall es. (GenSym :> es, State StringLiterals :> es)
  => Name -> Prim Val -> Exp
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
lowerPrim r p e k =
  telescope (lowerVal <$> p) \case
    (preview binaryPrim -> Just (bop,a,b)) ->
      getSmallInt a \a' ->
        getSmallInt b \b' ->
          smallIntHelper bop a' b' \c ->
            makeSmallInt' (lowerName r) c \_ ->
              lower' e k
    PrimCons x y -> lowerCons r x y e k
    PrimCar x -> lowerCar r x e k
    PrimCdr x -> lowerCdr r x e k
    PrimWrite x -> lowerWrite r x e k
    PrimNewline -> lowerNewline r k

lower'
  :: forall es. (GenSym :> es, State StringLiterals :> es)
  => Exp
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder

lower' (ExpVal v) k = lowerVal v k

lower' (ExpLetPrim r p e) k = lowerPrim r p e k

lower' (ExpLetApply r f xs e) k =
  telescope (lowerVal @es <$> (f:|xs)) \(f':|xs') ->
    Emit [ QBE.Call
            (Just (lowerName r, QBE.AbiBaseTy QBE.Long))
            f'
            Nothing
            (QBE.Arg (QBE.AbiBaseTy QBE.Long) <$> xs')
            []
          ]
      <$> lower' e k

lower' (ExpBegin (x:xs)) k = fold1 <$> traverse low (x:|xs)
  where low e = lower' @es e (pure . Exit . QBE.Ret . Just)

lower' _ k = _

lower
  :: (GenSym :> es, State StringLiterals :> es)
  => QBE.Ident QBE.Label
  -> Exp
  -> Eff es QBE.Block
lower n e = buildBlock n <$> lower' e (pure . Exit . QBE.Ret . Just)

lowerStringLiterals =
  ifoldMapOf itraversed \s v ->
    [ QBE.DataDef [] v Nothing
      [QBE.FieldExtTy QBE.Byte [QBE.String (T.encodeUtf8 s)]]]

lowerProgram 
  :: (GenSym :> es, Traversable t)
  => t Exp -> Eff es QBE.Program
lowerProgram anfs =
  case toList anfs of
    -- hack for dev convenience: if there's only one expression, let
    -- it be the entry point.
    [e] -> do
      (b,stringLits) <- runState mempty . lower "start" $ e
      let f = wrapFunction @NonEmpty "main" [b]
          dataDefs = lowerStringLiterals stringLits
      pure $ QBE.Program [] dataDefs [f]
    _ -> do
      let low e = do
            bl <- gensym' "b"
            fl <- gensym' "f"
            b <- lower bl e
            pure $ wrapFunction @NonEmpty fl [b]
      (fs,stringLits) <- runState mempty $ traverse low anfs
      pure $ QBE.Program [] (lowerStringLiterals stringLits) (fs ^.. traversed)

wrapFunction
  :: Foldable1 t
  => QBE.Ident 'QBE.Global -> t QBE.Block -> QBE.FuncDef
wrapFunction l bs =
  QBE.FuncDef [QBE.Export]
    (Just (QBE.AbiBaseTy QBE.Word))
    l Nothing [] QBE.NoVariadic (toNonEmpty bs)

wrapProgram :: Foldable1 t => t QBE.Block -> QBE.Program
wrapProgram bs = prims <> QBE.Program [] [] [main] where
  main = QBE.FuncDef [QBE.Export]
    (Just (QBE.AbiBaseTy QBE.Word))
    "main" Nothing [] QBE.NoVariadic (toNonEmpty bs)