module Lib.CompileExp

import Prelude
import Lib.Common
import Lib.Types -- Name / Tm
import Lib.TopContext
import Lib.Prettier
import Lib.Util
import Lib.Ref2
import Data.SortedMap

-- REVIEW Separate pass for constructor magic?
-- ConCase SuccCon will be replaced by Default CLet,
-- but we would need to fix up zero, since we collapse extra constructors into a default case.
-- But should be ok becaon CLitAlt doesn't bind.

CExp : U

data CAlt : U where
  CConAlt : Nat → String → ConInfo → List String → CExp → CAlt
  -- REVIEW keep var name?
  CDefAlt : CExp -> CAlt
  -- literal
  CLitAlt : Literal -> CExp -> CAlt

data CExp : U where
  CBnd : Int -> CExp
  -- How is CLam different from CFun with one arg?
  CLam : Name -> CExp -> CExp
  CFun : List (Quant × Name) -> CExp -> CExp
  CAppRef : QName -> List CExp -> List Quant -> CExp
  CApp : CExp -> CExp -> CExp
  CCase : CExp -> List CAlt -> CExp
  CRef : QName -> CExp
  CMeta : Int -> CExp
  CLit : Literal -> CExp
  CLet : Name -> CExp -> CExp -> CExp
  CLetRec : Name -> CExp -> CExp -> CExp
  CErased : CExp
  -- Data / type constructor
  CConstr : Nat → Name → List CExp → List Quant → CExp
  -- Raw javascript for `pfunc`
  CRaw : String -> List QName -> CExp
  -- Need this for magic Nat
  -- TODO - use for primitives too
  CPrimOp : String → CExp → CExp -> CExp

-- I'm counting Lam in the term for arity.  This matches what I need in
-- code gen.

lamArity : Tm -> List Quant
lamArity (Lam _ _ _ quant t) = quant :: (lamArity t)
lamArity _ = Nil

-- It would be nice to be able to declare these
compilePrimOp : String → List CExp → Maybe CExp
compilePrimOp "Prelude.addString" (x :: y :: Nil) = Just (CPrimOp "+" x y)
compilePrimOp "Prelude.addInt" (x :: y :: Nil) = Just (CPrimOp "+" x y)
compilePrimOp "Prelude.mulInt" (x :: y :: Nil) = Just (CPrimOp "*" x y)
compilePrimOp "Prelude.subInt" (x :: y :: Nil) = Just (CPrimOp "-" x y)
compilePrimOp "Prelude._&&_" (x :: y :: Nil) = Just (CPrimOp "&&" x y)
compilePrimOp "Prelude._||_" (x :: y :: Nil) = Just (CPrimOp "||" x y)
-- Assumes Bool is in the right order!
compilePrimOp "Prelude.jsEq" (_ :: x :: y :: Nil) = Just (CPrimOp "==" x y)
compilePrimOp "Prelude.divInt" (x :: y :: Nil) = Just (CPrimOp "|" (CPrimOp "/" x y) (CLit $ LInt 0))
compilePrimOp _ _ = Nothing

-- This is how much we want to curry at top level
-- leading lambda Arity is used for function defs and metas
-- TODO - figure out how this will work with erasure
arityForName : {{Ref2 Defs St}} → FC -> QName -> M (List Quant)
arityForName fc nm = do
  defs <- getRef Defs
  case lookupMap' nm defs of
    Nothing => error fc "Name \{show nm} not in scope"
    (Just Axiom) => pure Nil
    (Just (PrimOp _)) => pure $ Many :: Many :: Nil
    (Just (TCon arity strs)) => pure $ replicate' (cast arity) Many
    (Just (DCon _ _ arity str)) => pure arity
    (Just (Fn t)) => pure $ lamArity t
    (Just (PrimTCon arity)) => pure $ replicate' (cast arity) Many
    (Just (PrimFn t arity used)) => pure $ replicate' arity Many
  where


any : ∀ a. (a → Bool) → List a → Bool
any f Nil = False
any f (x :: xs) = if f x then True else any f xs

-- apply an expression at an arity to a list of args
-- CAppRef will specify any missing args, for eta conversion later
-- and any extra args get individual CApp.
apply : QName -> List CExp -> List Quant -> M CExp
-- out of args, make one up (fix that last arg)
apply qn args quants = pure $ CAppRef qn args quants
  -- go (CAppRef qn args quants) args quants
  where
    go : CExp -> List CExp -> List Quant → M CExp
    go t (arg :: args) (q :: qs) = go t args qs
    go t Nil _ = pure t
    go t (arg :: args) Nil = go (CApp t arg) args Nil

