module Lib.CompileScheme

import Prelude
import Lib.Common
import Lib.Types
import Lib.Eval
import Lib.Prettier
import Lib.CompileExp
import Lib.TopContext
import Lib.LiftWhere
import Lib.LiftLambda
import Lib.TCO
import Lib.Ref2
import Lib.Erasure
import Data.String
import Data.Int
import Data.SortedMap
import Data.IORef
import Data.SnocList

-- FIXME SnocList
SCEnv : U
SCEnv = List String

getEnv : Int → SCEnv → String
getEnv ix env = case getAt' ix env of
  Just e => e
  Nothing => fatalError "Bad bounds \{show ix}"

scbind : String → SCEnv → (String × SCEnv)
scbind nm env = let
  x = length' env
  nm' = "nm-\{show x}"
  in (nm', nm' :: env)

-- TODO
scmKeywords : List String
scmKeywords = "lambda" :: Nil

scmIdent : String → String
scmIdent id = if elem id scmKeywords then "$" ++ id else pack $ fix (unpack id)
  where
    fix : List Char -> List Char
    fix Nil = Nil
    fix (x :: xs) =
      if isAlphaNum x || x == '_' || x == '.' then
        x :: fix xs
      -- make qualified names more readable
      else if x == ',' then '_' :: fix xs
      else if x == ' ' then '_' :: fix xs
      else if x == '$' then
        '$' :: '$' :: fix xs
      else
        '$' :: (toHex (cast x)) ++ fix xs


scmLit : Literal → String
scmLit (LString str) = quoteString str
scmLit (LInt n) = show n
scmLit (LChar c) = pack $ '#' :: '\\' :: c :: Nil -- FIXME
scmLit (LBool b) = ite b "#t" "#f"

scmName : QName → String
scmName qn = scmIdent $ show qn

cexpToScm : SCEnv → CExp → String

argsToScm : SCEnv → List CExp → List Quant → List String
argsToScm env Nil Nil = Nil
argsToScm env (e :: es) (Many :: qs) = cexpToScm env e :: argsToScm env es qs
argsToScm env (e :: es) (Zero :: qs) = argsToScm env es qs
argsToScm env _ _ = fatalError "Arg count mismatch"

withVar : SCEnv → CExp → (String → String) → String
-- don't rebind a variable
withVar env (CBnd n) f = f $ getEnv n env
withVar env t f = let nm = "wv$\{show $ length' env}"
  in "(let ((\{nm} \{cexpToScm env t})) \{f nm})"

cexpToScm env (CBnd n) = getEnv n env
cexpToScm env (CLam nm t) = case scbind nm env of
  (nm', env') => "(lambda (\{nm'}) \{cexpToScm env' t})"
cexpToScm env (CFun args body) = case bindAll args Lin env  of
    (nms,env') => "(lambda (\{joinBy " " nms}) \{cexpToScm env' body})"
  where
    bindAll : List (Quant × String) → SnocList String → SCEnv → List String × SCEnv
    bindAll Nil acc env = (acc <>> Nil, env)
    bindAll ((Many,nm) :: rest) acc env = case scbind nm env of
      (nm', env') => bindAll rest (acc :< nm') env'
    bindAll ((Zero,nm) :: rest) acc env = bindAll rest acc ("#f" :: env)
cexpToScm env (CApp t u) = "(\{cexpToScm env t} \{cexpToScm env u})"
cexpToScm env (CAppRef nm args Nil) = go (scmName nm) $ map (cexpToScm env) args
  where
    go : String → List String → String
    go acc Nil = acc
    go acc (arg :: args) = go "(\{acc} \{arg})" args
-- If we're short, we need to eta expand
cexpToScm env (CAppRef nm args quants) =
  go env "(\{scmName nm}" (map (cexpToScm env) args) quants
  where
    -- here `acc` is always missing a closing paren
    go : SCEnv → String → List String → List Quant → String
    go env acc Nil Nil = acc ++ ")"
    -- extra args are applied
    go env acc (arg :: args) Nil = go env "(\{acc}) \{arg}" args Nil
    -- missing args get eta expanded
    go env acc Nil (q :: qs) = case scbind "_" env of
      (nm, env') => let acc = "\{acc} \{nm}" in "(lambda (\{nm}) \{go env' acc Nil qs})"
    -- TODO / REVIEW Only for Many?
    go env acc (arg :: args) (Many :: qs) = go env "\{acc} \{arg}" args qs
    go env acc (arg :: args) (Zero :: qs) = go env acc args qs
    -- go env acc (arg :: args) (q :: qs) = go env acc args qs
