module Parser.Parse

import Text.Lexer
import Text.Parser

import Data.Vect
import Data.List
import Data.List1
import Data.Maybe

-- eventually start using elaborator representation here, or just a parser
-- exclusive one
import Core.Term
import Core.Misc

%default total

data PiTokenKind
  = PTLambda
  | PTIdentifier
  | PTDot
  | PTArrow
  | PTPi
  | PTSigma
  | PTComma
  | PTLParen
  | PTRParen
  | PTColon
  | PTLet
  | PTDefEq
  | PTIn
  | PTIgnore
  | PTType
  | PTTop
  | PTStar
  | PTTopInd
  | PTBot
  | PTBotInd
  | PTNat
  | PTNum
  | PTSuc
  | PTNatInd
  | PTSigInd
  | PTId
  | PTRefl
  | PTJ

Eq PiTokenKind where
  (==) PTLambda     PTLambda     = True
  (==) PTIdentifier PTIdentifier = True
  (==) PTDot        PTDot        = True
  (==) PTArrow      PTArrow      = True
  (==) PTPi         PTPi         = True
  (==) PTSigma      PTSigma      = True
  (==) PTComma      PTComma      = True
  (==) PTLParen     PTLParen     = True
  (==) PTRParen     PTRParen     = True
  (==) PTColon      PTColon      = True
  (==) PTLet        PTLet        = True
  (==) PTDefEq      PTDefEq      = True
  (==) PTIn         PTIn         = True
  (==) PTIgnore     PTIgnore     = True
  (==) PTType       PTType       = True
  (==) PTTop        PTTop        = True
  (==) PTStar       PTStar       = True
  (==) PTTopInd     PTTopInd     = True
  (==) PTBot        PTBot        = True
  (==) PTBotInd     PTBotInd     = True
  (==) PTNat        PTNat        = True
  (==) PTNum        PTNum        = True
  (==) PTSuc        PTSuc        = True
  (==) PTNatInd     PTNatInd     = True
  (==) PTSigInd     PTSigInd     = True
  (==) PTId         PTId         = True
  (==) PTRefl       PTRefl       = True
  (==) PTJ          PTJ          = True
  (==) _            _            = False

Show PiTokenKind where
  show PTLambda     = "PTLambda"
  show PTIdentifier = "PTIdentifier"
  show PTDot        = "PTDot"
  show PTArrow      = "PTArrow"
  show PTPi         = "PTPi"
  show PTSigma      = "PTSigma"
  show PTComma      = "PTComma"
  show PTLParen     = "PTLParen"
  show PTRParen     = "PTRParen"
  show PTColon      = "PTColon"
  show PTLet        = "PTLet"
  show PTDefEq      = "PTDefEq"
  show PTIn         = "PTIn"
  show PTIgnore     = "PTIgnore"
  show PTType       = "PTType"
  show PTTop        = "PTTop"
  show PTStar       = "PTSTar"
  show PTTopInd     = "PTTopInd"
  show PTBot        = "PTBot"
  show PTBotInd     = "PTBotInd"
  show PTNat        = "PTNat"
  show PTNum        = "PTNum "
  show PTSuc        = "PTSuc"
  show PTNatInd     = "PTNatInd"
  show PTSigInd     = "PTSigInd"
  show PTId         = "PTId"
  show PTRefl       = "PTRefl"
  show PTJ          = "PTJ"

PiToken : Type
PiToken = Token PiTokenKind

Show PiToken where
  show (Tok kind text) = "Tok kind: " ++ show kind ++ " text: " ++ text

TokenKind PiTokenKind where
  TokType PTIdentifier = String
  TokType PTNum        = Nat
  TokType _ = ()

  tokValue PTLambda     _ = ()
  tokValue PTIdentifier s = s
  tokValue PTDot        _ = ()
  tokValue PTArrow      _ = ()
  tokValue PTPi         _ = ()
  tokValue PTSigma      _ = ()
  tokValue PTComma      _ = ()
  tokValue PTLParen     _ = ()
  tokValue PTRParen     _ = ()
  tokValue PTColon      _ = ()
  tokValue PTLet        _ = ()
  tokValue PTDefEq      _ = ()
  tokValue PTIn         _ = ()
  tokValue PTIgnore     _ = ()
  tokValue PTType       _ = ()
  tokValue PTTop        _ = ()
  tokValue PTStar       _ = ()
  tokValue PTTopInd     _ = ()
  tokValue PTBot        _ = ()
  tokValue PTBotInd     _ = ()
  tokValue PTNat        _ = ()
  tokValue PTNum        s = cast s
  tokValue PTSuc        _ = ()
  tokValue PTNatInd     _ = ()
  tokValue PTSigInd     _ = ()
  tokValue PTId         _ = ()
  tokValue PTRefl       _ = ()
  tokValue PTJ          _ = ()

