module Main

import CC.Name
import CC.Term
import CC.Term.Raw
import CC.Thinning

import Data.String

import System
import System.File.ReadWrite
import System.File.Virtual

import Text.Lexer
import Text.Parser
import Text.PrettyPrint.Prettyprinter
import Text.PrettyPrint.Prettyprinter.Render.Terminal

--------------------------------------------------------------------------------

app : Closure ctx -> Lazy (Value ctx) -> Value ctx
eval : Env local global -> Term local -> Value global

app (Close n env t) v = assert_total $ eval (env :< (n, v)) t

eval env (Var k n prf) = index env k prf
eval env (Abs n t) = VAbs (Close n env t)
eval env (App t u) =
  case eval env t of
    Ntrl n => Ntrl $ VApp n (eval env u)
    VAbs close => app close (eval env u)
eval env (Let n t u) = eval (env :< (n, eval env t)) u

--------------------------------------------------------------------------------

quote : {ctx : _} -> Value ctx -> Term ctx
quoteNtrl : {ctx : _} -> Neutral ctx -> Term ctx

quote (Ntrl n) = quoteNtrl n
quote (VAbs close@(Close n _ _)) =
  Abs n $
  assert_total $
  quote $
  app (wknClose close (drop id)) (Ntrl $ VVar 0 n $ Here $ reflexive {rel = (~~)})

quoteNtrl (VVar k n prf) = Var k n prf
quoteNtrl (VApp n v) = App (quoteNtrl n) (quote v)

normalise : {global : _} -> Env local global -> Term local -> Term global
normalise env t = quote $ eval env t

--------------------------------------------------------------------------------

putErrLn : HasIO io => String -> io ()
putErrLn str = do
  Right () <- fPutStrLn stderr str
    | Left e => pure ()
  pure ()

putBoundErrLn : HasIO io => Maybe Bounds -> String -> io ()
putBoundErrLn Nothing str = putErrLn str
putBoundErrLn (Just bounds) str =
  let startLine = padLeft 3 '0' $ show (1 + bounds.startLine) in
  let startCol = padLeft 2 '0' $ show bounds.startCol in
  let endLine = padLeft 3 '0' $ show (1 + bounds.endLine) in
  let endCol = padLeft 2 '0' $ show bounds.endCol in
  putErrLn "\{startLine}:\{startCol}--\{endLine}:\{endCol}: \{str}"

data CCTokenKind
  = TIgnore
  | TIdent
  -- Keywords
  | TLet
  -- Symbols
  | TLambda
  | TPOpen
  | TPClose
  | TComma
  | TFatArrow
  | TDef
  | TSemi

%name CCTokenKind k

Eq CCTokenKind where
  TIgnore == TIgnore = True
  TIdent == TIdent = True
  TLet == TLet = True
  TLambda == TLambda = True
  TPOpen == TPOpen = True
  TPClose == TPClose = True
  TComma == TComma = True
  TFatArrow == TFatArrow = True
  TDef == TDef = True
  TSemi == TSemi = True
  _ == _ = False

TokenKind CCTokenKind where
  TokType TIdent = String
  TokType _ = ()

  tokValue TIgnore s = ()
  tokValue TIdent s = s
  tokValue TLet s = ()
  tokValue TLambda s = ()
  tokValue TPOpen s = ()
  tokValue TPClose s = ()
  tokValue TComma s = ()
  tokValue TFatArrow s = ()
  tokValue TDef s = ()
  tokValue TSemi s = ()

CCToken = Token CCTokenKind

ignored : List Lexer
ignored = [spaces, lineComment (exact "--"), blockComment (exact "{-") (exact "-}")]

identifier : Lexer
identifier = alpha <+> many alphaNum

keywords : List (String, CCTokenKind)
keywords = [("let", TLet)]

tokenMap : TokenMap CCToken
tokenMap =
  toTokenMap (map (, TIgnore) ignored) ++
  [(identifier, \s =>
    case lookup s keywords of
      Just kind => Tok kind s
      Nothing => Tok TIdent s)
  ] ++ toTokenMap [
    (exact "\\", TLambda),
    (exact "\x03BB", TLambda),
    (exact "(", TPOpen),
    (exact ")", TPClose),
    (exact ",", TComma),
    (exact "=>", TFatArrow),
    (exact "=", TDef),
    (exact ";", TSemi)
  ]

%tcinline
parens : {b : Bool} -> Grammar state CCToken b ty -> Grammar state CCToken True ty
parens p = match TPOpen >> commit >> p >>= (\t => match TPClose *> commit *> pure t)

parseTerm : Grammar state CCToken True RawTerm
parseLambda : Grammar state CCToken True RawTerm
parseLet : Grammar state CCToken True RawTerm
parseSpine : Grammar state CCToken True RawTerm
parseAtom : Grammar state CCToken True RawTerm

parseTerm = parseLambda <|> parseLet <|> parseSpine

parseLambda = do
  match TLambda
  commit
  names <- sepBy1 (match TComma >> commit) (MkName <$> (bounds $ match TIdent))
  commit
  match TFatArrow
  commit
  t <- parseTerm
  pure $ foldr1By RAbs (flip RAbs t) names

