module Lib.LiftWhere

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

-- track quantity, depth, and whether we need to replace Bnd with a top level call.
-- We _could_ keep track of the support and elide that as well.
LiftEnv : U
LiftEnv = List (Quant × Maybe (QName × List Quant))

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 ((quant, 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 ((quant, 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 ((Many, 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) =
      -- TODO - can we pull quantities off of these arguments?
      -- We should switch to SnocList first, the args are backwards here
      CaseCons qn args <$> liftWhereTm name ((map (const (Many, Nothing)) args) ++ env) tm
-- This is where the magic happens
liftWhereTm name env (LetRec fc nm ty t u) = do
  -- let l = length env
  qn <- getName name nm
  -- I'll erase this one, since we're replacing all of the usages
  let qs = Zero :: map fst env
  let env' = ((Zero, Just (qn, qs)) :: 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 qs t'))

  -- The rest
  u' <- liftWhereTm qn env' u
  pure $ LetRec fc nm (Erased fc) (Erased fc) u'
  where
    -- TODO might be nice if we could nest the names (Foo.go.go) for nested 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 ++ "'")

    wrapLam : List Quant → Tm → Tm
    wrapLam Nil tm = tm
    wrapLam (q :: qs) tm = wrapLam qs $ Lam fc "_" Explicit q 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) (length' v) (Ref fc qn)
  _ => pure tm
  where
    apply : Int → Int → Tm → Tm
    apply l 0 tm = tm
    -- (l - k) is like lvl2ix, but with 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')
liftWhereFn _ = pure MkUnit

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