ignored : WithBounds PiToken -> Bool
ignored (MkBounded (Tok PTIgnore _) _ _) = True
ignored _ = False

identifier : Lexer
identifier = some (alphaNum <|> oneOf "-_")

keywords : List (String, PiTokenKind)
keywords = [
  ("in",    PTIn),
  ("let",   PTLet),
  ("suc",   PTSuc),
  ("Type",  PTType),
  ("J",     PTJ),
  ("Id",    PTId),
  ("refl",  PTRefl)
]

tokenmap : List (Lexer, PiTokenKind)
tokenmap = [
  (exact "⊤-ind", PTTopInd),
  (exact "⊥-ind", PTBotInd),
  (exact "ℕ-ind", PTNatInd),
  (exact "Σ-ind", PTSigInd),
  (is 'λ',        PTLambda),
  (is '.',        PTDot),
  (is '→',        PTArrow),
  (is 'Π',        PTPi),
  (is '∏',        PTPi),
  (is 'Σ',        PTSigma),
  (is '∑',        PTSigma),
  (is ',',        PTComma),
  (is '(',        PTLParen),
  (is ')',        PTRParen),
  (is ':',        PTColon),
  (is '≔',        PTDefEq),
  (is '⊤',        PTTop),
  (is '⊥',        PTBot),
  (is 'ℕ',        PTNat),
  (is '★',        PTStar),
  (digits,        PTNum)
]

piTokenMap : TokenMap PiToken
piTokenMap =
  (spaces, Tok PTIgnore) ::
  (lineComment  (exact "--"), Tok PTIgnore) ::
  (blockComment (exact "{-") (exact "-}"), Tok PTIgnore) ::
  toTokenMap tokenmap ++
  [(identifier, \s =>
      case lookup s keywords of
        (Just kind) => Tok kind s
        Nothing => Tok PTIdentifier s
    )
  ]

lexPi : String -> Maybe (List (WithBounds PiToken))
lexPi str =
  case lex piTokenMap str of
    (tokens, _, _, "") => Just tokens
    _ => Nothing

