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
  -- Whitespace
  = OTNewlines
  | OTSpaces
  -- Variables
  | OTIdent
  -- Universes
  | OTProp
  | OTSet
  | OTNat
  -- Term Forms
  | OTBool
  | OTTrue
  | OTFalse
  | OTUnit
  | OTPoint
  | OTVoid
  -- Head Forms
  | OTFst
  | OTSnd
  | OTMkContainer
  | OTShape
  | OTPosition
  | OTNextIndex
  | OTExtension
  | OTBox
  | OTMkBox
  | OTUnbox
  | OTAbsurd
  | OTRefl
  | OTTransp
  | OTCast
  | OTCastRefl
  -- Lambda Head Forms
  | OTIf
  -- Decl Forms
  | OTThinArrow
  | OTProduct
  -- Special Forms
  | OTContainer
  -- Lambda Form
  | OTBackslash
  | OTFatArrow
  -- Grouping Symbols
  | OTPOpen
  | OTPClose
  | OTAOpen
  | OTAClose
  -- Other Symbols
  | OTPostulate
  | 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

  OTBool == OTBool = True
  OTTrue == OTTrue = True
  OTFalse == OTFalse = True
  OTUnit == OTUnit = True
  OTPoint == OTPoint = True
  OTVoid == OTVoid = True

  OTFst == OTFst = True
  OTSnd == OTSnd = True
  OTMkContainer == OTMkContainer = True
  OTShape == OTShape = True
  OTPosition == OTPosition = True
  OTNextIndex == OTNextIndex = True
  OTExtension == OTExtension = True
  OTBox == OTBox = True
  OTMkBox == OTMkBox = True
  OTUnbox == OTUnbox = True
  OTAbsurd == OTAbsurd = True
  OTRefl == OTRefl = True
  OTTransp == OTTransp = True
  OTCast == OTCast = True
  OTCastRefl == OTCastRefl = True

  OTIf == OTIf = True

  OTThinArrow == OTThinArrow = True
  OTProduct == OTProduct = True

  OTContainer == OTContainer = True

  OTBackslash == OTBackslash = True
  OTFatArrow == OTFatArrow = True

  OTPOpen == OTPOpen = True
  OTPClose == OTPClose = True
  OTAOpen == OTAOpen = True
  OTAClose == OTAClose = True

  OTPostulate == OTPostulate = 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 OTBool = "Bool"
  show OTTrue = "True"
  show OTFalse = "False"
  show OTUnit = "()"
  show OTPoint = "tt"
  show OTVoid = "Void"

  show OTFst = "fst"
  show OTSnd = "snd"
  show OTMkContainer = "MkContainer"
  show OTShape = "Shape"
  show OTPosition = "Position"
  show OTNextIndex = "nextIndex"
  show OTExtension = "extension"
  show OTBox = "Box"
  show OTMkBox = "box"
  show OTUnbox = "unbox"
  show OTAbsurd = "absurd"
  show OTRefl = "refl"
  show OTTransp = "transp"
  show OTCast = "cast"
  show OTCastRefl = "castRefl"

  show OTIf = "if"

  show OTThinArrow = "->"
  show OTProduct = "**"

  show OTContainer = "Container"

  show OTBackslash = "\\"
  show OTFatArrow = "=>"

  show OTPOpen = "("
  show OTPClose = ")"
  show OTAOpen = "<"
  show OTAClose = ">"

  show OTPostulate = "postulate"
  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 OTBool s = ()
  tokValue OTTrue s = ()
  tokValue OTFalse s = ()
  tokValue OTUnit s = ()
  tokValue OTPoint s = ()
  tokValue OTVoid s = ()

  tokValue OTFst s = ()
  tokValue OTSnd s = ()
  tokValue OTMkContainer s = ()
  tokValue OTShape s = ()
  tokValue OTPosition s = ()
  tokValue OTNextIndex s = ()
  tokValue OTExtension s = ()
  tokValue OTBox s = ()
  tokValue OTMkBox s = ()
  tokValue OTUnbox s = ()
  tokValue OTAbsurd s = ()
  tokValue OTRefl s = ()
  tokValue OTTransp s = ()
  tokValue OTCast s = ()
  tokValue OTCastRefl s = ()

  tokValue OTIf s = ()

  tokValue OTThinArrow s = ()
  tokValue OTProduct s = ()

  tokValue OTContainer s = ()

  tokValue OTBackslash s = ()
  tokValue OTFatArrow s = ()

  tokValue OTPOpen s = ()
  tokValue OTPClose s = ()
  tokValue OTAOpen s = ()
  tokValue OTAClose s = ()

  tokValue OTPostulate s = ()
  tokValue OTComma s = ()
  tokValue OTTilde s = ()
  tokValue OTEqual s = ()
  tokValue OTColon s = ()

