module Main

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

import Data.DPair
import Data.List
import Data.Nat
import Data.Singleton
import Data.String

import Decidable.Equality

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

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

data RawTerm : Type where
  RVar : (n : Name) -> RawTerm
  RAbs : (n : Name) -> RawTerm -> RawTerm
  RApp : RawTerm -> RawTerm -> RawTerm
  RLet : (n : Name) -> RawTerm -> RawTerm -> RawTerm

%name RawTerm t, u

data NameError : Type -> Type where
  Ok : ty -> NameError ty
  Err : List1 Name -> NameError ty

Functor NameError where
  map f (Ok v) = Ok (f v)
  map f (Err errs) = Err errs

Applicative NameError where
  pure = Ok

  Ok f <*> Ok v = Ok (f v)
  Ok f <*> Err ys = Err ys
  Err xs <*> Ok v = Err xs
  Err xs <*> Err ys = Err (xs ++ ys)

lookup : (ctx : Context) -> (n : Name) -> NameError $ Subset Nat (\k => IsVar k n ctx)
lookup [<] n = Err (n ::: [])
lookup (ctx :< m) n =
  case decEquiv n m of
    Yes eq => Ok (Element 0 (Here eq))
    No _ => map (\p => Element (S p.fst) (There p.snd)) (lookup ctx n)

translate : (ctx : Context) -> RawTerm -> NameError (Term ctx)
translate ctx (RVar n) = (\p => Var p.fst n p.snd) <$> lookup ctx n
translate ctx (RAbs n t) = pure (Abs n) <*> translate (ctx :< n) t
translate ctx (RApp t u) = [| App (translate ctx t) (translate ctx u) |]
translate ctx (RLet n t u) = pure (Let n) <*> translate ctx t <*> translate (ctx :< n) u

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

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 Ok term = translate [<] rawTerm
    | Err 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)