-- TODO Audit how much "outside" stuff could pile up in the continuation.
module Lib.Compile

import Lib.Types
import Lib.Prettier
import Lib.CompileExp
import Lib.TopContext
import Data.String
import Data.Maybe
import Data.Int

data StKind = Plain | Return | Assign String

JSStmt : StKind -> U
JSExp : U

data JAlt : U where
  JConAlt : ∀ e. String -> JSStmt e -> JAlt
  JDefAlt : ∀ e. JSStmt e -> JAlt
  JLitAlt : ∀ e. JSExp -> JSStmt e -> JAlt

data JSExp : U where
  LitArray : List JSExp -> JSExp
  LitObject : List (String × JSExp) -> JSExp
  LitString : String -> JSExp
  LitInt : Int -> JSExp
  Apply : JSExp -> List JSExp -> JSExp
  Var : String -> JSExp
  JLam : List String -> JSStmt Return -> JSExp
  JUndefined : JSExp
  Index : JSExp -> JSExp -> JSExp
  Dot : JSExp -> String -> JSExp

data JSStmt : StKind -> U where
  -- Maybe make this a snoc...
  JSnoc : ∀ a. JSStmt Plain -> JSStmt a -> JSStmt a
  JPlain : JSExp -> JSStmt Plain
  JConst : (nm : String) -> JSExp -> JSStmt Plain
  JReturn : JSExp -> JSStmt Return
  JLet : (nm : String) -> JSStmt (Assign nm) -> JSStmt Plain  -- need somebody to assign
  JAssign : (nm : String) -> JSExp -> JSStmt (Assign nm)
  -- TODO - switch to Int tags
  -- FIXME add e to JAlt (or just drop it?)
  JCase : ∀ a. JSExp -> List JAlt -> JSStmt a
  -- throw can't be used
  JError : ∀ a. String -> JSStmt a

Cont : StKind → U
Cont e = JSExp -> JSStmt e

-- JSEnv contains `Var` for binders or `Dot` for destructured data. It
-- used to translate binders
record JSEnv where
  constructor MkEnv
  jsenv : List JSExp
  depth : Int

-- this was like this, are we not using depth?
push : JSEnv -> JSExp -> JSEnv
push (MkEnv env depth) exp = MkEnv (exp :: env) depth

emptyJSEnv : JSEnv
emptyJSEnv = MkEnv Nil 0

litToJS : Literal -> JSExp
litToJS (LString str) = LitString str
litToJS (LChar c) = LitString $ pack (c :: Nil)
litToJS (LInt i) = LitInt i

-- Stuff nm.h1, nm.h2, ... into environment
-- TODO consider JSExp instead of nm, so we can have $foo.h1 instead of assigning a sc.
mkEnv : String -> Int -> JSEnv -> List String -> JSEnv
mkEnv nm k env Nil = env
mkEnv nm k env (x :: xs) = mkEnv nm (1 + k) (push env (Dot (Var nm) "h\{show k}")) xs

envNames : Env -> List String

-- given a name, find a similar one that doesn't shadow in Env
freshName : String -> JSEnv -> String
freshName nm env = if free env.jsenv nm then nm else go nm 1
  where
    free : List JSExp -> String -> Bool
    free Nil nm = True
    free (Var n :: xs) nm = if n == nm then False else free xs nm
    free (_ :: xs) nm = free xs nm

    go : String -> Int -> String
    go nm k = let nm' = "\{nm}\{show k}" in if free env.jsenv nm' then nm' else go nm (1 + k)

freshName' : String -> JSEnv -> (String × JSEnv)
freshName' nm env =
  let nm' = freshName nm env -- "\{nm}$\{show $ length env}"
      env' = push env (Var nm')
  in (nm', env')

freshNames : List String -> JSEnv -> (List String × JSEnv)
freshNames nms env = go nms env Lin
  where
    go : List Name -> JSEnv -> SnocList Name -> (List String × JSEnv)
    go Nil env acc = (acc <>> Nil, env)
    go (n :: ns) env acc =
      let (n', env') = freshName' n env
      in go ns env' (acc :< n')

-- This is inspired by A-normalization, look into the continuation monad
-- There is an index on JSStmt, adopted from Stefan Hoeck's code.
--
-- Here we turn a Term into a statement (which may be a sequence of statements), there
-- is a continuation, which turns the final JSExpr into a JSStmt, and the function returns
-- a JSStmt, wrapping recursive calls in JSnoc if necessary.
termToJS : ∀ e. JSEnv -> CExp -> Cont e -> JSStmt e
termToJS env (CBnd k) f = case getAt (cast k) env.jsenv of
  (Just e) => f e
  Nothing => fatalError "Bad bounds"