{- Basic idea, parsing has a list of the identifiers to convert to -}
{- de bruijn indecies, and a Nat to keep track of context size -}
mutual
  expr : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  expr defs n env = tpi defs n env
                <|> tsigma defs n env
                <|> tarr defs n env
                <|> (do 
                  e <- expr1 defs n env
                  tapp defs n env e <|> pure e)


  expr1 : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  expr1 defs n env = ttopind defs n env
                 <|> tbotind defs n env
                 <|> tsuc defs n env
                 <|> tnatind defs n env
                 <|> tsigind defs n env
                 <|> tid defs n env
                 <|> trefl defs n env
                 <|> tj defs n env
                 <|> (do
                   t <- term defs n env
                   tapp defs n env t <|> pure t)

  term : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  term defs n env = ttype
                <|> ttop
                <|> tstar
                <|> tbot
                <|> tnat
                <|> tnum
                <|> tpair defs n env
                <|> tlet defs n env
                <|> tlam defs n env
                <|> tvar defs n env
                <|> paren defs n env

  ttype : Grammar () PiToken True (Term n)
  ttype = match PTType >> pure TType

  ttop : Grammar () PiToken True (Term n)
  ttop = match PTTop >> pure TTop

  tstar : Grammar () PiToken True (Term n)
  tstar = match PTStar >> pure TStar

  ttopind : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  ttopind defs n env = do
    match PTTopInd
    commit
    c  <- term defs n env
    st <- term defs n env
    pure (TTopInd c st)

  tbot : Grammar () PiToken True (Term n)
  tbot = match PTBot >> pure TBot

  tbotind : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tbotind defs n env = do
    match PTBotInd
    commit
    c <- term defs n env
    pure (TBotInd c)

  tnat : Grammar () PiToken True (Term n)
  tnat = match PTNat >> pure TNat

  tnum : Grammar () PiToken True (Term n)
  tnum = do
    n <- match PTNum
    pure (conv n)
    where
      conv : Nat -> Term n
      conv 0     = TZero
      conv (S n) = TSuc (conv n)

  tsuc : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tsuc defs n env = do
    match PTSuc
    commit
    TSuc <$> term defs n env

  tnatind : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tnatind defs n env = do
    match PTNatInd
    commit
    c <- term defs n env
    z <- term defs n env
    s <- term defs n env
    pure (TNatInd c z s)

  tsigma : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tsigma defs n env = do
    match PTSigma
    commit
    match PTLParen
    arg <- match PTIdentifier
    match PTColon
    a <- expr defs n env
    match PTRParen
    b <- expr defs (S n) (arg :: env)
    pure (TSigma a (TLam b))

  tpair : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tpair defs n env = do
    match PTLParen
    x <- expr defs n env
    match PTComma
    commit
    y <- expr defs n env
    match PTRParen
    pure (TPair x y)

  tsigind : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tsigind defs n env = do
    match PTSigInd
    commit
    a <- term defs n env
    b <- term defs n env
    c <- term defs n env
    f <- term defs n env
    pure (TSigInd a b c f)

  tid : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tid defs n env = do
    match PTId
    commit
    ty <- term defs n env
    a  <- term defs n env
    b  <- term defs n env
    pure (TId ty a b)

  trefl : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  trefl defs n env = do
    match PTRefl
    commit
    ty <- term defs n env
    tr <- term defs n env
    pure (TRefl ty tr)

  tj : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tj defs n env = do
    match PTJ
    commit
    ty <- term defs n env
    a  <- term defs n env
    b  <- term defs n env
    c  <- term defs n env
    d  <- term defs n env
    pure (TJ ty a b c d)

  tlet : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tlet defs n env = do
    match PTLet
    commit
    arg <- match PTIdentifier
    match PTColon
    ty <- expr defs n env
    match PTDefEq
    tr <- expr defs n env
    match PTIn
    tri <- expr defs (S n) (arg :: env)
    pure (TLet ty tr tri)

  tlam : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tlam defs n env = do
    match PTLambda
    commit
    arg <- match PTIdentifier
    match PTDot
    e <- expr defs (S n) (arg :: env)
    pure (TLam e)

  tpi : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tpi defs n env = do
    match PTPi
    commit
    match PTLParen
    arg <- match PTIdentifier
    match PTColon
    a <- expr defs n env
    match PTRParen
    b <- expr defs (S n) (arg :: env)
    pure (TPi a b)

  tarr : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tarr defs n env = do
    l <- expr1 defs n env
    match PTArrow
    commit
    r <- expr defs (S n) ("" :: env)
    pure (TPi l r)

  tapp : List String -> (n : Nat) -> Vect n String -> Term n -> Grammar () PiToken True (Term n)
  tapp defs n env e1 = do
    e2 <- term defs n env
    tapp1 defs n env (TApp e1 e2)

  tapp1 : List String -> (n : Nat) -> Vect n String -> Term n -> Grammar () PiToken False (Term n)
  tapp1 defs n env e = tapp defs n env e <|> pure e

  tvar : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  tvar defs n env = do
    str <- match PTIdentifier
    commit
    fromMaybe (fromMaybe (fail ("'" ++ str ++ "' not idefs n env"))
                         (pure . TDef . cast . finToInteger <$> findIndex (== str) defs))
              (pure . TVar <$> findIndex (== str) env)

  paren : List String -> (n : Nat) -> Vect n String -> Grammar () PiToken True (Term n)
  paren defs n env = do
    match PTLParen
    commit
    e <- expr defs n env
    match PTRParen
    pure e

definitions : List String -> Grammar () PiToken True (List String, List (Term 0, Term 0))
definitions defs = do
    match PTLet
    commit
    arg <- match PTIdentifier
    match PTColon
    ty <- expr defs 0 []
    match PTDefEq
    tr <- expr defs 0 []
    next <- definitions (arg :: defs) <|> pure ([], [])
    pure (arg :: fst next, (ty, tr) :: snd next)

parsePi : List String -> Grammar () PiToken True a -> List (WithBounds PiToken) -> Either String a
parsePi defs parseEntry toks =
  case parse parseEntry $ filter (not . ignored) toks of
    Right (l, []) => Right l
    Right (_, l) => Left ("contains tokens that were not consumed: " ++ show l)
    Left e  => (Left . show . map getErr) e
  where
    getErr : ParsingError tok -> String
    getErr (Error "Unrecognised input" _) = ""
    getErr (Error s _) = s


public export
parse : List String -> Grammar () PiToken True a -> String -> Either String a
parse defs parseEntry x =
  case lexPi x of
    Just toks => parsePi defs parseEntry toks
    Nothing => Left "Failed to lex."

public export
parse0 : String -> Either String (Term 0)
parse0 = parse [] (expr [] 0 [])

public export
parseEnv : List String -> String -> Either String (Term 0)
parseEnv defs = parse defs (expr defs 0 [])

public export
toplevel : String -> Either String (List String, List (Term 0, Term 0))
toplevel = parse [] (definitions [])