module Obs.Parser

import Data.String

import Obs.Logging
import Obs.Sort
import Obs.Syntax

import System
import System.File

import Text.Lexer
import Text.Parser

%default total

data ObsTokenKind
  = OTNewline
  | OTIgnore
  | OTIdent
  -- Keywords
  | OTProp
  | OTSet
  | OTNat
  | OTFst
  | OTSnd
  | OTPoint
  | OTVoid
  | OTAbsurd
  | OTRefl
  | OTTransp
  | OTCast
  | OTCastRefl
  -- Special symbols
  | OTUnit
  -- Grouping symbols
  | OTPOpen
  | OTPClose
  | OTAOpen
  | OTAClose
  -- Definition characters
  | OTBackslash
  | OTThinArrow
  | OTFatArrow
  | OTProduct
  | OTComma
  | OTTilde
  | OTEqual
  | OTColon

Eq ObsTokenKind where
  OTNewline == OTNewline = True
  OTIgnore == OTIgnore = True
  OTIdent == OTIdent = True
  OTProp == OTProp = True
  OTSet == OTSet = True
  OTNat == OTNat = True
  OTFst == OTFst = True
  OTSnd == OTSnd = True
  OTPoint == OTPoint = True
  OTVoid == OTVoid = True
  OTAbsurd == OTAbsurd = True
  OTRefl == OTRefl = True
  OTTransp == OTTransp = True
  OTCast == OTCast = True
  OTCastRefl == OTCastRefl = True
  OTUnit == OTUnit = True
  OTPOpen == OTPOpen = True
  OTPClose == OTPClose = True
  OTAOpen == OTAOpen = True
  OTAClose == OTAClose = True
  OTBackslash == OTBackslash = True
  OTThinArrow == OTThinArrow = True
  OTFatArrow == OTFatArrow = True
  OTProduct == OTProduct = True
  OTComma == OTComma = True
  OTTilde == OTTilde = True
  OTEqual == OTEqual = True
  OTColon == OTColon = True
  _ == _ = False

TokenKind ObsTokenKind where
  TokType OTIdent = String
  TokType OTNat   = Nat
  TokType _     = ()

  tokValue OTNewline s = ()
  tokValue OTIgnore s = ()
  tokValue OTIdent s = s
  tokValue OTProp s = ()
  tokValue OTSet s = ()
  tokValue OTNat s = stringToNatOrZ s
  tokValue OTFst s = ()
  tokValue OTSnd s = ()
  tokValue OTPoint s = ()
  tokValue OTVoid s = ()
  tokValue OTAbsurd s = ()
  tokValue OTRefl s = ()
  tokValue OTTransp s = ()
  tokValue OTCast s = ()
  tokValue OTCastRefl s = ()
  tokValue OTUnit s = ()
  tokValue OTPOpen s = ()
  tokValue OTPClose s = ()
  tokValue OTAOpen s = ()
  tokValue OTAClose s = ()
  tokValue OTBackslash s = ()
  tokValue OTThinArrow s = ()
  tokValue OTFatArrow s = ()
  tokValue OTProduct s = ()
  tokValue OTComma s = ()
  tokValue OTTilde s = ()
  tokValue OTEqual s = ()
  tokValue OTColon s = ()

ObsToken : Type
ObsToken = Token ObsTokenKind

ignored : WithBounds ObsToken -> Bool
ignored (MkBounded (Tok OTIgnore _) _ _) = True
ignored _ = False

identifier : Lexer
identifier = alpha <+> many alphaNum

keywords : List (String, ObsTokenKind)
keywords =
  [ ("Set", OTSet)
  , ("Prop", OTProp)
  , ("fst", OTFst)
  , ("snd", OTSnd)
  , ("tt", OTPoint)
  , ("Void", OTVoid)
  , ("absurd", OTAbsurd)
  , ("refl", OTRefl)
  , ("transp", OTTransp)
  , ("cast", OTCast)
  , ("castRefl", OTCastRefl)
  ]

