{-# 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


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' 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
  . (Data.Bits..|. 2)
  . (Data.Bits..<<. 2) . fromIntegral

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

lowerVal (ValLit (LitInt n)) k = k . lowerInt $ n
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 = 16

lowerCons
  :: (GenSym :> 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)) ]
              []
            ]
    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.Val -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
smallIntHelper bop v k = do
  r <- gensym
  Emit [ QBE.BinaryOp (r QBE.:= QBE.Long)
         bop v (QBE.ValConst (QBE.CInt 2)) ]
    <$> k (QBE.ValTemporary r)

makeSmallInt
  :: forall es. (GenSym :> es)
  => QBE.Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder
makeSmallInt n k =
  smallIntHelper QBE.Shl n \n' ->
    smallIntHelper QBE.And n' k

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

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

lowerPrim 
  :: forall es. (GenSym :> 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)) -> do
      r1 <- gensym
      r2 <- gensym
      Emit [ QBE.BinaryOp (r1 QBE.:= QBE.Long) bop a b
           , QBE.BinaryOp (r2 QBE.:= QBE.Long) QBE.And
             (QBE.ValTemporary r1)
             (QBE.ValConst (QBE.CInt 0xffff_ffff_ffff_fffd))
           , QBE.BinaryOp (lowerName r QBE.:= QBE.Long) QBE.Or
             (QBE.ValTemporary r2)
             (QBE.ValConst (QBE.CInt 0b10))
           ]
        <$> lower' e k
    PrimCons x y -> lowerCons r x y e k
    PrimWrite x -> lowerWrite r x e k

lower'
  :: forall es. (GenSym :> 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 => QBE.Ident QBE.Label -> Exp -> Eff es QBE.Block
lower n e = buildBlock n <$> lower' e (pure . Exit . QBE.Ret . Just)

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 <- lower "start" e
      let f = wrapFunction @NonEmpty "main" [b]
      pure $ QBE.Program [] [] [f]
    _ -> do
      let low e = do
            bl <- gensym' "b"
            fl <- gensym' "f"
            b <- lower bl e
            pure $ wrapFunction @NonEmpty fl [b]
      fs <- traverse low anfs
      pure $ QBE.Program [] [] (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)