-- so... we're not giving scrutinee a deBruijn index, but we may
-- need to let it so we can pull data off for the CConAlt
cexpToScm env (CCase sc alts) = do
    -- assign sc, might need to do more to make the name unique, we only
    -- get one of these per env at the moment
    -- Not vector-ref for CLitAlt...
    withVar env sc $ \ nm => doCase nm alts
  where
    -- add a bunch of lets then do body
    conAlt : SCEnv → String → Int → List String → List Quant → CExp → String
    conAlt env nm ix Nil Nil body = cexpToScm env body
    conAlt env nm ix (arg :: args) (Many :: qs) body =
        conAlt ("(vector-ref \{nm} \{show ix})" :: env) nm (ix + 1) args qs body
    conAlt env nm ix (arg :: args) (Zero :: qs) body = conAlt ("#f" :: env) nm ix args qs body
    conAlt env nm ix _ _ body = fatalError "arg/qs mismatch in conAlt"

    nilAlt : SCEnv → List Quant → CExp → String
    nilAlt env Nil body = cexpToScm env body
    nilAlt env (Zero :: qs) body = nilAlt ("#f" :: env) qs body
    nilAlt env (Many :: qs) body = fatalError "Non-empty field on nil constructor"

    consAlt' : SCEnv → List String → List Quant → CExp → String
    consAlt' env nms Nil body = cexpToScm env body
    consAlt' env nms (Zero :: qs) body = consAlt' ("#f" :: env) nms qs body
    consAlt' env Nil (Many :: qs) body = fatalError "Too many fields on cons constructor"
    consAlt' env (nm :: nms) (Many :: qs) body = consAlt' (nm :: env) nms qs body

    consAlt : SCEnv → String → List Quant → CExp → String
    consAlt env nm qs body = consAlt' env ("(car \{nm})" :: "(cdr \{nm})" :: Nil) qs body

    -- an alt branch in a `case` statement
    -- for the Nil/Cons case, we're not inside a `case`.
    doAlt : String → CAlt → String
    doAlt nm (CConAlt tag cname _ args qs body) = "((\{show tag}) \{conAlt env nm 1 args qs body})"
    doAlt nm (CDefAlt body) = "(else \{cexpToScm env body})"
    doAlt nm (CLitAlt lit body) = "((\{scmLit lit}) \{cexpToScm env body})"

    -- doCase decides the top level path - case/cond/if/...
    doCase : String → List CAlt → String
    doCase nm (CDefAlt body :: Nil) = cexpToScm env body
    doCase nm (cons@(CConAlt tag cname ConsCon args qs body ) :: rest) =
      let consBranch = consAlt env nm qs body in
      case rest of
        Nil => consBranch
        (CDefAlt body :: Nil) => "(if (null? \{nm}) \{cexpToScm env body} \{consBranch})"
        (CConAlt _ _ NilCon _ qs body :: _) => "(if (null? \{nm}) \{nilAlt env qs body} \{consBranch})"
        (CConAlt _ _ info _ _ _ :: _) => fatalError "\{show info} alt after cons"
        (CLitAlt _ _ :: _) => fatalError "lit alt after cons"
        _ => fatalError "too many alts after cons"
    doCase nm (cons@(CConAlt tag cname NilCon args qs body ) :: rest) =
      let nilBranch = consAlt env nm qs body in
      case rest of
        Nil => nilBranch
        (CDefAlt body :: Nil) => "(if (null? \{nm}) \{nilBranch} \{cexpToScm env body})"
        (CConAlt _ _ ConsCon _ qs body :: _) => "(if (null? \{nm}) \{nilBranch} \{consAlt env nm qs body})"
        (CConAlt _ _ info _ _ _ :: _) => fatalError "\{show info} alt after nil"
        (CLitAlt _ _ :: _) => fatalError "lit alt after nil"
        _ => fatalError "too many alts after cons"
    doCase nm (CConAlt tag cname _ args qs body :: Nil) = conAlt env nm 1 args qs body
    doCase nm (CLitAlt _ body :: Nil) = cexpToScm env body
    doCase nm (CDefAlt body :: Nil) = cexpToScm env body
    doCase nm alts@(CLitAlt _ _ :: _) = "(case \{nm} \{joinBy " " $ map (doAlt nm) alts})"
    doCase nm alts = "(case (vector-ref \{cexpToScm env sc} 0) \{joinBy " " $ map (doAlt nm) alts})"

cexpToScm env (CRef nm) = scmName nm
cexpToScm env (CMeta _) = fatalError "meta in code gen"
cexpToScm env (CLit lit) = scmLit lit
cexpToScm env (CLet nm t u) =
  let (nm',env') = scbind nm env in
  "(let ((\{nm'} \{cexpToScm env t})) \{cexpToScm env' u})"
cexpToScm env (CLetRec nm t u) =
  let (nm',env') = scbind nm env in
  "(let ((\{nm'} \{cexpToScm env' t})) \{cexpToScm env' u})"