obsTokenMap : TokenMap ObsToken
obsTokenMap = toTokenMap [(newline, OTNewline), (spaces, OTIgnore)] ++
  [ (identifier, \s =>
        case lookup s keywords of
          (Just kind) => Tok kind s
          Nothing => Tok OTIdent s)
  ] ++ toTokenMap
  [ (digits, OTNat)
  , (exact "(", OTPOpen)
  , (exact ")", OTPClose)
  , (exact "<", OTAOpen)
  , (exact ">", OTAClose)
  , (exact "=>", OTFatArrow)
  , (exact "->", OTThinArrow)
  , (exact "\\", OTBackslash)
  , (exact "**", OTProduct)
  , (exact ",", OTComma)
  , (exact "~", OTTilde)
  , (exact "=", OTEqual)
  , (exact ":", OTColon)
  ]

newlines : Grammar state ObsToken False ()
newlines = map (\_ => ()) $ many (match OTNewline)

parens : {c : _} -> Grammar state ObsToken c a -> Grammar state ObsToken True a
parens g =
  match OTPOpen *>
  newlines *>
  g <*
  newlines <*
  match OTPClose

ident : Grammar state ObsToken True (WithBounds String)
ident = bounds $ match OTIdent

var : Grammar state ObsToken True Syntax
var = [| Var (ident <* commit) |]

prop : Grammar state ObsToken True Syntax
prop = do
  prop <- bounds $ match OTProp <* commit
  pure (Sort prop.bounds Prop)

set0 : Grammar state ObsToken True Syntax
set0 = do
  set <- bounds $ match OTSet <* commit
  pure (Sort set.bounds (Set 0))

top : Grammar state ObsToken True Syntax
top = do
  unit <- bounds $ match OTUnit <* commit
  pure (Top unit.bounds)

point : Grammar state ObsToken True Syntax
point = do
  point <- bounds $ match OTPoint <* commit
  pure (Point point.bounds)

bottom : Grammar state ObsToken True Syntax
bottom = do
  void <- bounds $ match OTVoid <* commit
  pure (Bottom void.bounds)

atomicNoSet : Grammar state ObsToken True Syntax
atomicNoSet = (var <|> prop <|> top <|> point <|> bottom) <* commit

atomic : Grammar state ObsToken True Syntax
atomic = (atomicNoSet <|> set0) <* commit

set : Grammar state ObsToken True Syntax
set = do
  i <- [| mergeBounds (bounds $ match OTSet) (bounds $ map Set $ option 0 $ match OTNat) |] <* commit
  pure (Sort i.bounds i.val)

lambda : Grammar state ObsToken True (Syntax -> Syntax)
lambda = do
  name <- bounds $ match OTBackslash *> commit *> ident <* newlines <* match OTFatArrow <* newlines
  pure (Lambda name.bounds name.val)

projFst : Grammar state ObsToken True (Syntax -> Syntax)
projFst = do
  fst <- bounds $ match OTFst <* commit
  pure (Fst fst.bounds)

projSnd : Grammar state ObsToken True (Syntax -> Syntax)
projSnd = do
  snd <- bounds $ match OTSnd <* commit
  pure (Snd snd.bounds)

botElim : Grammar state ObsToken True (Syntax -> Syntax -> Syntax)
botElim = do
  absurd <- bounds $ match OTAbsurd <* commit
  pure (Absurd absurd.bounds)

refl : Grammar state ObsToken True (Syntax -> Syntax)
refl = do
  refl <- bounds $ match OTRefl <* commit
  pure (Refl refl.bounds)

transp : Grammar state ObsToken True (Syntax -> Syntax -> Syntax -> Syntax -> Syntax -> Syntax)
transp = do
  transp <- bounds $ match OTTransp <* commit
  pure (Transp transp.bounds)

cast : Grammar state ObsToken True (Syntax -> Syntax -> Syntax -> Syntax)
cast = do
  cast <- bounds $ match OTCast <* commit
  pure (Cast cast.bounds)

castRefl : Grammar state ObsToken True (Syntax -> Syntax)
castRefl = do
  castRefl <- bounds $ match OTCastRefl <* commit
  pure (CastId castRefl.bounds)