termToJS env CErased f = f JUndefined
termToJS env (CLam nm t) f =
  let (nm',env') = freshName' nm env -- "\{nm}$\{show $ length env}"
  in f $ JLam (nm' :: Nil) (termToJS env' t JReturn)
termToJS env (CFun nms t) f =
  let (nms', env') = freshNames nms env
  in f $ JLam nms' (termToJS env' t JReturn)
termToJS env (CRef nm) f = f $ Var nm
termToJS env (CMeta k) f = f $ LitString "META \{show k}"
termToJS env (CLit lit) f = f (litToJS lit)
-- if it's a var, just use the original
termToJS env (CLet nm (CBnd k) u) f = case getAt (cast k) env.jsenv of
  Just e  => termToJS (push env e) u f
  Nothing => fatalError "bad bounds"
termToJS env (CLet nm t u) f =
  let nm' = freshName nm env
      env' = push env (Var nm')
  -- If it's a simple term, use const
  in case termToJS env t (JAssign nm') of
    (JAssign _ exp) => JSnoc (JConst nm' exp) (termToJS env' u f)
    t' => JSnoc (JLet nm' t') (termToJS env' u f)
termToJS env (CLetRec nm t u) f =
  let nm' = freshName nm env
      env' = push env (Var nm')
  -- If it's a simple term, use const
  in case termToJS env' t (JAssign nm') of
    (JAssign _ exp) => JSnoc (JConst nm' exp) (termToJS env' u f)
    t' => JSnoc (JLet nm' t') (termToJS env' u f)

termToJS env (CApp t args etas) f = termToJS env t (\ t' => (argsToJS t' args Lin f)) --  (f (Apply t' args'))))
  where
    etaExpand : JSEnv -> Nat -> SnocList JSExp -> JSExp -> JSExp
    etaExpand env Z args tm = Apply tm (args <>> Nil)
    etaExpand env (S etas) args tm =
      let nm' = freshName "eta" env
          env' = push env (Var nm')
       in JLam (nm' :: Nil) $ JReturn $ etaExpand (push env (Var nm')) etas (args :< Var nm') tm

    argsToJS : ∀ e. JSExp -> List CExp -> SnocList JSExp -> (JSExp -> JSStmt e) -> JSStmt e
    argsToJS tm Nil acc k = k (etaExpand env (cast etas) acc tm)
    -- k (acc <>> Nil)
    argsToJS tm (x :: xs) acc k = termToJS env x (\ x' => argsToJS tm xs (acc :< x') k)


termToJS {e} env (CCase t alts) f =
  -- need to assign the scrutinee to a variable (unless it is a var already?)
  -- and add (Bnd -> JSExpr map)
  -- TODO default case, let's drop the extra field.

  termToJS env t $ \case
    (Var nm) => maybeCaseStmt env nm alts
    t' => do
      -- TODO refactor nm to be a JSExp with Var{} or Dot{}
      -- FIXME sc$ seemed to shadow something else, lets get this straightened out
      -- we need freshName names that are not in env (i.e. do not play in debruijn)
      let nm = "_sc$\{show env.depth}"
      let env' = MkEnv env.jsenv (1 + env.depth)
      JSnoc (JConst nm t') (maybeCaseStmt env' nm alts)

  where
    termToJSAlt : JSEnv -> String -> CAlt -> JAlt
    termToJSAlt env nm (CConAlt name args u) = JConAlt name (termToJS (mkEnv nm 0 env args) u f)
    -- intentionally reusing scrutinee name here
    termToJSAlt env nm (CDefAlt u) = JDefAlt (termToJS  (env) u f)
    termToJSAlt env nm (CLitAlt lit u) = JLitAlt (litToJS lit) (termToJS env u f)

    maybeCaseStmt : JSEnv -> String -> List CAlt -> JSStmt e
    -- If there is a single alt, assume it matched
    maybeCaseStmt env nm ((CConAlt _ args u) :: Nil) = (termToJS (mkEnv nm 0 env args) u f)
    maybeCaseStmt env nm alts@(CLitAlt _ _  :: _) =
        (JCase (Var nm) (map (termToJSAlt env nm) alts))
    maybeCaseStmt env nm alts =
          (JCase (Dot (Var nm) "tag") (map (termToJSAlt env nm) alts))

jsKeywords : List String
jsKeywords = (
  "break" :: "case" :: "catch" :: "continue" :: "debugger" :: "default" :: "delete" :: "do" :: "else" ::
  "finally" :: "for" :: "function" :: "if" :: "in" :: "instanceof" :: "new" :: "return" :: "switch" ::
  "this" :: "throw" :: "try" :: "typeof" :: "var" :: "void" :: "while" :: "with" ::
  "class" :: "const" :: "enum" :: "export" :: "extends" :: "import" :: "super" ::
  "implements" :: "class" :: "let" :: "package" :: "private" :: "protected" :: "public" ::
  "static" :: "yield" ::
  "null" :: "true" :: "false" ::
  -- might not be a big issue with namespaces on names now.
  "String" :: "Number" :: "Array" :: "BigInt" :: Nil)