lookupDef : {{Ref2 Defs St}} → FC → QName → M Def
lookupDef fc nm = do
  defs <- getRef Defs
  case lookupMap' nm defs of
    Nothing => error fc "\{show nm} not in scope"
    Just def => pure def

getBody : CAlt → CExp
getBody (CConAlt _ _ _ _ t) = t
getBody (CLitAlt _ t) = t
getBody (CDefAlt t) = t

compileTerm : {{Ref2 Defs St}} → Tm -> M CExp
compileTerm (Bnd _ k) = pure $ CBnd k
-- need to eta expand to arity
compileTerm t@(Ref fc nm@(QN _ tag)) = do
  arity <- arityForName fc nm
  defs <- getRef Defs
  case arity of
    Nil =>
      case lookupMap' nm defs : Maybe Def of
        Just (DCon ix EnumCon _ _) => pure $ CLit $ LInt $ cast ix
        Just (DCon ix FalseCon _ _) => pure $ CLit $ LBool False
        Just (DCon ix TrueCon _ _) => pure $ CLit $ LBool True
        Just (DCon _ ZeroCon _ _) => pure $ CLit $ LInt 0
        Just (DCon _ SuccCon _ _) =>
          pure $ CLam "x" $ CPrimOp "+" (CLit $ LInt 1) (CBnd 0)
        _ => pure $ CRef nm
    _ => apply nm Nil arity

compileTerm (Meta fc k) = error fc "Compiling meta \{show k}"
compileTerm (Lam _ nm _ _ t) = CLam nm <$> compileTerm t
compileTerm tm@(App _ _ _) = case funArgs tm of
  (Meta _ k, args) => do
        error (getFC tm) "Compiling an unsolved meta \{show tm}"
        -- info (getFC tm) "Compiling an unsolved meta \{show tm}"
        -- pure $ CAppRef  "Meta\{show k}" Nil 0
  (t@(Ref fc nm), args) => do
        defs <- getRef Defs
        args' <- traverse compileTerm args
        arity <- arityForName fc nm
        let (Nothing) = compilePrimOp (show nm) args'
          | Just cexp => pure cexp
        case lookupMap' nm defs : Maybe Def of
          Just (DCon _ SuccCon _ _) => applySucc args'
          _ => apply nm args' arity
  -- REVIEW maybe we want a different constructor for non-Ref applications?
  (t, args) => do
        debug $ \ _ => "apply other \{render 90 $ pprint Nil t}"
        t' <- compileTerm t
        args' <- traverse compileTerm args
        pure $ foldl CApp t' args'
  where
    applySucc : List CExp → M CExp
    applySucc Nil = pure $ CLam "x" $ CPrimOp "+" (CLit $ LInt 1) (CBnd 0)
    applySucc (t :: Nil) = pure $ CPrimOp "+" (CLit $ LInt 1) t
    applySucc _ = error emptyFC "overapplied Succ \{show tm}"
compileTerm (UU _) = pure $ CRef (QN Nil "U")
compileTerm (Pi _ nm icit rig t u) = do
  t' <- compileTerm t
  u' <- compileTerm u
  pure $ CAppRef (QN primNS "PiType") (t' :: CLam nm u' :: Nil) (Many :: Many :: Nil)