parseLet = do
  match TLet
  commit
  n <- bounds $ match TIdent
  commit
  match TDef
  commit
  t <- parseTerm
  match TSemi
  commit
  u <- parseTerm
  pure $ RLet (MkName n) t u

parseSpine = foldl1 RApp <$> some parseAtom

parseAtom =
  (RVar <$> MkName <$> (bounds $ match TIdent <* commit)) <|>
  (parens parseTerm)

parseSrc : Grammar state CCToken True RawTerm
parseSrc = parseTerm <* eof

parseString : String -> Either (List1 (ParsingError CCToken)) RawTerm
parseString str = do
  let (tokens, _, _, "") = lex tokenMap str
    | (_, line, col, rest) =>
      let bounds = MkBounds { startCol = col, startLine = line, endCol = col, endLine = line } in
      Left (Error "failed to lex input" (Just bounds) ::: [])
  (val, []) <- parse parseSrc $ filter (\tok => tok.val.kind /= TIgnore) tokens
    | (_, fst :: rest) => Left (Error "unexpected token" (Just fst.bounds) ::: [])
  pure val

parse : String -> IO RawTerm
parse str =
  case parseString str of
    Left errs => for_ errs putErr >> exitFailure
    Right v => pure v
  where
  putErr : ParsingError CCToken -> IO ()
  putErr (Error str bounds) = putBoundErrLn bounds ("parse error: \{str}" )

partial
parseStdin : IO RawTerm
parseStdin = do
  Right str <- fRead stdin
    | Left err => putErrLn "failed to read input: \{show err}" >> exitFailure
  parse str

--------------------------------------------------------------------------------

startPrec, leftAppPrec, appPrec : Prec
startPrec = Open
leftAppPrec = Backtick
appPrec = App

prettyVar : (ctx : Context) -> (k : Nat) -> (0 _ : IsVar k n ctx) -> Doc ann
prettyVar (_ :< n) 0 (Here eq) = pretty n
prettyVar (ctx :< _) (S k) (There prf) = prettyVar ctx k prf

prettyTerm : {ctx : Context} -> Prec -> Term ctx -> Doc ann
prettyTerm d (Var k _ prf) = prettyVar ctx k prf
prettyTerm d t@(Abs _ _) =
  let (names ** t) = getNames t in
  parenthesise (d > startPrec) $ group $
    backslash <+> prettyNames names <++> pretty "=>" <+> softline <+>
    (assert_total $ prettyTerm startPrec t)
  where
  getNames : forall ctx. Term ctx -> (local ** Term (ctx <>< local))
  getNames (Abs n t) = let (names ** t) = getNames t in (n :: names ** t)
  getNames t = ([] ** t)

  prettyNames : List Name -> Doc ann
  prettyNames [] = neutral
  prettyNames [n] = pretty n
  prettyNames (n :: ns) = pretty n <+> comma <+> line <+> prettyNames ns
prettyTerm d (App t u) =
  parenthesise (d >= appPrec) $ group $
  prettyTerm leftAppPrec t <++> prettyTerm appPrec u
prettyTerm d (Let n t u) =
  parenthesise (d > startPrec) $ group $ align $
    pretty "let" <++>
    (group $ align $ hang 2 $ pretty n <++> pretty "=" <+> line <+> prettyTerm startPrec t) <+>
    pretty ";" <+> line <+>
    prettyTerm startPrec u

{ctx : Context} -> Pretty (Term ctx) where
  prettyPrec = prettyTerm

--------------------------------------------------------------------------------

-- -- ("\\" works for lambda as well)
-- ex = main' "nf" $ unlines [
--   "let λ λ 1 (1 (1 (1 (1 0))));",    -- five = λ s z. s (s (s (s (s z))))
--   "let λ λ λ λ 3 1 (2 1 0);",        -- add  = λ a b s z. a s (b s z)
--   "let λ λ λ λ 3 (2 1) 0;",          -- mul  = λ a b s z. a (b s) z
--   "let 1 2 2;",                      -- ten  = add five five
--   "let 1 0 0;",                      -- hundred = mul ten ten
--   "let 2 1 0;",                      -- thousand = mul ten hundred
--   "0"                                -- thousand

helpMsg : String
helpMsg = unlines [
  "usage: cc-eval [--help|nf]",
  "  --help : display this message",
  "  nf     : read expression from stdin, print its normal form"
  ]

mainWith : IO (List String) -> IO RawTerm -> IO ()
mainWith getOpt getTerm = do
  [_, "nf"] <- getOpt
    | [_, "--help"] => putStrLn helpMsg
    | _ => putErrLn helpMsg >> exitFailure
  rawTerm <- getTerm
  let Right term = runNameError $ translate [<] rawTerm
    | Left errs => for_ errs putErr >> exitFailure
  let nf : Term [<] = normalise [<] term
  putDoc (pretty nf)
  where
  putErr : Name -> IO ()
  putErr n = putBoundErrLn (bounds n) "undefined name: \{n}"

partial
main : IO ()
main = mainWith getArgs parseStdin

main' : String -> String -> IO ()
main' mode src = mainWith (pure [mode]) (parse src)