module Quox.Parser.Parser

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

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


export
res : (str : String) -> (0 _ : IsReserved str) => Grammar True ()
res str = terminal "expecting \"\{str}\"" $ guard . (== Reserved str)

export
resC : (str : String) -> (0 _ : IsReserved str) => Grammar True ()
resC str = do res str; commit


parameters {default True commit : Bool}
  private
  maybeCommit : Grammar False ()
  maybeCommit = when commit Core.commit

  export
  betweenR : {c : Bool} -> (op, cl : String) ->
             (0 _ : IsReserved op) => (0 _ : IsReserved cl) =>
             Grammar c a -> Grammar True a
  betweenR o c p = res o *> maybeCommit *> p <* res c <* maybeCommit

  export
  parens, bracks, braces : {c : Bool} -> Grammar c a -> Grammar True a
  parens = betweenR "(" ")"
  bracks = betweenR "[" "]"
  braces = betweenR "{" "}"


export
commaSep : {c : Bool} -> Grammar c a -> Grammar False (List a)
commaSep = sepEndBy (res ",")
  -- don't commit because of the possible terminating ","

export
semiSep : {c : Bool} -> Grammar c a -> Grammar False (List a)
semiSep = sepEndBy (res ";")

export
commaSep1 : {c : Bool} -> Grammar c a -> Grammar c (List1 a)
commaSep1 = sepEndBy1 (res ",")

export
darr : Grammar True ()
darr = resC "⇒"


export
name : Grammar True PName
name = terminal "expecting name" $
  \case Name i => Just i; _ => Nothing

export
mods : Grammar True Mods
mods = name <&> \n => n.mods :< n.base

export
baseName : Grammar True String
baseName = terminal "expecting unqualified name" $
  \case Name i => guard (null i.mods) $> i.base
        _      => Nothing

export
nat : Grammar True Nat
nat = terminal "expecting natural number" $
  \case Nat n => pure n; _ => Nothing

export
string : Grammar True String
string = terminal "expecting string literal" $
  \case Str s => pure s; _ => Nothing

export
tag : Grammar True String
tag = terminal "expecting tag constructor" $
  \case Tag t => pure t; _ => Nothing

export
bareTag : Grammar True String
bareTag = string <|> [|toDotsP name|]

export
universe : Grammar True PUniverse
universe = terminal "expecting type universe" $
  \case TYPE u => Just u; _ => Nothing

export
bname : Grammar True BName
bname = Nothing <$ res "_"
    <|> [|Just baseName|]

export
dimConst : Grammar True DimConst
dimConst = terminal "expecting 0 or 1" $
  \case Nat 0 => Just Zero; Nat 1 => Just One; _ => Nothing

export
qty : Grammar True Qty
qty = terminal "expecting quantity" $
  \case Nat 0 => Just Zero; Nat 1 => Just One; Reserved "ω" => Just Any
        _ => Nothing

export
zero : Grammar True ()
zero = terminal "expecting 0" $ guard . (== Nat 0)


public export
AllReserved : List (a, String) -> Type
AllReserved = foldr (\x, ps => (IsReserved $ snd x, ps)) ()

export
symbols : (lst : List (a, String)) -> (0 _ : AllReserved lst) =>
          Grammar True a
symbols []                 = fail "no symbols found"
symbols [(x, str)]         = x <$ res str
symbols ((x, str) :: rest) = x <$ res str <|> symbols rest

export
symbolsC : (lst : List (a, String)) -> (0 _ : AllReserved lst) =>
           Grammar True a
symbolsC lst = symbols lst <* commit


private covering
qtys : Grammar False (List Qty)
qtys = option [] [|toList $ some qty <* res "."|]


export
dim : Grammar True PDim
dim = [|V baseName|] <|> [|K dimConst|]

private
0 PBinderHead : Nat -> Type
PBinderHead n = (Vect n Qty, BName, PTerm)

private
toVect : List a -> (n ** Vect n a)
toVect []        = (_ ** [])
toVect (x :: xs) = (_ ** x :: snd (toVect xs))


private
lamIntro : Grammar True (BName -> PTerm -> PTerm)
lamIntro = symbolsC [(Lam, "λ"), (DLam, "δ")]

private covering
caseIntro : Grammar True Qty
caseIntro = symbols [(Zero, "case0"), (One, "case1"), (Any, "caseω")]
        <|> resC "case" *>
              (qty <* resC "." <|>
               fatalError {c = True} "missing quantity on 'case'")

private
optNameBinder : Grammar False BName
optNameBinder = [|join $ optional $ bname <* darr|]

