module Main

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

infix 4 ~~

Name : Type
Name = WithBounds String

(~~) : Name -> Name -> Type
m ~~ n = m.val = n.val

Context : Type
Context = SnocList Name

data IsVar : Nat -> Name -> Context -> Type where
  Here : (eq : m ~~ n) -> IsVar 0 m (ctx :< n)
  There : IsVar k n ctx -> IsVar (S k) n (ctx :< _)

data Term : Context -> Type where
  Var : (k : Nat) -> (0 n : Name) -> (0 prf : IsVar k n ctx) -> Term ctx
  Abs : (n : Name) -> Term (ctx :< n) -> Term ctx
  App : Term ctx -> Term ctx -> Term ctx
  Let : (n : Name) -> Term ctx -> Term (ctx :< n) -> Term ctx

data Value : Context -> Type
data Neutral : Context -> Type
data Closure : Context -> Type

data Env : (local, global : Context) -> Type where
  Lin  : Env [<] global
  (:<) : Env local global -> (p : (Name, Lazy (Value global))) -> Env (local :< fst p) global

data Closure : Context -> Type where
  Close : (n : Name) -> Env local global -> Term (local :< n) -> Closure global

data Value where
  Ntrl : Neutral ctx -> Value ctx
  VAbs : Closure ctx -> Value ctx

data Neutral where
  VVar : (k : Nat) -> (0 n : Name) -> (0 prf : IsVar k n ctx) -> Neutral ctx
  VApp : Neutral ctx -> Lazy (Value ctx) -> Neutral ctx

%name IsVar prf
%name Term t, u
%name Value v
%name Neutral n
%name Closure close
%name Env env

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

index : Env local global -> (k : Nat) -> (0 prf : IsVar k n local) -> Value global
index (env :< (n, v)) 0 (Here eq) = v
index (env :< _) (S k) (There prf) = index env k prf

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

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

data Thins : (src, tgt : Context) -> Type where
  Empty : [<] `Thins` [<]
  Drop : src `Thins` tgt -> src `Thins` tgt :< n
  Keep : src `Thins` tgt -> src :< n `Thins` tgt :< n

%name Thins thin

id : {ctx : Context} -> ctx `Thins` ctx
id {ctx = [<]} = Empty
id {ctx = _ :< _} = Keep id

distance : src `Thins` tgt -> {k : Nat} -> (0 prf : IsVar k n src) -> Nat
distance (Drop thin) prf = S (distance thin prf)
distance (Keep thin) {k = 0} (Here eq) = 0
distance (Keep thin) {k = S k} (There prf) = distance thin prf

0 wknIsVar :
  (prf : IsVar k n src) ->
  (thin : src `Thins` tgt) ->
  IsVar (distance thin prf + k) n tgt
wknIsVar prf Empty impossible
wknIsVar prf (Drop thin) = There (wknIsVar prf thin)
wknIsVar (Here eq) (Keep thin) = Here eq
wknIsVar (There {k} prf) (Keep thin) =
  rewrite sym $ plusSuccRightSucc (distance thin prf) k in
  There (wknIsVar prf thin)

wknTerm : Term src -> src `Thins` tgt -> Term tgt
wknTerm (Var k n prf) thin = Var (distance thin prf + k) n (wknIsVar prf thin)
wknTerm (Abs n t) thin = Abs n (wknTerm t (Keep thin))
wknTerm (App t u) thin = App (wknTerm t thin) (wknTerm u thin)
wknTerm (Let n t u) thin = Let n (wknTerm t thin) (wknTerm u (Keep thin))

wknNtrl : Neutral src -> src `Thins` tgt -> Neutral tgt
wknVal : Value src -> src `Thins` tgt -> Value tgt
wknClose : Closure src -> src `Thins` tgt -> Closure tgt
wknEnv : Env local src -> src `Thins` tgt -> Env local tgt

wknNtrl (VVar k n prf) thin = VVar (distance thin prf + k) n (wknIsVar prf thin)
wknNtrl (VApp n v) thin = VApp (wknNtrl n thin) (wknVal v thin)

wknVal (Ntrl n) thin = Ntrl (wknNtrl n thin)
wknVal (VAbs close) thin = VAbs (wknClose close thin)

wknClose (Close n env t) thin = Close n (wknEnv env thin) t

wknEnv [<] thin = [<]
wknEnv (env :< (n, v)) thin = wknEnv env thin :< (n, wknVal v thin)

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

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 Refl))

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 decEq n.val m.val 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) (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 n t u

parseSpine = foldl1 RApp <$> some parseAtom

parseAtom =
  (RVar <$> (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

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

recomputeName : {ctx : Context} -> (k : Nat) -> (0 _ : IsVar k n ctx) -> Singleton n.val
recomputeName {ctx = _ :< n} 0 (Here eq) = rewrite eq in Val n.val
recomputeName (S k) (There prf) = recomputeName k prf

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

prettySingleton : Pretty a => Singleton {a} x -> Doc ann
prettySingleton (Val x) = pretty x

prettyTerm : {ctx : Context} -> Prec -> Term ctx -> Doc ann
prettyTerm d (Var k _ prf) = prettySingleton $ recomputeName 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.val
  prettyNames (n :: ns) = pretty n.val <+> 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.val <++> 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 (Just n.bounds) "undefined name: \{n.val}"

partial
main : IO ()
main = mainWith getArgs parseStdin

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