module Obs.Parser

import Data.String
import Data.Vect

import Obs.Logging
import Obs.Universe
import Obs.Syntax

import System
import System.File

import Text.Lexer
import Text.Parser

%default total

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

Eq ObsTokenKind where
  OTNewlines == OTNewlines = True
  OTSpaces == OTSpaces = True
  OTIdent == OTIdent = True
  OTProp == OTProp = True
  OTSet == OTSet = True
  OTNat == OTNat = True
  OTFst == OTFst = True
  OTSnd == OTSnd = True
  OTBool == OTBool = True
  OTTrue == OTTrue = True
  OTFalse == OTFalse = True
  OTIf == OTIf = 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
  OTSlash == OTSlash = True
  OTThinArrow == OTThinArrow = True
  OTFatArrow == OTFatArrow = True
  OTProduct == OTProduct = True
  OTComma == OTComma = True
  OTTilde == OTTilde = True
  OTEqual == OTEqual = True
  OTColon == OTColon = True
  _ == _ = False

Show ObsTokenKind where
  show OTNewlines = "\\n"
  show OTSpaces = " "
  show OTIdent = "ident"
  show OTProp = "Prop"
  show OTSet = "Set"
  show OTNat = "nat"
  show OTFst = "fst"
  show OTSnd = "snd"
  show OTBool = "Bool"
  show OTTrue = "True"
  show OTFalse = "False"
  show OTIf = "if"
  show OTPoint = "tt"
  show OTVoid = "Void"
  show OTAbsurd = "absurd"
  show OTRefl = "refl"
  show OTTransp = "transp"
  show OTCast = "cast"
  show OTCastRefl = "castRefl"
  show OTUnit = "()"
  show OTPOpen = "("
  show OTPClose = ")"
  show OTAOpen = "<"
  show OTAClose = ">"
  show OTBackslash = "\\"
  show OTSlash = "/"
  show OTThinArrow = "->"
  show OTFatArrow = "=>"
  show OTProduct = "**"
  show OTComma = ","
  show OTTilde = "~"
  show OTEqual = "="
  show OTColon = ":"

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

  tokValue OTNewlines s = ()
  tokValue OTSpaces s = ()
  tokValue OTIdent s = s
  tokValue OTProp s = ()
  tokValue OTSet s = ()
  tokValue OTNat s = stringToNatOrZ s
  tokValue OTFst s = ()
  tokValue OTSnd s = ()
  tokValue OTBool s = ()
  tokValue OTTrue s = ()
  tokValue OTFalse s = ()
  tokValue OTIf 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 OTSlash 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

Show ObsToken where
  show (Tok OTNewlines s) = "\\n"
  show (Tok OTSpaces s) = " "
  show (Tok OTIdent s) = s
  show (Tok OTProp s) = "Prop"
  show (Tok OTSet s) = "Set"
  show (Tok OTNat s) = s
  show (Tok OTFst s) = "fst"
  show (Tok OTSnd s) = "snd"
  show (Tok OTBool s) = "Bool"
  show (Tok OTTrue s) = "True"
  show (Tok OTFalse s) = "False"
  show (Tok OTIf s) = "if"
  show (Tok OTPoint s) = "tt"
  show (Tok OTVoid s) = "Void"
  show (Tok OTAbsurd s) = "absurd"
  show (Tok OTRefl s) = "refl"
  show (Tok OTTransp s) = "transp"
  show (Tok OTCast s) = "cast"
  show (Tok OTCastRefl s) = "castRefl"
  show (Tok OTUnit s) = "()"
  show (Tok OTPOpen s) = "("
  show (Tok OTPClose s) = ")"
  show (Tok OTAOpen s) = "<"
  show (Tok OTAClose s) = ">"
  show (Tok OTBackslash s) = "\\"
  show (Tok OTSlash s) = "/"
  show (Tok OTThinArrow s) = "->"
  show (Tok OTFatArrow s) = "=>"
  show (Tok OTProduct s) = "**"
  show (Tok OTComma s) = ","
  show (Tok OTTilde s) = "~"
  show (Tok OTEqual s) = "="
  show (Tok OTColon s) = ":"