mutual
  export covering
  term : Grammar True PTerm
  term = lamTerm <|> caseTerm <|> coeTerm <|> compTerm <|> bindTerm <|> annTerm

  private covering
  lamTerm : Grammar True PTerm
  lamTerm = flip . foldr <$> lamIntro <*> some bname <* darr <*> term

  private covering
  caseTerm : Grammar True PTerm
  caseTerm =
    [|Case caseIntro term
           (resC "return" *> optBinderTerm)
           (resC "of"     *> caseBody)|]

  private covering
  coeTerm : Grammar True PTerm
  coeTerm = resC "coe" *> [|Coe typeLine dimArg dimArg aTerm|]

  private covering
  compTerm : Grammar True PTerm
  compTerm = resC "comp" *> do
    head <- [|Comp typeLine dimArg dimArg aTerm dimArg|]
    uncurry head <$> compBody
  where
    compBranch : Grammar True (DimConst, BName, PTerm)
    compBranch = [|(,,) dimConst bname (resC "⇒" *> term)|]

    zeroOne : (DimConst, a) -> (DimConst, a) -> Grammar False (a, a)
    zeroOne (Zero, x) (One,  y) = pure (x, y)
    zeroOne (One,  x) (Zero, y) = pure (y, x)
    zeroOne _ _ = fatalError "body of comp needs one 0 branch and one 1 branch"

    compBody : Grammar True ((BName, PTerm), (BName, PTerm))
    compBody = braces $ do
      et0 <- compBranch <* resC ";"
      et1 <- compBranch <* optional (resC ";")
      zeroOne et0 et1

  private covering
  caseBody : Grammar True PCaseBody
  caseBody = braces $
        [|CasePair (pairPat <* darr) (term <* optSemi)|]
    <|> [|CaseBox (bracks bname <* darr) (term <* optSemi)|]
    <|> CaseNat <$> zeroCase <* resC ";" <*> succCase <* optSemi
    <|> flip CaseNat <$> succCase <* resC ";" <*> zeroCase <* optSemi
    <|> [|CaseEnum $ semiSep [|MkPair tag (darr *> term)|]|]
  where
    optSemi = optional $ res ";"
    pairPat = parens [|MkPair bname (resC "," *> bname)|]

    zeroCase : Grammar True PTerm
    zeroCase = (resC "zero" <|> zero) *> darr *> term

    succCase : Grammar True (BName, Qty, BName, PTerm)
    succCase = do
      resC "succ"
      n  <- bname
      ih <- option (Zero, Nothing) $
        resC "," *> [|MkPair qty (resC "." *> bname)|]
      rhs <- darr *> term
      pure $ (n, fst ih, snd ih, rhs)

  private covering
  bindTerm : Grammar True PTerm
  bindTerm = pi <|> sigma
  where
    binderHead = parens {commit = False} [|MkPair bname (resC ":" *> term)|]

    pi, sigma : Grammar True PTerm
    pi = [|makePi (qty <* res ".") domain (resC "→" *> term)|]
    where
      makePi : Qty -> (BName, PTerm) -> PTerm -> PTerm
      makePi q (x, s) t = Pi q x s t
      domain = binderHead <|> [|(Nothing,) aTerm|]

    sigma = [|makeSigma binderHead (resC "×" *> annTerm)|]
    where
      makeSigma : (BName, PTerm) -> PTerm -> PTerm
      makeSigma (x, s) t = Sig x s t

  private covering
  annTerm : Grammar True PTerm
  annTerm = [|annotate infixEqTerm (optional $ resC "∷" *> term)|]
  where annotate s a = maybe s (s :#) a

  private covering
  infixEqTerm : Grammar True PTerm
  infixEqTerm = do
    l <- infixTimesTerm
    rty <- optional [|MkPair (resC "≡" *> term) (resC ":" *> infixTimesTerm)|]
    pure $ maybe l (\rty => Eq (Nothing, snd rty) l (fst rty)) rty

  private covering
  infixTimesTerm : Grammar True PTerm
  infixTimesTerm = foldr1 (Sig Nothing) <$> sepBy1 (resC "×") appTerm

  private covering
  appTerm : Grammar True PTerm
  appTerm = resC "★"    *> [|TYPE nat|]
        <|> resC "Eq"   *> [|Eq typeLine aTerm aTerm|]
        <|> resC "succ" *> [|Succ aTerm|]
        <|> [|apply aTerm (many appArg)|]
  where
    data PArg = TermArg PTerm | DimArg PDim

    appArg : Grammar True PArg
    appArg = [|DimArg dimArg|] <|> [|TermArg aTerm|]

    apply : PTerm -> List PArg -> PTerm
    apply = foldl apply1 where
      apply1 : PTerm -> PArg -> PTerm
      apply1 f (TermArg x) = f :@ x
      apply1 f (DimArg  p) = f :% p

  private covering
  typeLine : Grammar True (BName, PTerm)
  typeLine = bracks optBinderTerm

  private covering
  dimArg : Grammar True PDim
  dimArg = resC "@" *> dim

  private covering
  aTerm : Grammar True PTerm
  aTerm = [|Enum $ braces $ commaSep bareTag|]
      <|> foldr1 Pair <$> parens (commaSep1 term)
      <|> boxTerm
      <|> [|TYPE universe|]
      <|> Nat  <$ resC "ℕ"
      <|> Zero <$ resC "zero"
      <|> (nat <&> \n => fromNat n :# Nat)
      <|> [|V name|]
      <|> [|Tag tag|]

  private covering
  boxTerm : Grammar True PTerm
  boxTerm = bracks $
        [|BOX (qty <* resC ".") term|]
    <|> [|Box term|]

  private covering
  optBinderTerm : Grammar True (BName, PTerm)
  optBinderTerm = [|MkPair optNameBinder term|]


export covering
defIntro : Grammar True Qty
defIntro = symbols [(Zero, "def0"), (Any, "defω")]
       <|> resC "def" *> option Any (qty <* resC ".")

export covering
definition : Grammar True PDefinition
definition =
  [|MkPDef defIntro name (optional (resC ":" *> term)) (resC "=" *> term)|]

export
load : Grammar True String
load = resC "load" *> string

mutual
  export covering
  namespace_ : Grammar True PNamespace
  namespace_ =
    [|MkPNamespace (resC "namespace" *> mods) (braces $ many decl)|]

  export covering
  decl : Grammar True PDecl
  decl = [|PDef definition|]
     <|> [|PNs  namespace_|]

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

export covering
input : Grammar False (List PTopLevel)
input = skipSemis *> many (topLevel <* skipSemis)
  where skipSemis = ignore $ many $ resC ";"


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

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

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