-- escape identifiers for js
jsIdent : String -> Doc
jsIdent id = if elem id jsKeywords then text ("$" ++ id) else text $ pack $ fix (unpack id)
  where
    fix : List Char -> List Char
    fix Nil = Nil
    fix (x :: xs) =
      if isAlphaNum x || x == '_' then
        x :: fix xs
      -- make qualified names more readable
      else if x == '.' then '_' :: fix xs
      else if x == '$' then
        '$' :: '$' :: fix xs
      else
        '$' :: (toHex (cast x)) ++ fix xs

stmtToDoc : ∀ e. JSStmt e -> Doc


expToDoc : JSExp -> Doc
expToDoc (LitArray xs) = fatalError "TODO - LitArray to doc"
expToDoc (LitObject xs) = text "{" <+> folddoc (\ a e => a ++ text ", " <+/> e) (map entry xs) <+> text "}"
  where
    entry : (String × JSExp) -> Doc
    -- TODO quote if needed
    entry (nm, exp) = jsIdent nm ++ text ":" <+> expToDoc exp

expToDoc (LitString str) = text $ quoteString str
expToDoc (LitInt i) = text $ show i
-- TODO add precedence
expToDoc (Apply x@(JLam _ _) xs) = text "(" ++ expToDoc x ++ text ")" ++ text "(" ++ nest 2 (commaSep (map expToDoc xs)) ++ text ")"
expToDoc (Apply x xs) = expToDoc x ++ text "(" ++ nest 2 (commaSep (map expToDoc xs)) ++ text ")"
expToDoc (Var nm) = jsIdent nm
expToDoc (JLam nms (JReturn exp)) = text "(" <+> commaSep (map jsIdent nms) <+> text ") =>" <+> text "(" ++ expToDoc exp ++ text ")"
expToDoc (JLam nms body) = text "(" <+> commaSep (map jsIdent nms) <+> text ") =>" <+> bracket "{" (stmtToDoc body) "}"
expToDoc JUndefined = text "undefined"
expToDoc (Index obj ix) = expToDoc obj ++ text "(" ++ expToDoc ix ++ text " :: Nil)"
expToDoc (Dot obj nm) = expToDoc obj ++ text "." ++ jsIdent nm

caseBody : ∀ e. JSStmt e -> Doc
caseBody stmt@(JReturn x) = nest 2 (line ++ stmtToDoc stmt)
-- caseBody {e = Return} stmt@(JCase{}) = nest 2 (line ++ stmtToDoc stmt)
caseBody {e} stmt@(JCase _ _) = nest 2 (line ++ stmtToDoc stmt </> text "break;")
caseBody stmt = line ++ text "{" ++ nest 2 (line ++ stmtToDoc stmt </> text "break;") </> text "}"

altToDoc : JAlt -> Doc
altToDoc (JConAlt nm stmt) = text "case" <+> text (quoteString nm) ++ text ":" ++ caseBody stmt
altToDoc (JDefAlt stmt) = text "default" ++ text ":" ++ caseBody stmt
altToDoc (JLitAlt a stmt) = text "case" <+> expToDoc a ++ text ":" ++ caseBody stmt

stmtToDoc (JSnoc x y) = stmtToDoc x </> stmtToDoc y
stmtToDoc (JPlain x) = expToDoc x ++ text ";"
-- I might not need these split yet.
stmtToDoc (JLet nm body) = text "let" <+> jsIdent nm ++ text ";" </> stmtToDoc body
stmtToDoc (JAssign nm expr) = jsIdent nm <+> text "=" <+> expToDoc expr ++ text ";"
stmtToDoc (JConst nm x) = text "const" <+> jsIdent nm <+> nest 2 (text "=" <+/> expToDoc x ++ text ";")
stmtToDoc (JReturn x) = text "return" <+> expToDoc x ++ text ";"
stmtToDoc (JError str) = text "throw new Error(" ++ text (quoteString str) ++ text ");"
stmtToDoc (JCase sc alts) =
  text "switch (" ++ expToDoc sc ++ text ")" <+> bracket "{" (stack $ map altToDoc alts) "}"

mkArgs : Nat -> List String -> List String
mkArgs Z acc = acc
mkArgs (S k) acc = mkArgs k ("h\{show k}" :: acc)

