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
  -- Special symbols
  | OTUnit
  | OTPoint
  -- Grouping symbols
  | OTPOpen
  | OTPClose
  -- Definition characters
  | 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
  OTPOpen   == OTPOpen   = True
  OTPClose  == OTPClose  = 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 OTPOpen   s = ()
  tokValue OTPClose  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)
  ]

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 "=", 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 = do
  id <- ident
  pure (Var id.bounds id.val)

prop : Grammar state ObsToken True (WithBounds Sort)
prop = do
  p <- bounds $ match OTProp
  pure (map (\_ => Prop) p)

set : Grammar state ObsToken True (WithBounds Sort)
set = do
  s <- bounds $ match OTSet
  i <- bounds $ option 0 (match OTNat)
  pure (map Set $ mergeBounds s i)

sort : Grammar state ObsToken True Syntax
sort = do
  sort <- prop <|> set
  pure (Sort sort.bounds sort.val)

top : Grammar state ObsToken True Syntax
top = map (\v => Top v.bounds) $ bounds $ match OTUnit

point : Grammar state ObsToken True Syntax
point = map (\v => Point v.bounds) $ bounds $ match OTPoint

partial
exp  : Grammar state ObsToken True Syntax
partial
term : Grammar state ObsToken True Syntax

exp = term
term = var <|> sort <|> top <|> point <|> parens exp

partial
decl : Grammar state ObsToken True (WithBounds String, Syntax)
decl = [| MkPair ident (match OTColon *> 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) *> sepEndBy (some (match OTNewline)) fullDef

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