identifier : Lexer
identifier = alpha <+> many alphaNum

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

obsTokenMap : TokenMap ObsToken
obsTokenMap = toTokenMap [(newlines, OTNewlines), (spaces, OTSpaces)] ++
  [ (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 "/", OTSlash)
  , (exact "**", OTProduct)
  , (exact ",", OTComma)
  , (exact "~", OTTilde)
  , (exact "=", OTEqual)
  , (exact ":", OTColon)
  ]

op : Type -> Type -> Nat -> Type
op a b 0 = b
op a b (S n) = a -> op a b n

uncurry : (n : _) -> op a b n -> Vect n a -> b
uncurry 0     f []        = f
uncurry (S n) f (x :: xs) = uncurry n (f x) xs

termForms : List (ObsTokenKind, Syntax)
termForms =
  [ (OTBool, Bool)
  , (OTTrue, True)
  , (OTFalse, False)
  , (OTPoint, Point)
  , (OTVoid, Bottom)
  , (OTUnit, Top)
  ]

headForms : List (ObsTokenKind, (n ** op (WithBounds Syntax) Syntax (S n)))
headForms =
  [ (OTFst, (0 ** First))
  , (OTSnd, (0 ** Second))
  , (OTAbsurd, (0 ** Absurd))
  , (OTRefl, (0 ** Refl))
  , (OTTransp, (4 ** Transp))
  , (OTCast, (3 ** Cast))
  , (OTCastRefl, (0 ** CastId))
  ]

headLambdaForms : List (ObsTokenKind, (n ** (LambdaForm -> op (WithBounds Syntax) Syntax n)))
headLambdaForms = [(OTIf, (3 ** If))]

declForms : List (ObsTokenKind, (n ** op DeclForm (WithBounds Syntax -> Syntax) (S n)))
declForms =
  [ (OTThinArrow, (0 ** Pi))
  , (OTProduct, (0 ** Sigma))
  ]

count : (n : _) -> Grammar state tok True a -> Grammar state tok (isSucc n) (Vect n a)
count 0     p = pure []
count (S n) p = [| p :: count n p |]

match' : (tok : ObsTokenKind) -> Grammar state ObsToken True (TokType tok)
match' tok = do
  pos <- position
  token <- peek
  let True = token.kind == tok
    | False => failLoc pos "expected \{show tok}"
  match tok

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