-- NOTE: these functions are all total.
partial
pi : WithBounds (WithBounds String, Syntax) -> Grammar state ObsToken True (Syntax -> Syntax)
partial
sigma : WithBounds (WithBounds String, Syntax) -> Grammar state ObsToken True (Syntax -> Syntax)
partial
pair : Grammar state ObsToken True Syntax -> Grammar state ObsToken True Syntax -> Grammar state ObsToken True Syntax
partial
term : Grammar state ObsToken True Syntax
partial
precGteApp : Grammar state ObsToken True Syntax
partial
precGteEqual : Grammar state ObsToken True Syntax
partial
exp' : Grammar state ObsToken True Syntax
partial
exp : Grammar state ObsToken True Syntax
partial
decl : Grammar state ObsToken True (WithBounds String, Syntax)

pi decl = do
  _ <- match OTThinArrow <* commit <* newlines
  Parser.Core.pure (uncurry (Pi decl.bounds) decl.val)

sigma decl = do
  _ <- match OTProduct <* commit <* newlines
  Parser.Core.pure (uncurry (Sigma decl.bounds) decl.val)

pair fst snd = do
  pair <- bounds $ [| MkPair (match OTAOpen *> fst <* newlines <* match OTComma <* newlines) (snd <* newlines <* match OTAClose) |]
  pure (uncurry (Pair pair.bounds) pair.val)

term = atomic <|> pair exp exp <|> parens exp
precGteApp = atomicNoSet <|> set <|> pair exp exp <|> parens exp
precGteEqual =
  map (\ty => let (t, t') = ty.val in maybe t (Equal ty.bounds t) t') $
  bounds [| MkPair precGteApp (optional (newlines *> match OTTilde *> commit *> newlines *> precGteApp)) |]
exp' =
  precGteEqual <|>
  (projFst <*> term) <|>
  (projSnd <*> term) <|>
  (botElim <*> term <*> term) <|>
  (refl <*> term) <|>
  (transp <*> term <*> term <*> term <*> term <*> term) <|>
  (Parser.cast <*> term <*> term <*> term) <|>
  (castRefl <*> term)
exp =
  (bounds (parens decl) <* many (match OTNewline) >>= (\decl => (pi decl <|> sigma decl) <*> exp)) <|>
  (lambda <*> exp) <|>
  (map (\(t, ts) => foldl1 App (t :: ts)) [| MkPair exp' (many term) |])
decl = [| MkPair ident (match OTColon *> commit *> exp) |]

partial
def : Grammar state ObsToken True (WithBounds String, Syntax)
def = [| MkPair ident (match OTEqual *> exp) |]

partial
fullDef : Grammar state ObsToken True Definition
fullDef =
  seq
    [| MkPair decl (match OTNewline *> def) |]
    (\((name, ty), (name', tm)) =>
      if name.val == name'.val
        then pure (MkDefinition {bounds = name.bounds, name = name.val, ty, tm})
        else fatalLoc name.bounds "name mismatch for declaration and definition")

partial
file : Grammar state ObsToken False (List Definition)
file = many (match OTNewline) *> sepBy (some (match OTNewline)) fullDef <* many (match OTNewline) <* eof

postprocess : List (WithBounds ObsToken) -> List (WithBounds ObsToken)
postprocess [] = []
postprocess [tok] = if ignored tok then [] else [tok]
postprocess (tok :: tail@(tok' :: toks)) = case (tok.val.kind, tok'.val.kind) of
  (OTPOpen, OTPClose) => map (\_ => Tok OTUnit "()") (mergeBounds tok tok') :: postprocess toks
  _ => if ignored tok then postprocess tail else tok :: postprocess tail

castErr : Either (List1 (ParsingError tok)) a -> Logging ann a
castErr (Right x) = pure x
castErr (Left errs) = do
  for_ {b = ()} errs
    (\(Error msg bounds) =>
      error $
      maybe
        (irrelevantBounds $ msg)
        (\bounds => MkBounded msg False bounds) $
      bounds)
  abort

partial
export
parse : String -> Logging ann (List Definition)
parse str = do
  let (toks, _, _, "") = lex obsTokenMap str
    | (_, l, c, rem) => inScope "lex" $ fatal "failed to lex input"
  (defs, []) <- inScope "parse" $ castErr $ parse file $ postprocess toks
    | (_, (t :: _)) => inScope "parse" $ fatal "unparsed tokens"
  pure defs