cexpToScm env (CLetLoop _ _) = fatalError "CLetLoop in scheme codegen"
cexpToScm env (CLoop _ _) = fatalError "CLoop in scheme codegen"
cexpToScm env CErased = "#f"
cexpToScm env (CConstr tag nm args quants NilCon) = "'()"
cexpToScm env (CConstr tag nm args quants ConsCon) = case argsToScm env args quants of
  (a :: b :: Nil) => "(cons \{a} \{b})"
  _ => fatalError "Wrong number of args on a ConsCon"
cexpToScm env (CConstr tag nm args quants info) = "(vector \{show tag} \{unwords $ argsToScm env args quants})"
-- Should be handled by the caller
cexpToScm env (CRaw _ _) = fatalError "Stray CRaw"
-- TODO We still have a couple of these in CompileExp, for the magic Nat
cexpToScm env (CPrimOp op a b) = "(\{op} \{cexpToScm env a} \{cexpToScm env b})"

-- Collect the QNames used in a term
getNames : Tm -> List QName -> List QName
getNames (Ref x nm) acc = nm :: acc
getNames (Lam x str _ _ t) acc = getNames t acc
getNames (App x t u) acc = getNames u $ getNames t acc
getNames (Pi x str icit y t u) acc = getNames u $ getNames t $ QN primNS "PiType" :: acc
getNames (Let x str t u) acc = getNames u $ getNames t acc
getNames (LetRec x str _ t u) acc = getNames u $  getNames t acc
getNames (Case x t alts) acc = getNames t $ foldl getAltNames acc alts
  where
    getAltNames : List QName -> CaseAlt -> List QName
    getAltNames acc (CaseDefault t) = getNames t acc
    getAltNames acc (CaseCons name args t) = name :: getNames t acc
    getAltNames acc (CaseLit lit t) = getNames t acc
getNames _ acc = acc

-- returns a QName -> Def of in-use entries
-- This will be what we work on for optimization passes
getEntries : SortedMap QName Def → QName → M (SortedMap QName Def)
getEntries acc name = do
  top <- getTop
  case lookup name top of
    Nothing => do
      putStrLn "bad name \{show name}"
      pure acc
    Just (MkEntry _ name type def@(Fn exp) _) => case lookupMap' name acc of
      Just _ => pure acc
      Nothing => do
        top <- getTop
        exp <- zonk top 0 Nil exp
        let acc = updateMap name (Fn exp) acc
        foldlM getEntries acc $ getNames exp Nil
    Just (MkEntry _ name type def@(PrimFn _ _ used) _) =>
      let acc = updateMap name def acc in
      foldlM getEntries acc used
    Just entry => pure $ updateMap name entry.def acc

/-

## Sort names by dependencies

This code is _way too subtle. The problem goes away if I wrap `() =>` around 0-ary top level functions. But I'm
stubborn, so I'm giving it a try. Changing codegen may repeatedly break this and require switching to `() =>`.

The idea here is to get a list of names to emit in order of dependencies. But top level 0-ary functions can reference
and call names at startup. They can't reference something that hasn't been declared yet and can't call something that
hasn't been defined.

As an example, a recursive show instance:
- References the `show` function
- `show` function references the instance under a lambda (not inlining this yet)
- We hit the bare function first, it depends on the instance (because of recursion), which depends on the
  function, but loop prevention cuts.

We have main at the top, it is called so we visit it deep. We do a depth-first traversal, but will distinguish whether
we're visiting shallow or deep. We're trying to avoid hitting issues with indirect circular references.

- Anything we visit deep, we ensure is visited shallow first
- Shallow doesn't go into function bodies, but we do go into lambdas
- Anything invoked with arguments while shallow is visited deep, anything referenced or partially applied is still shallow.
- We keep track of both shallow and deep visits in our accumuulator
- Shallow represents the declaration, so we filter to those at the end

TODO this could be made faster by keeping a map of the done information
REVIEW could I avoid most of this by using `function` instead of arrow functions?