whitespace : Grammar state ObsToken True ()
whitespace = map (\_ => ()) $ some (optional (match' OTNewlines) *> match' OTSpaces)

parens : Lazy (Grammar state ObsToken True a) -> Grammar state ObsToken True a
parens p = match' OTPOpen *> optional whitespace *> p <* optional whitespace <* match' OTPClose

var : Grammar state ObsToken True (WithBounds Syntax)
var = map (map Var) ident

noArgUniverse : Grammar state ObsToken True (WithBounds Universe)
noArgUniverse = bounds $
  map (const Prop) (match' OTProp <* commit) <|>
  map (const (Set 0)) (match' OTSet <* commit)

universe : Grammar state ObsToken True (WithBounds Universe)
universe = bounds $
  map (const Prop) (match' OTProp <* commit) <|>
  map Set (match' OTSet *> commit *> option 0 (match' OTNat <* commit))

decl : Lazy (Grammar state ObsToken True (WithBounds Syntax)) ->
  Grammar state ObsToken True DeclForm
decl p = [| MkDecl (ident <* whitespace <* match' OTColon <* whitespace) p |]

pair : Lazy (Grammar state ObsToken True a) -> Grammar state ObsToken True (a, a)
pair p =
  match' OTAOpen *> commit *>
  [| MkPair
      (optional whitespace *> p <* optional whitespace <* match' OTComma <* commit)
      (optional whitespace *> p <* optional whitespace) |]
  <* match' OTAClose <* commit

lambda : Lazy (Grammar state ObsToken True (WithBounds Syntax)) ->
  Grammar state ObsToken True LambdaForm
lambda p =
  match' OTBackslash *> commit *>
  [| MkLambda (ident <* whitespace <* match' OTFatArrow <* whitespace) p |]

partial
noUniversesTerm : Grammar state ObsToken True (WithBounds Syntax)
partial
term : Grammar state ObsToken True (WithBounds Syntax)
partial
head : Grammar state ObsToken True (WithBounds Syntax)
partial
spine : Grammar state ObsToken True (WithBounds Syntax)
partial
equals : Grammar state ObsToken True (WithBounds Syntax)
partial
exp : Grammar state ObsToken True (WithBounds Syntax)

noUniversesTerm =
  parens exp <|>
  var <|>
  bounds (map (uncurry Pair) (pair exp)) <|>
  bounds (choiceMap (\(k, s) => map (const s) (match' k)) termForms)

term = noUniversesTerm <|> map (map Universe) noArgUniverse

head =
  noUniversesTerm <|>
  map (map Universe) universe <|>
  bounds
    (choiceMap
      (\(hd, (n ** f)) =>
        match' hd *> commit *> whitespace *>
        pure (uncurry n . f) <*> parens (lambda exp) <*> count n (whitespace *> term))
      headLambdaForms) <|>
  bounds
    (choiceMap
      (\(hd, (n ** f)) => match' hd *> commit *>
        pure (uncurry (S n) f) <*> count (S n) (whitespace *> term))
      headForms)

spine = map (uncurry $ foldl (\t, u => joinBounds (map (\_ => map (\_ => App t u) u) t))) $
  [| MkPair head (many (whitespace *> term)) |]

equals = map (\(t, u) => maybe t (\u => joinBounds (map (\_ => map (\_ => Equal t u) u) t)) u) $
  [| MkPair spine (optional (whitespace *> match' OTTilde *> commit *> whitespace *> spine)) |]

exp =
  equals <|>
  bounds (map Lambda $ lambda exp) <|>
  bounds
    (choiceMap
      (\(sep, (n ** f)) =>
        pure (uncurry (S n) f) <*>
        count (S n) (parens (decl exp) <* whitespace <* match' sep <* commit <* whitespace) <*>
        exp)
      declForms)

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

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

partial
file : Grammar state ObsToken False (List Definition)
file =
  optional (whitespace *> match' OTNewlines) *>
  sepEndBy (optional whitespace *> match' OTNewlines) (fullDef <* commit)

whitespaceIrrelevant : WithBounds ObsToken -> WithBounds ObsToken
whitespaceIrrelevant a = case (a.val.kind) of
  OTSpaces => removeIrrelevance a
  OTNewlines => removeIrrelevance a
  _ => a

-- Combines tokens for nicer parsing
mergeToks : List (WithBounds ObsToken) -> List (WithBounds ObsToken)
mergeToks (a :: tail@(b :: rest)) = case (a.val.kind, b.val.kind) of
  (OTSpaces, OTNat) => mergeBounds a b :: mergeToks rest
  (OTPOpen, OTPClose) => map (\_ => (Tok OTUnit "()")) (mergeBounds a b) :: mergeToks rest
  _ => a :: mergeToks tail
mergeToks toks = toks

preprocess : List (WithBounds ObsToken) -> List (WithBounds ObsToken)
preprocess = mergeToks . map whitespaceIrrelevant

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

partial
export
parse : String -> Logging ann (List Definition)
parse str = do
  let (toks, _, _, "") = lex obsTokenMap str
    | (toks, l, c, rem) => inScope "lex" $ do
      error "failed to lex input"
      trace (map (\tok => show tok.val) toks)
      trace (show rem)
      abort
  let toks = preprocess toks
  (defs, []) <- inScope "parse" $ castErr $ parse file $ toks
    | (_, ts) => inScope "parse" $ do
      trace (map (show . val) ts)
      fatal "unparsed tokens"
  pure defs