dcon : QName -> Nat -> Doc
dcon qn@(QN ns nm) Z = stmtToDoc $ JConst (show qn) $ LitObject (("tag", LitString nm) :: Nil)
dcon qn@(QN ns nm) arity =
  let args = mkArgs arity Nil
      obj = ("tag", LitString nm) :: map (\x => (x, Var x)) args
  in stmtToDoc $ JConst (show qn) (JLam args (JReturn (LitObject obj)))

-- use iife to turn stmts into expr
maybeWrap : JSStmt Return -> JSExp
maybeWrap (JReturn exp) = exp
maybeWrap stmt = Apply (JLam Nil stmt) Nil

entryToDoc : TopEntry -> M Doc
entryToDoc (MkEntry _ name ty (Fn tm)) = do
  debug $ \ _ => "compileFun \{render 90 $ pprint Nil tm}"
  ct <- compileFun tm
  let exp = maybeWrap $ termToJS emptyJSEnv ct JReturn
  pure $ text "const" <+> jsIdent (show name) <+> text "=" <+/> expToDoc exp ++ text ";"
entryToDoc (MkEntry _ name type Axiom) = pure $ text ""
entryToDoc (MkEntry _ name type (TCon strs)) = pure $ dcon name (piArity type)
entryToDoc (MkEntry _ name type (DCon arity str)) = pure $ dcon name (cast arity)
entryToDoc (MkEntry _ name type PrimTCon) = pure $ dcon name (piArity type)
entryToDoc (MkEntry _ name _ (PrimFn src _)) = pure $ text "const" <+> jsIdent (show name) <+> text "=" <+> text src


process : (List QName × List Doc) -> QName -> M (List QName × List Doc)
walkTm : Tm -> (List QName × List Doc) -> M (List QName × List Doc)
walkAlt : (List QName × List Doc) -> CaseAlt -> M (List QName × List Doc)

walkAlt acc (CaseDefault t) = walkTm t acc
walkAlt acc (CaseCons name args t) = walkTm t acc
walkAlt acc (CaseLit lit t) = walkTm t acc


walkTm (Ref x nm y) acc = process acc nm
walkTm (Lam x str _ _ t) acc = walkTm t acc
walkTm (App x t u) acc = walkTm u acc >>= walkTm t
walkTm (Pi x str icit y t u) acc = walkTm u acc >>= walkTm t
walkTm (Let x str t u) acc = walkTm u acc >>= walkTm t
walkTm (LetRec x str _ t u) acc = walkTm u acc >>=  walkTm t
walkTm (Case x t alts) acc = foldlM walkAlt acc alts
walkTm _ acc = pure acc

-- This version (call `reverse ∘ snd <$> process "main" (Nil × Nil)`) will do dead
-- code elimination, but the Prelude js primitives are reaching for
-- stuff like True, False, MkUnit, fs which get eliminated
process (done,docs) nm = do
  let (False) = elem nm done | _ => pure (done,docs)
  top <- get
  case lookup nm top of
    Nothing => error emptyFC "\{show nm} not in scope"
    Just entry@(MkEntry _ name ty (PrimFn src used)) => do
      (done,docs) <- foldlM assign (nm :: done, docs) used
      edoc <- entryToDoc entry
      pure (done, edoc :: docs)
    Just (MkEntry _ name ty (Fn tm)) => do
                debug $ \ _ => "compileFun \{render 90 $ pprint Nil tm}"
                ct <- compileFun tm
                -- If ct has zero arity and is a compount expression, this fails..
                let exp = maybeWrap $ termToJS emptyJSEnv ct JReturn
                let doc = text "const" <+> jsIdent (show name) <+> text "=" <+/> expToDoc exp ++ text ";"
                (done,docs) <- walkTm tm (nm :: done, docs)
                pure (done, doc :: docs)
    Just entry => do
      edoc <- entryToDoc entry
      pure (nm :: done, edoc :: docs)
  where
    assign : (List QName × List Doc) -> String -> M (List QName × List Doc)
    assign (done, docs) nm = do
      top <- get
      case lookupRaw nm top of
        Nothing => pure (done, docs)
        (Just (MkEntry fc name type def)) => do
          let tag = QN Nil nm
          let (False) = elem tag done | _ => pure (done,docs)
          (done,docs) <- process (done, docs) name
          let doc = text "const" <+> jsIdent nm <+> text "=" <+> jsIdent (show name) ++ text ";"
          pure (tag :: done, doc :: docs)




compile : M (List Doc)
compile = do
  top <- get
  case lookupRaw "main" top of
    Just (MkEntry fc name type def) => do
      tmp <- snd <$> process (Nil, Nil) name
      let exec = stmtToDoc $ JPlain $ Apply (Var $ show name) Nil
      pure $ reverse (exec :: tmp)
    -- If there is no main, compile everything for the benefit of the playground
    Nothing => do
      top <- get
      traverse entryToDoc $ map snd $ toList top.defs