-/
-- TODO factor out to CompilerCommon
sortedNames : SortedMap QName CExp → List QName → List QName
sortedNames defs names =
  map snd $ filter (not ∘ fst) $ foldl (go Nil) Nil $ map (True,) names
  where
    getBody : CAlt → CExp
    getBody (CConAlt _ _ _ _ _ t) = t
    getBody (CLitAlt _ t) = t
    getBody (CDefAlt t) = t

    -- deep if this qn is being applied to something
    getNames : (deep : Bool) → List (Bool × QName) → CExp →  List (Bool × QName)
    -- liftIO calls a lambda statically
    getNames deep acc (CLam _ t) = getNames deep acc t
    getNames deep acc (CLetLoop _ t) = getNames deep acc t
    -- top level 0-ary function, doesn't happen
    getNames deep acc (CFun _ t) = if deep then getNames deep acc t else acc
    -- REVIEW - True or deep?
    getNames deep acc (CLoop args qs) = foldl (getNames True) acc args
    getNames deep acc (CAppRef nm args qs) =
      if length' args == length' qs
        then case args of
          Nil => (True, nm) :: acc
          ts => foldl (getNames True) ((True, nm) :: acc) ts
        else
          foldl (getNames deep) ((deep, nm) :: acc) args
   -- TODO look at which cases generate CApp
    getNames deep acc (CApp t u) = getNames True (getNames deep acc u) t
    getNames deep acc (CCase t alts) = foldl (getNames deep) acc $ t :: map getBody alts
    -- we're not calling it
    getNames deep acc (CRef qn) = (deep, qn) :: acc
    getNames deep acc (CLet _ t u) = getNames deep (getNames deep acc t) u
    getNames deep acc (CLetRec _ t u) = getNames deep (getNames deep acc t) u
    getNames deep acc (CConstr _ _ ts _ info) = foldl (getNames deep) acc ts
    -- if the CRaw is called, then the deps are called
    getNames deep acc (CRaw _ deps) = acc -- map (_,_ deep) deps ++ acc
    -- wrote these out so I get an error when I add a new constructor
    getNames deep acc (CLit _) = acc
    getNames deep acc (CMeta _) = acc
    getNames deep acc (CBnd _) = acc
    getNames deep acc CErased = acc
    getNames deep acc (CPrimOp op t u) = getNames deep (getNames deep acc t) u

    -- recurse on all dependencies, pushing onto acc
    go : List (Bool × QName) → List (Bool × QName) → (Bool × QName) → List (Bool × QName)
      -- Need to force shallow if we're doing deep and haven't done shallow.
    go loop acc this@(deep, qn) =
      -- there is a subtle issue here with an existing (False, qn) vs (True, qn)
      let acc = if deep && not (elem (False, qn) acc) && not (elem (False, qn) loop)
                  then go loop acc (False, qn)
                  else acc
      in if elem this loop
        then acc
        else if elem this acc then acc
        else
          case lookupMap' qn defs of
            Nothing => acc   -- only `bouncer`
            Just exp => this :: foldl (go $ this :: loop) acc (getNames deep Nil exp)

eraseEntries : {{Ref2 Defs St}} → M Unit
eraseEntries = do
  defs <- getRef Defs
  ignore $ traverse go $ toList defs
  where
    go : {{Ref2 Defs St}} → (QName × Def) → M Unit
    go (qn, Fn tm) = do
      tm' <- erase Nil tm Nil
      modifyRef Defs $ updateMap qn (Fn tm')
    go _ = pure MkUnit

-- given a initial function, return a dependency-ordered list of javascript source
process : List QName → M (List String)
process names = do
  top <- getTop
  entries <- foldlM getEntries emptyMap names

  -- Maybe move this dance into liftWhere
  ref <- newIORef entries
  let foo = MkRef ref -- for the autos below
  eraseEntries
  liftWhere
  entries <- readIORef ref
  -- Now working with defs
  exprs <- mapM defToCExp $ toList entries
  let cexpMap = foldMap const emptyMap exprs
  -- not needed for scheme
  -- cexpMap <- tailCallOpt cexpMap
  -- Not needed for scheme
  -- cexpMap <- liftLambda cexpMap
  let names = sortedNames cexpMap names
  pure $ mapMaybe (go cexpMap) names
  where
    go : ExpMap → QName → Maybe String
    go cexpMap name = do
      (qn, cexp) <- lookupMap name cexpMap
      case cexp of
        CRaw _ _ => Nothing
        _ => Just $ "(define \{scmName qn} \{cexpToScm Nil cexp})"

compileScheme : M (List String)
compileScheme = do
  top <- getTop
  let exports = getExports Nil $ listValues top.currentMod.modDefs
  let mainName = (QN top.currentMod.modName "main")
  let main = lookup mainName top
  let todo = case main of
    Nothing => exports
    Just _ => mainName :: exports
  defs <- process todo
  case lookup mainName top of
    Just _ => -- tack on call to main function
      let exec = "(\{show mainName} 'world)"
      in pure $ "(import (chezscheme))" :: "(include \"prim.ss\")" :: reverse (exec :: defs)
    Nothing => pure $ reverse defs
  where
    getExports : List QName → List TopEntry → List QName
    getExports acc Nil = acc
    getExports acc ((MkEntry fc name@(QN ns nm) type def eflags) :: rest) =
      let acc = if elem Export eflags then name :: acc else acc
      in getExports acc rest