ObsToken : Type
ObsToken = Token ObsTokenKind

Show ObsToken where
  show (Tok OTIdent s) = s
  show (Tok OTNat s) = s

  show (Tok kind s) = show kind

identifier : Lexer
identifier = alpha <+> many alphaNum

keywords : List (String, ObsTokenKind)
keywords =
  [ ("Prop", OTProp)
  , ("Set", OTSet)

  , ("Bool", OTBool)
  , ("True", OTTrue)
  , ("False", OTFalse)
  , ("tt", OTPoint)
  , ("Void", OTVoid)

  , ("fst", OTFst)
  , ("snd", OTSnd)
  , ("MkContainer", OTMkContainer)
  , ("Shape", OTShape)
  , ("Position", OTPosition)
  , ("nextIndex", OTNextIndex)
  , ("extension", OTExtension)
  , ("Box", OTBox)
  , ("box", OTMkBox)
  , ("unbox", OTUnbox)
  , ("absurd", OTAbsurd)
  , ("refl", OTRefl)
  , ("transp", OTTransp)
  , ("cast", OTCast)
  , ("castRefl", OTCastRefl)

  , ("if", OTIf)

  , ("Container", OTContainer)

  , ("postulate", OTPostulate)
  ]

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 "->", OTThinArrow)
  , (exact "**", OTProduct)

  , (exact "\\", OTBackslash)
  , (exact "=>", OTFatArrow)

  , (exact "(", OTPOpen)
  , (exact ")", OTPClose)
  , (exact "<", OTAOpen)
  , (exact ">", OTAClose)

  , (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)
  , (OTUnit, Top)
  , (OTPoint, Point)
  , (OTVoid, Bottom)
  ]

headForms : List (ObsTokenKind, (n ** op (WithBounds Syntax) Syntax (S n)))
headForms =
  [ (OTFst, (0 ** First))
  , (OTSnd, (0 ** Second))
  , (OTMkContainer, (2 ** MkContainer))
  , (OTShape, (0 ** Shape))
  , (OTPosition, (0 ** Position))
  , (OTNextIndex, (0 ** Next))
  , (OTBox, (0 ** Box))
  , (OTMkBox, (0 ** MkBox))
  , (OTUnbox, (0 ** Unbox))
  , (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 = [(OTExtension, (1 ** Sem)), (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))

termUniverse : Grammar state ObsToken True (WithBounds Universe)
termUniverse = noArgUniverse <|> parens universe

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 (WithBounds Syntax)) ->
  Grammar state ObsToken True (WithBounds Syntax)
pair p = bounds $
  match' OTAOpen *> commit *>
  [| Pair
      (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 <|>
  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) <|>
  bounds
    (match OTContainer *> commit *>
      [| Container
          (whitespace *> term)
          (whitespace *> term)
          (whitespace *> termUniverse)
          (whitespace *> termUniverse)
      |])

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
param : Grammar state ObsToken True Parameter
param = map (\(MkDecl name ty) => MkParameter {name, ty}) $
  match OTPostulate *> commit *> whitespace *> decl 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, ty, tm})
        else fatalLoc name.bounds "name mismatch' for declaration and definition")

partial
file : Grammar state ObsToken False (List (Either Parameter Definition))
file = many (whitespace *> match' OTNewlines) *>
  sepEndBy (optional whitespace *> match' OTNewlines) (choose param 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 (Either Parameter 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