{-# LANGUAGE RequiredTypeArguments #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE PartialTypeSignatures #-}
module Gyehoek.QBE.Parse where

import Language.QBE as QBE
import Effectful.State.Dynamic
import Effectful.Dispatch.Dynamic
import Effectful
import Numeric.Natural
import Data.String (IsString(fromString))
import Prettyprinter (Pretty(pretty), layoutPretty, defaultLayoutOptions)
import Data.Text (Text)
import Data.Data
import Prettyprinter.Render.Text (renderStrict)
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L
import Data.Void (Void)
import Data.Char (isAlpha, isAlphaNum)
import Control.Lens.Wrapped
import Data.Functor.Contravariant (Predicate(Predicate))
import qualified Data.Text as T
import Data.Functor
import Data.List (List)
import Data.Foldable (fold)
import Data.Maybe (isJust, fromMaybe)
import Control.Monad.Fix (MonadFix(mfix))
import Data.List.NonEmpty (fromList)
import Language.Haskell.TH qualified as TH
import Language.Haskell.TH.Quote
import Data.Proxy
import Data.Kind (Type)


type P = Parsec Void Text

sc :: P ()
sc = L.space hspace1 (L.skipLineComment "#") empty

lexeme :: P a -> P a
lexeme = L.lexeme sc

symbol :: Text -> P Text
symbol = L.symbol sc

infixr 8 .:
(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
(.:) f g x y = f (g x y)

rawIdent :: P QBE.RawIdent
rawIdent = (fromString .: (:) <$> lead <*> trail) <?> "ident"
  where
    lead = satisfy \x -> isAlpha x || (x=='.') || (x=='_')
    trail = fmap T.unpack . takeWhileP Nothing $ \x ->
      isAlphaNum x || (x=='.') || (x=='_')

class ParseIdent (s :: Sigil) where
  ident :: P (QBE.Ident s)

rawIdentWithSigil :: Char -> P (Ident t)
rawIdentWithSigil c = Ident <$> lexeme (char c *> rawIdent)

instance ParseIdent AggregateTy where ident = rawIdentWithSigil ':'
instance ParseIdent Global      where ident = rawIdentWithSigil '$'
instance ParseIdent Temporary   where ident = rawIdentWithSigil '%'
instance ParseIdent QBE.Label   where ident = rawIdentWithSigil '@'

const :: P QBE.Const
const = cint <|> csingle <|> cdouble <|> cglobal <?> "const"
  where
    cint = CInt <$> lexeme (L.signed empty L.decimal) <?> "integer"
    csingle = empty <?> "single-precision float"
    cdouble = empty <?> "double-precision float"
    cglobal = CGlobal <$> ident <?> "global symbol"

val :: P QBE.Val
val = vconst <|> vtemp <?> "val"
  where
    vconst = ValConst <$> Gyehoek.QBE.Parse.const
    vtemp = ValTemporary <$> ident <?> "temporary symbol"

assignment :: P QBE.Assignment
assignment =
  Assignment <$> ident <*> (char '=' *> basety)

basety :: P QBE.BaseTy
basety = lexeme $ char 'w' $> Word
              <|> char 'l' $> Long
              <|> char 's' $> Single
              <|> char 'd' $> Double

abity :: P AbiTy
abity = AbiBaseTy <$> basety
    <|> AbiAggregateTy <$> ident

binaryOp :: P QBE.BinaryOp
binaryOp = lexeme $ "add" $> Add
                <|> "sub" $> Sub
                <|> "mul" $> Mul
                <|> "div" $> Div Signed

comma :: P a -> P a
comma p = symbol "," *> p

inst :: P QBE.Inst
inst = try binaryOpInst <|> negInst <?> "inst"
  where
    binaryOpInst =
      BinaryOp
        <$> assignment
        <*> binaryOp
        <*> val <*> comma val
    negInst = Neg <$> assignment <*> (symbol "neg" *> val)

jump :: P QBE.Jump
jump = jmp <|> jnz <|> ret <|> hlt <?> "jump"
  where
    jmp = symbol "jmp" *> (Jmp <$> ident)
    jnz = symbol "jnz" *> (Jnz <$> val <*> ident <*> comma ident)
    ret = symbol "ret" *> (Ret <$> optional val)
    hlt = empty

nl :: P ()
nl = void (some (newline *> sc)) <?> "newline"

phi :: P QBE.Phi
phi = empty

block :: P QBE.Block
block = Block
  <$> (ident <* nl)
  <*> sepBy phi nl
  <*> sepBy inst nl
  <*> jump

sepByTry :: MonadParsec e s m => m a -> m sep -> m (List a)
sepByTry p sep = do
  x <- p
  xs <- many (try $ sep *> p)
  pure (x:xs)

paramList :: P (Maybe (Ident Temporary), List Param, Variadic)
paramList = label "parameter list" $ between (symbol "(") (symbol ")") do
  e <- optional env
  ps <- optional . try $ do
    commaIf (isJust e)
    sepByTry reg (symbol ",")
  v <- optional do
    commaIf (isJust e || isJust ps)
    variadic
  pure (e, fromMaybe [] ps, fromMaybe NoVariadic v)
  where
    commaIf True = void $ symbol ","
    commaIf False = pure ()
    env = symbol "env" *> ident @Temporary <?> "environment parameter"
    reg = Param <$> abity <*> ident <?> "regular parameter"
    variadic = symbol "..." $> Variadic <?> "variadic parameter"

funcdef :: P QBE.FuncDef
funcdef = do
  linkages <- many linkage
  symbol "function"
  returnTy <- optional abity
  name <- ident @Global
  (env,params,variadic) <- paramList
  code <- fmap fromList . between (symbol "{" *> nl) (symbol "}") $
    sepEndBy1 block nl
  pure $ FuncDef linkages returnTy name env params variadic code

linkage :: P Linkage
linkage = symbol "export" $> Export

-- stripped :: P a -> P a
-- stripped p = optional nl *>



class Data a => Parser a where
  parser :: P a

instance Parser FuncDef where parser = funcdef

class ParseSeparator a where
  parseSeparator :: Proxy a -> P ()

instance (Parser a, ParseSeparator a) => Parser (List a) where
  parser = sepBy parser (parseSeparator @a Proxy)

instance ParseSeparator FuncDef where parseSeparator _ = nl
instance ParseSeparator Block where parseSeparator _ = nl