module Lib.LiftWhere

import Prelude
import Lib.Common
import Lib.Types
import Lib.TopContext
import Lib.Ref2
import Data.SortedMap
import Monad.State
import Data.IORef

-- track depth and whether we need to replace Bnd with a top level call
LiftEnv : U
LiftEnv = List (Maybe (QName × Nat))

liftWhereTm : {{Ref2 Defs St}} → QName → LiftEnv → Tm → M Tm
liftWhereTm name env (Lam fc nm icit quant t) =
  Lam fc nm icit quant <$> liftWhereTm name (Nothing :: env) t
liftWhereTm name env (App fc t u) =
  App fc <$> liftWhereTm name env t <*> liftWhereTm name env u
liftWhereTm name env (Pi fc nm icit quant t u) = do
  t <- liftWhereTm name env t
  u <- liftWhereTm name (Nothing :: env) u
  pure $ Pi fc nm icit quant t u
liftWhereTm name env (Let fc nm v sc) = do
  v <- liftWhereTm name env v
  sc <- liftWhereTm name (Nothing :: env) sc
  pure $ Let fc nm v sc
liftWhereTm name env tm@(Case fc t alts) = do
  t <- liftWhereTm name env t
  alts' <- traverse liftWhereAlt alts
  pure $ Case fc t alts'
  where
    -- This is where I wish I had put indexes on things
    liftWhereAlt : CaseAlt → M CaseAlt
    liftWhereAlt (CaseDefault tm) = CaseDefault <$> liftWhereTm name env tm
    liftWhereAlt (CaseLit l tm) = CaseLit l <$> liftWhereTm name env tm
    liftWhereAlt (CaseCons qn args tm) =
      CaseCons qn args <$> liftWhereTm name (map (const Nothing) args ++ env) tm
-- This is where the magic happens
liftWhereTm name env (LetRec fc nm ty t u) = do
  let l = length env
  -- FIXME we need a namespace and a name, these collide everywhere.
  qn <- getName name nm
  let env' = (Just (qn, S l) :: env)
  -- environment should subst this function (see next case)
  t' <- liftWhereTm qn env' t
  -- TODO we could have subst in this var and dropped the extra argument
  modifyRef Defs (updateMap qn (Fn $ wrapLam (S l) t'))

  -- The rest
  u' <- liftWhereTm qn env' u
  pure $ LetRec fc nm (Erased fc) (Erased fc) u'
  where
    getName : QName → String → M QName
    getName qn@(QN ns nm) ext = do
      let qn' = QN ns (nm ++ "." ++ ext)
      top <- getRef Defs
      let (Just _) = lookupMap qn' top | _ => pure qn'
      getName qn (ext ++ "'")

    -- Hacky - CompileExp expects a pi type that matches arity
    wrapPi : Nat → Tm → Tm
    wrapPi Z tm = tm
    wrapPi (S k) tm = Pi fc "_" Explicit Many (Erased fc) $ wrapPi k tm

    wrapLam : Nat → Tm → Tm
    wrapLam Z tm = tm
    -- REVIEW We've already erased, hopefully we don't need quantity
    wrapLam (S k) tm = Lam fc "_" Explicit Many $ wrapLam k tm

-- And where it lands
liftWhereTm name env tm@(Bnd fc k) = case getAt (cast k) env of
  Just (Just (qn, v)) => pure $ apply (length' env) (cast v) (Ref fc qn)
  _ => pure tm
  where
    apply : Int → Int → Tm → Tm
    apply l 0 tm = tm
    -- (l - k) is like lvl2ix, but witih k one bigger
    apply l k tm = App fc (apply l (k - 1) tm) (Bnd fc (l - k))
liftWhereTm name env tm = pure tm

liftWhereFn : {{Ref2 Defs St}} → QName × Def → M Unit
liftWhereFn (name, Fn tm) = do
  tm' <- liftWhereTm name Nil tm
  modifyRef Defs $ updateMap name (Fn tm')
  -- updateDef name fc type (Fn tm')
liftWhereFn _ = pure MkUnit

liftWhere : {{Ref2 Defs St}} → M Unit
liftWhere = do
  defs <- getRef Defs
  ignore $ traverse liftWhereFn $ toList defs