module Quox.Parser.Parser

import public Quox.Parser.Lexer
import public Quox.Parser.Syntax

import Data.Bool
import Data.Fin
import Data.Vect
import public Text.Parser
import Derive.Prelude

%language ElabReflection
%default total


public export
0 Grammar : Bool -> Type -> Type
Grammar = Core.Grammar () Token
%hide Core.Grammar


public export
data Error =
    LexError   Lexer.Error
  | ParseError (List1 (ParsingError Token))
%hide Lexer.Error
%runElab derive "Error" [Show]

export
lexParseWith : {c : Bool} -> Grammar c a -> String -> Either Error a
lexParseWith grm input = do
  toks <- mapFst LexError $ lex input
  bimap ParseError fst $ parse (grm <* eof) toks


||| reserved token, like punctuation or keywords etc
export
res : (str : String) -> (0 _ : IsReserved str) => Grammar True ()
res str = terminal "expected \"\{str}\"" $ guard . (== Reserved str)

||| optional reserved token, e.g. trailing comma
export
optRes : (str : String) -> (0 _ : IsReserved str) => Grammar False ()
optRes str = ignore $ optional $ res str

||| reserved token, then commit
export
resC : (str : String) -> (0 _ : IsReserved str) => Grammar True ()
resC str = do res str; commit

||| reserved token or fatal error
export
needRes : (str : String) -> (0 _ : IsReserved str) => Grammar True ()
needRes str = resC str <|> fatalError "expected \"\{str}\""