compileTerm (Case fc t alts) = do
  t' <- compileTerm t
  alts' <- for alts $ \case
    CaseDefault tm => CDefAlt <$> compileTerm tm
    -- we use the base name for the tag, some primitives assume this
    CaseCons qn@(QN ns nm) args tm => do
      defs <- getRef Defs
      def <- lookupDef emptyFC qn
      case def of
        DCon ix info _ _ => CConAlt ix nm info args <$> compileTerm tm
        _ => error fc "\{show nm} is not constructor"

    CaseLit lit tm => CLitAlt lit <$> compileTerm tm
  pure $ CCase t' $ fancyCons t' alts'
  where
    numAltP : CAlt → Bool
    numAltP (CConAlt _ _ SuccCon _ _) = True
    numAltP (CConAlt _ _ ZeroCon _ _) = True
    numAltP _ = False

    enumAlt : CAlt → CAlt
    enumAlt (CConAlt ix nm EnumCon args tm) = CLitAlt (LInt $ cast ix) tm
    enumAlt (CConAlt ix nm FalseCon args tm) = CLitAlt (LBool False) tm
    enumAlt (CConAlt ix nm TrueCon args tm) = CLitAlt (LBool True) tm
    enumAlt alt = alt

    isInfo : ConInfo → CAlt → Bool
    isInfo needle (CConAlt _ _  info _ _)  = needle == info
    isInfo _ _ = False

    isDef : CAlt → Bool
    isDef (CDefAlt _) = True
    isDef _ = False


    doNumCon : CExp → List CAlt → List CAlt
    doNumCon sc alts =
      let zeroAlt = case find (isInfo ZeroCon) alts of
            Just (CConAlt _ _ _ _ tm) => CLitAlt (LInt 0) tm :: Nil
            Just tm => fatalError "ERROR zeroAlt mismatch \{debugStr tm}"
            _ => case find isDef alts of
                  Just (CDefAlt tm) => CLitAlt (LInt 0) tm :: Nil
                  -- This happens if the zero alt is impossible
                  _ => Nil
      in
      let succAlt = case find (isInfo SuccCon) alts of
            Just (CConAlt _ _ _ _ tm) => CDefAlt (CLet "x" (CPrimOp "-" sc (CLit $ LInt 1)) tm) :: Nil
            Just tm => fatalError "ERROR succAlt mismatch \{debugStr tm}"
            _ => case find isDef alts of
              Just alt => alt :: Nil
              _ => Nil
      in zeroAlt ++ succAlt

    fancyCons : CExp → List CAlt → List CAlt
    fancyCons sc alts =
      if any numAltP alts
        then doNumCon sc alts
        else map enumAlt alts

compileTerm (Lit _ lit) = pure $ CLit lit
compileTerm (Let _ nm t u) = do
  t' <- compileTerm t
  u' <- compileTerm u
  pure $ CLet nm t' u'
compileTerm (LetRec _ nm _ t u) = do
  t' <- compileTerm t
  u' <- compileTerm u
  pure $ CLetRec nm t' u'
compileTerm (Erased _) = pure CErased

compileFun : {{Ref2 Defs St}} → Tm -> M CExp
compileFun tm = go tm Lin
  where
    go : Tm -> SnocList (Quant × String) -> M CExp
    go (Lam _ nm _ quant t) acc = go t (acc :< (quant, nm))
    go tm Lin = compileTerm tm
    go tm args = CFun (args <>> Nil) <$> compileTerm tm

compilePop : QName → M CExp
compilePop qn = do
  top <- getTop
  let (Just def) = lookup qn top | _ => error emptyFC "\{show qn} not found"
  pure $ CErased -- FIXME - not implemented

-- What are the Defs used for above? (Arity for name)
compileDCon : Nat → QName → ConInfo → List Quant → CExp
compileDCon ix (QN _ nm) EnumCon Nil = CLit $ LInt $ cast ix
compileDCon ix (QN _ nm) TrueCon Nil = CLit $ LBool True
compileDCon ix (QN _ nm) FalseCon Nil = CLit $ LBool False
compileDCon ix (QN _ nm) info Nil = CConstr ix nm Nil Nil
compileDCon ix (QN _ nm) info arity =
  -- so we're fully applying this here, but dropping the args later?
  -- The weird thing is that lambdas need the
  let args = mkArgs Z arity
      alen = length' arity
  in CFun args $ CConstr ix nm (map (\k => CBnd $ alen - k - 1) (range 0 alen)) arity
  where
    mkArgs : Nat → List Quant → List (Quant × String)
    mkArgs k (quant :: args) = (quant, "h\{show k}") :: mkArgs (S k) args
    mkArgs k Nil = Nil



-- probably want to drop the Ref2 when we can
defToCExp : {{Ref2 Defs St}} → (QName × Def) -> M (QName × CExp)
defToCExp (qn, Axiom)      = pure $ (qn, CErased)
defToCExp (qn, (PrimOp _)) = (_,_ qn) <$> compilePop qn
defToCExp (qn, DCon ix info arity _)   = pure $ (qn, compileDCon ix qn info arity)
-- We're not using these are runtime at the moment, no typecase
-- we need to sort out tag number if we do that
defToCExp (qn, TCon arity conNames) = pure $ (qn, compileDCon Z qn NormalCon (replicate' (cast arity) Many))
defToCExp (qn, PrimTCon arity) = pure $ (qn, compileDCon Z qn NormalCon (replicate' (cast arity) Many))
defToCExp (qn, PrimFn src _ deps) = pure $ (qn, CRaw src deps)
defToCExp (qn, Fn tm)          = (_,_ qn) <$> compileFun tm