{-# 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
  (toANF, lower)
  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.Syntax qualified as Lam
import Gyehoek.Syntax (Name, Prim(..), Lit(..))
import Gyehoek.GenSym
import Control.Monad.Cont
import Data.Foldable
import Data.List.NonEmpty (NonEmpty((:|)), toList)
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)
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


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



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 Nothing name Nothing ps QBE.NoVariadic

prims :: QBE.Program
prims = QBE.Program mempty mempty primfns where
  primfns = [ 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 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

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

lowerVal (ValLit (LitInt n)) k =
  k . QBE.ValConst . QBE.CInt . fromIntegral $ n
lowerVal (ValLit _) k = error "todo"
lowerVal (ValVar x) k = k . QBE.ValTemporary . lowerName $ x

lowerBinaryOp :: QBE.Assignment -> Prim QBE.Val -> QBE.Inst
lowerBinaryOp 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)
      _ -> _

-- lowerPrim 
--   :: forall es. (GenSym :> es)
--   => Prim Val
--   -> (QBE.Val -> Eff es BlockBuilder)
--   -> Eff es BlockBuilder
-- lowerPrim p k = telescope (lowerVal <$> p) \p' -> do
--   Emit [ lowerBinaryOp (r QBE.:= QBE.Long) p' ]
--     <$> k (QBE.ValTemporary r)

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 =
  telescope (lowerVal <$> p) \p' -> do
    Emit [ lowerBinaryOp (lowerName r QBE.:= QBE.Long) p' ]
      <$> lower' e k

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

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)

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