private
terminalMatchN_ : String -> List (TTImp, TTImp) -> Elab (Grammar True a)
terminalMatchN_ what matches = do
  func <- check $ lam (lambdaArg `{x}) $
    iCase `(x) implicitFalse $
    map (\(l, r) => patClause l `(Just ~(r))) matches ++
      [patClause `(_) `(Nothing)]
  pure $ terminal "expected \{what}" func

private %macro
terminalMatchN : String -> List (TTImp, TTImp) -> Elab (Grammar True a)
terminalMatchN = terminalMatchN_

private %macro
terminalMatch : String -> TTImp -> TTImp -> Elab (Grammar True a)
terminalMatch what l r = terminalMatchN_ what [(l, r)]

||| tag without leading `'`
export
bareTag : Grammar True TagVal
bareTag = terminalMatchN "bare tag" [(`(Name t), `(show t)), (`(Str s), `(s))]

||| tag with leading quote
export
tag : Grammar True TagVal
tag = terminalMatch "tag" `(Tag t) `(t)

||| natural number
export
nat : Grammar True Nat
nat = terminalMatch "natural number" `(Nat n) `(n)

||| string literal
export
strLit : Grammar True String
strLit = terminalMatch "string literal" `(Str s) `(s)

||| single-token universe, like ★₀ or Type1
export
universe1 : Grammar True Universe
universe1 = terminalMatch "universe" `(TYPE u) `(u)

||| possibly-qualified name
export
qname : Grammar True PName
qname = terminalMatch "name" `(Name n) `(n)

||| unqualified name
export
baseName : Grammar True PBaseName
baseName = terminalMatch "unqualified name" `(Name (MakePName [<] b)) `(b)

||| dimension constant (0 or 1)
export
dimConst : Grammar True DimConst
dimConst = terminalMatchN "dimension constant"
  [(`(Nat 0), `(Zero)), (`(Nat 1), `(One))]

||| quantity (0, 1, or ω)
export
qty : Grammar True Qty
qty = terminalMatchN "quantity"
  [(`(Nat 0), `(Zero)), (`(Nat 1), `(One)), (`(Reserved "ω"), `(Any))]


||| pattern var (unqualified name or _)
export
patVar : Grammar True BName
patVar = [|Just baseName|] <|> Nothing <$ res "_"


||| dimension (without `@` prefix)
export
dim : Grammar True PDim
dim = [|K dimConst|] <|> [|V baseName|]

||| dimension argument (with @)
export
dimArg : Grammar True PDim
dimArg = resC "@" *> mustWork dim


delim : (o, c : String) -> (0 _ : IsReserved o) => (0 _ : IsReserved c) =>
        {k : Bool} -> Grammar k a -> Grammar True a
delim o c p = resC o *> p <* needRes c

-- this stuff is Like This (rather than just being delim + sepEndBy{1})
-- so that it checks for the close bracket before trying another list element,
-- giving (imo) a better error
parameters (o, c, s : String)
           {auto 0 _ : IsReserved o} {auto 0 _ : IsReserved c}
           {auto 0 _ : IsReserved s}
           (p : Grammar True a)
  private
  dsBeforeDelim, dsAfterDelim : Grammar True (List a)
  dsBeforeDelim = [] <$ resC c <|> resC s *> assert_total dsAfterDelim
  dsAfterDelim  = [] <$ resC c <|> [|p :: assert_total dsBeforeDelim|]

  export
  delimSep1 : Grammar True (List1 a)
  delimSep1 = resC o *> [|p ::: dsBeforeDelim|]

  export
  delimSep : Grammar True (List a)
  delimSep = resC o *> dsAfterDelim


||| enum type, e.g. `{a, b, c.d, "e f g"}`
export
enumType : Grammar True (List TagVal)
enumType = delimSep "{" "}" "," bareTag

||| e.g. `case` or `case 1.`
export
caseIntro : Grammar True Qty
caseIntro =
      Zero <$ res "case0"
  <|> One  <$ res "case1"
  <|> Any  <$ res "caseω"
  <|> delim "case" "." qty

export
qtyPatVar : Grammar True (Qty, BName)
qtyPatVar = [|(,) qty (needRes "." *> patVar)|]


public export
data PCasePat =
    PPair BName BName
  | PTag  TagVal
  | PZero
  | PSucc BName Qty BName
  | PBox  BName
%runElab derive "PCasePat" [Eq, Ord, Show]

||| either `zero` or `0`
export
zeroPat : Grammar True ()
zeroPat = resC "zero" <|> terminal "expected '0'" (guard . (== Nat 0))

export
casePat : Grammar True PCasePat
casePat =
      delim "(" ")" [|PPair patVar (needRes "," *> patVar)|]
  <|> [|PTag tag|]
  <|> PZero <$ zeroPat
  <|> do p  <- resC "succ" *> patVar
         ih <- option (Zero, Nothing) $ resC "," *> qtyPatVar
         pure $ PSucc p (fst ih) (snd ih)
  <|> delim "[" "]" [|PBox patVar|]
  <|> fatalError "invalid pattern"


export covering
term : Grammar True PTerm

export covering
typeLine : Grammar True (BName, PTerm)
typeLine = do
  resC "["
  mustWork $ do
    i <- option Nothing $ patVar <* resC "⇒"
    t <- term <* needRes "]"
    pure (i, t)

||| box term `[t]` or type `[π.A]`
export covering
boxTerm : Grammar True PTerm
boxTerm = do
  res "["; commit
  q <- optional $ qty  <* res "."
  t <- mustWork $ term <* needRes "]"
  pure $ maybe (Box t) (\q => BOX q t) q

||| tuple term like `(a, b)`, or parenthesised single term.
||| allows terminating comma. more than two elements are nested on the right:
||| `(a, b, c, d) = (a, (b, (c, d)))`.
export covering
tupleTerm : Grammar True PTerm
tupleTerm = foldr1 Pair <$> delimSep1 "(" ")" "," term

||| argument/atomic term: single-token terms, or those with delimiters e.g.
||| `[t]`
export covering
termArg : Grammar True PTerm
termArg = [|TYPE universe1|]
      <|> [|Enum enumType|]
      <|> [|Tag tag|]
      <|> boxTerm
      <|> Nat  <$ res "ℕ"
      <|> Zero <$ res "zero"
      <|> [|fromNat nat|]
      <|> [|V qname|]
      <|> tupleTerm

export covering
coeTerm : Grammar True PTerm
coeTerm = resC "coe" *> mustWork [|Coe typeLine dimArg dimArg termArg|]

export covering
compBranch : Grammar True (DimConst, BName, PTerm)
compBranch = [|(,,) dimConst patVar (needRes "⇒" *> term)|]

export covering
compTerm : Grammar True PTerm
compTerm = do
  resC "comp"
  mustWork $ do
    a <- typeLine
    p <- dimArg;  q <- dimArg
    s <- termArg; r <- dimArg
    needRes "{"
    s0 <- compBranch; needRes ";"
    s1 <- compBranch; optRes  ";"
    needRes "}"
    case (fst s0, fst s1) of
      (Zero, One)  => pure $ Comp a p q s r (snd s0) (snd s1)
      (One,  Zero) => pure $ Comp a p q s r (snd s1) (snd s0)
      (_,    _)    =>
        fatalError "body of 'comp' needs one 0 case and one 1 case"

export covering
splitUniverseTerm : Grammar True PTerm
splitUniverseTerm = resC "★" *> mustWork [|TYPE nat|]

export covering
eqTerm : Grammar True PTerm
eqTerm = resC "Eq" *> mustWork [|Eq typeLine termArg termArg|]

export covering
succTerm : Grammar True PTerm
succTerm = resC "succ" *> mustWork [|Succ termArg|]

||| a dimension argument with an `@` prefix, or
||| a term argument with no prefix
export covering
anyArg : Grammar True (Either PDim PTerm)
anyArg = dimArg <||> termArg

export covering
normalAppTerm : Grammar True PTerm
normalAppTerm = foldl ap <$> termArg <*> many anyArg
where ap : PTerm -> Either PDim PTerm -> PTerm
      ap f (Left  p) = f :% p
      ap f (Right s) = f :@ s

||| application term `f x @y z`, or other terms that look like application
||| like `succ` or `coe`.
export covering
appTerm : Grammar True PTerm
appTerm = coeTerm
      <|> compTerm
      <|> splitUniverseTerm
      <|> eqTerm
      <|> succTerm
      <|> normalAppTerm

export covering
infixEqTerm : Grammar True PTerm
infixEqTerm = do
  l <- appTerm; commit
  rest <- optional $ res "≡" *> [|(,) term (needRes ":" *> appTerm)|]
  pure $ maybe l (\rest => Eq (Nothing, snd rest) l (fst rest)) rest

export covering
annTerm : Grammar True PTerm
annTerm = do
  tm <- infixEqTerm; commit
  ty <- optional $ res "∷" *> term
  pure $ maybe tm (tm :#) ty

export covering
lamTerm : Grammar True PTerm
lamTerm = do
  k <- DLam <$ res "δ" <|> Lam <$ res "λ"
  mustWork $ do
    xs   <- some patVar; needRes "⇒"
    body <- term;        commit
    pure $ foldr k body xs

-- [todo] fix the backtracking in e.g. (F x y z × B)
export covering
properBinders : Grammar True (List1 BName, PTerm)
properBinders = do
  res "("
  xs <- some patVar; resC ":"
  t  <- term;        needRes ")"
  pure (xs, t)

export covering
piTerm : Grammar True PTerm
piTerm = do
  q   <- qty;      resC "."
  dom <- piBinder; needRes "→"
  cod <- term;     commit
  pure $ foldr (\x => Pi q x (snd dom)) cod (fst dom)
where
  piBinder : Grammar True (List1 BName, PTerm)
  piBinder = properBinders <|> [|(singleton Nothing,) termArg|]

export covering
sigmaTerm : Grammar True PTerm
sigmaTerm =
      (properBinders >>= continueDep)
  <|> (annTerm       >>= continueNondep)
where
  continueDep : (List1 BName, PTerm) -> Grammar True PTerm
  continueDep (names, dom) = do
    cod <- needRes "×" *> sigmaTerm
    pure $ foldr (\x => Sig x dom) cod names
  continueNondep : PTerm -> Grammar False PTerm
  continueNondep dom = do
    rest <- optional $ resC "×" *> sepBy1 (res "×") annTerm
    Core.pure $ foldr1 (Sig Nothing) $ dom ::: maybe [] forget rest


public export
PCaseArm : Type
PCaseArm = (PCasePat, PTerm)

export
checkCaseArms : List PCaseArm -> Grammar False PCaseBody
checkCaseArms [] =
  pure $ CaseEnum []
checkCaseArms ((PPair x y, rhs) :: rest) = do
  let [] = rest | _ => fatalError "unexpected pattern after pair"
  pure $ CasePair (x, y) rhs
checkCaseArms ((PTag tag, rhs1) :: rest) = do
  let rest = for rest $ \case
        (PTag tag, rhs) => Just (tag, rhs)
        _               => Nothing
  maybe (fatalError "expected all patterns to be tags")
        (pure . CaseEnum . ((tag, rhs1) ::)) rest
checkCaseArms ((PZero, rhs1) :: rest) = do
  let [(PSucc p q ih, rhs2)] = rest
    | _ => fatalError "expected succ pattern after zero"
  pure $ CaseNat rhs1 (p, q, ih, rhs2)
checkCaseArms ((PSucc p q ih, rhs1) :: rest) = do
  let [(PZero, rhs2)] = rest
    | _ => fatalError "expected zero pattern after succ"
  pure $ CaseNat rhs2 (p, q, ih, rhs1)
checkCaseArms ((PBox x, rhs) :: rest) = do
  let [] = rest | _ => fatalError "unexpected pattern after box"
  pure $ CaseBox x rhs

export covering
caseArm : Grammar True PCaseArm
caseArm = [|(,) casePat (needRes "⇒" *> term)|]

export covering
caseBody : Grammar True PCaseBody
caseBody = delimSep "{" "}" ";" caseArm >>= checkCaseArms

export covering
caseTerm : Grammar True PTerm
caseTerm = do
  qty <- caseIntro; commit
  head <- mustWork term; needRes "return"
  ret  <- mustWork [|(,) (option Nothing $ patVar <* resC "⇒") term|]
  needRes "of"
  body <- mustWork caseBody
  pure $ Case qty head ret body


term = lamTerm
   <|> caseTerm
   <|> piTerm
   <|> sigmaTerm


||| `def` alone means `defω`
export
defIntro : Grammar True Qty
defIntro = Zero <$ resC "def0"
       <|> Any  <$ resC "defω"
       <|> resC "def" *> option Any (qty <* needRes ".")


export covering
decl : Grammar True PDecl

export covering
definition : Grammar True PDefinition
definition = do
  qty  <- defIntro
  name <- baseName
  type <- optional $ resC ":" *> mustWork term
  term <- needRes "=" *> mustWork term
  optRes ";"
  pure $ MkPDef {qty, name, type, term}

export covering
namespace_ : Grammar True PNamespace
namespace_ = do
  ns    <- resC "namespace" *> qname; needRes "{"
  decls <- nsInner; optRes ";"
  pure $ MkPNamespace (ns.mods :< ns.base) decls
where
  nsInner : Grammar True (List PDecl)
  nsInner = [] <$ resC "}"
        <|> [|(decl <* commit) :: nsInner|]

decl = [|PDef definition|] <|> [|PNs namespace_|]

export covering
load : Grammar True String
load = resC "load" *> mustWork strLit <* optRes ";"

export covering
topLevel : Grammar True PTopLevel
topLevel = [|PLoad load|] <|> [|PD decl|]

export covering
input : Grammar False (List PTopLevel)
input = [] <$ eof
    <|> [|(topLevel <* commit) :: input|]

export covering
lexParseTerm : String -> Either Error PTerm
lexParseTerm = lexParseWith term

export covering
lexParseInput : String -> Either Error (List PTopLevel)
lexParseInput = lexParseWith input