module Obs.Parser

import Data.String

import Obs.Sort
import Obs.Syntax

import System
import System.File

import Text.Lexer
import Text.Parser

%default total

Error : Type -> Type -> Type
Error tok = Either (List1 (ParsingError tok))

report : List1 (ParsingError tok) -> Error tok a
report = Left

data ObsTokenKind
  = OTNewline
  | OTIgnore
  | OTIdent
  -- Keywords
  | OTProp
  | OTSet
  | OTNat
  | OTVoid
  | OTAbsurd
  -- Special symbols
  | OTUnit
  | OTPoint
  -- Grouping symbols
  | OTPOpen
  | OTPClose
  -- Definition characters
  | OTBackslash
  | OTThinArrow
  | OTFatArrow
  | OTEqual
  | OTColon

Eq ObsTokenKind where
  OTNewline == OTNewline = True
  OTIgnore == OTIgnore = True
  OTIdent == OTIdent = True
  OTProp == OTProp = True
  OTSet == OTSet = True
  OTNat == OTNat = True
  OTUnit == OTUnit = True
  OTPoint == OTPoint = True
  OTVoid == OTVoid = True
  OTAbsurd == OTAbsurd = True
  OTPOpen == OTPOpen = True
  OTPClose == OTPClose = True
  OTBackslash == OTBackslash = True
  OTThinArrow == OTThinArrow = True
  OTFatArrow == OTFatArrow = 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 OTUnit s = ()
  tokValue OTPoint s = ()
  tokValue OTVoid s = ()
  tokValue OTAbsurd s = ()
  tokValue OTPOpen s = ()
  tokValue OTPClose s = ()
  tokValue OTBackslash s = ()
  tokValue OTThinArrow s = ()
  tokValue OTFatArrow 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)
  , ("Void", OTVoid)
  , ("absurd", OTAbsurd)
  ]

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 "*", OTPoint)
  , (exact "(", OTPOpen)
  , (exact ")", OTPClose)
  , (exact "=>", OTFatArrow)
  , (exact "->", OTThinArrow)
  , (exact "\\", OTBackslash)
  , (exact "=", OTEqual)
  , (exact ":", OTColon)
  ]

parens : {c : _} -> Grammar state ObsToken c a -> Grammar state ObsToken True a
parens g =
  match OTPOpen *>
  many (match OTNewline) *>
  g <*
  many (match OTNewline) <*
  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 <* match OTFatArrow
  pure (Lambda name.bounds name.val)

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

-- NOTE: these functions are all total.
partial
pi : Grammar state ObsToken True (Syntax -> Syntax)
partial
term : Grammar state ObsToken True Syntax
partial
precGteApp : 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 = do
  decl <- bounds $ parens decl <* match OTThinArrow <* commit
  pure (uncurry (Pi decl.bounds) decl.val)

term = atomic <|> parens exp
precGteApp = atomicNoSet <|> set <|> parens exp
exp' = precGteApp <|> (botElim <*> term <*> term)
exp = (pi <*> 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

partial
export
parse : String -> Error ObsToken (List Definition)
parse str = do
  let (toks, _, _, "") = lex obsTokenMap str
    | (_, l, c, rem) => report (Error "failed to lex input" (Just $ MkBounds { startLine = l, startCol = c, endLine = l, endCol = c }) ::: [])
  (defs, []) <- parse file $ postprocess toks
    | (_, (t :: _)) => report (Error "unparsed tokens" (Just t.bounds) ::: [])
  pure defs

printErr : ParsingError tok -> IO ()
printErr (Error str Nothing)  = putStrLn "error: \{str}"
printErr (Error str (Just b)) = putStrLn
  "error: \{show b.startLine}:\{show b.startCol}-\{show b.endLine}:\{show b.endCol}: \{str}"

reportErrs : List (ParsingError tok) -> IO a
reportErrs []          = exitFailure
reportErrs (e :: errs) = do
  () <- printErr e
  reportErrs errs

partial
export
parseFile : (fname : String) -> IO (List Definition)
parseFile fname = do
  Right str <- readFile fname
    | Left err => do
      () <- putStrLn "error: \{show err}"
      exitFailure
  let Right defs = parse str
    | Left err => reportErrs (forget err)
  pure defs