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
  | OTContainer
  | OTTag
  | OTPosition
  | OTNext
  | OTSem
  | OTPoint
  | OTVoid
  | OTAbsurd
  | OTRefl
  | OTTransp
  | OTCast
  | OTCastRefl
  -- Special symbols
  | OTUnit
  | OTTriangle
  -- Grouping symbols
  | OTPOpen
  | OTPClose
  | OTAOpen
  | OTAClose
  -- Definition characters
  | OTBackslash
  | OTSlash
  | OTThinArrow
  | OTFatArrow
  | OTProduct
  | OTComma
  | OTTilde
  | OTEqual
  | OTColon
  | OTPipe

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
  OTContainer == OTContainer = True
  OTTag == OTTag = True
  OTPosition == OTPosition = True
  OTNext == OTNext = True
  OTSem == OTSem = 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
  OTTriangle == OTTriangle = 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
  OTPipe == OTPipe = True
  _ == _ = False

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 OTContainer s = ()
  tokValue OTTag s = ()
  tokValue OTPosition s = ()
  tokValue OTNext s = ()
  tokValue OTSem 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 OTTriangle 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 = ()
  tokValue OTPipe 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 OTContainer s) = "Container"
  show (Tok OTTag s) = "tag"
  show (Tok OTPosition s) = "position"
  show (Tok OTNext s) = "next"
  show (Tok OTSem s) = "sem"
  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 OTTriangle 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) = ":"
  show (Tok OTPipe 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)
  , ("Container", OTContainer)
  , ("tag", OTTag)
  , ("position", OTPosition)
  , ("next", OTNext)
  , ("sem", OTSem)
  , ("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)
  , (exact "|", OTPipe)
  ]

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 ** (Vect (S n) (WithBounds Syntax) -> Syntax)))
headForms =
  [ (OTFst, (0 ** uncurry 1 First))
  , (OTSnd, (0 ** uncurry 1 Second))
  , (OTAbsurd, (0 ** uncurry 1 Absurd))
  , (OTRefl, (0 ** uncurry 1 Refl))
  , (OTTransp, (4 ** (uncurry 5 Transp)))
  , (OTCast, (3 ** uncurry 4 Cast))
  , (OTCastRefl, (0 ** uncurry 1 CastId))
  ]

headLambdaForms : List (ObsTokenKind, (n ** ((WithBounds String, WithBounds Syntax) -> Vect n (WithBounds Syntax) -> Syntax)))
headLambdaForms =
  [ (OTIf, (3 ** (uncurry 3 . uncurry If)))
  ]

declForms : List (ObsTokenKind, (WithBounds String -> op (WithBounds Syntax) Syntax 2))
declForms = [(OTThinArrow, Pi), (OTProduct, 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 |]

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 a) -> Grammar state ObsToken True (WithBounds String, a)
decl p = [| MkPair (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 a) -> Grammar state ObsToken True (WithBounds String, a)
lambda p =
  match OTBackslash *> commit *>
  [| MkPair (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
container : 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 *>
        [| f (whitespace *> parens (lambda exp)) (count n (whitespace *> term))|])
      headLambdaForms) <|>
  bounds
    (choiceMap
      (\(hd, (n ** f)) => match hd *> commit *> [| 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)) |]

container = spine

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

exp =
  equals <|>
  bounds (map (uncurry Lambda) $ lambda exp) <|>
  bounds
    (map (\((var, a), (b, f)) => f var a b) $
     [| MkPair
          (parens (decl exp) <* whitespace)
          (choiceMap
            (\(op, f) => match op *> commit *> whitespace *> [| MkPair exp (pure f) |])
            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) |]
    (\((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)
  -- <* eof

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
  (OTAOpen, OTPipe) => map (\_ => (Tok OTTriangle "<|")) (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 True) 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