module Lib.ProcessDecl

import Prelude
import Data.IORef
import Data.String

import Lib.Common
import Lib.Elab
import Lib.Parser
import Lib.Syntax
import Data.SortedMap
import Lib.TopContext
import Lib.Eval
import Lib.Prettier
import Lib.Types
import Lib.Util
import Lib.Erasure

dumpEnv : Context -> M String
dumpEnv ctx =
  unlines ∘ reverse <$> go (names ctx) 0 (reverse $ zip ctx.env ctx.types) Nil
  where
    isVar : Int -> Val -> Bool
    isVar k (VVar _ k' Lin) = k == k'
    isVar _ _ = False

    go : List String -> Int -> List (Val × String × Val) -> List String -> M (List String)
    go _ _ Nil acc = pure acc
    go names k ((v, n, ty) :: xs) acc = if isVar k v
      -- TODO - use Doc and add <+/> as appropriate to printing
      then do
        ty' <- quote ctx.lvl ty
        go names (1 + k) xs ("  \{n} : \{render 90 $ pprint names ty'}":: acc)
      else do
        v' <- quote ctx.lvl v
        ty' <- quote ctx.lvl ty
        go names (1 + k) xs ("  \{n} = \{render 90 $ pprint names v'} : \{render 90 $ pprint names ty'}":: acc)


logMetas : List MetaEntry -> M Unit
logMetas Nil = pure MkUnit
logMetas (OutOfScope :: rest) = logMetas rest
logMetas (Solved fc k soln :: rest) = logMetas rest
logMetas (Unsolved fc k ctx ty User cons :: rest) = do
      ty' <- quote ctx.lvl ty
      let names = map fst ctx.types
      env <- dumpEnv ctx
      let msg = "\{env}  -----------\n  \{render 90 $ pprint names ty'}"
      info fc "User Hole\n\{msg}"
      logMetas rest
logMetas (Unsolved fc k ctx ty kind cons :: rest) = do
      ty' <- forceMeta ty
      tm <- quote ctx.lvl ty'
      -- Now that we're collecting errors, maybe we simply check at the end
      -- TODO - log constraints?
      -- FIXME in Combinatory, the val doesn't match environment?
      let msg = "Unsolved meta \{show k} \{show kind} type \{render 90 $ pprint (names ctx) tm} \{show $ length' cons} constraints"
      msgs <- for cons $ \case
        (MkMc fc env sp val) => do
            pure "  * (m\{show k} (\{unwords $ map show $ sp <>> Nil}) =?= \{show val}"
      sols <- case kind of
        AutoSolve => do
          x <- quote ctx.lvl ty
          ty <- eval ctx.env CBN x
          debug $ \ _ => "AUTO ---> \{show ty}"
          -- we want the context here too.
          top <- get
          -- matches <- case !(contextMatches ctx ty) of
          --   Nil => findMatches ctx ty $ toList top.defs
          --   xs => pure xs
          matches <- findMatches ctx ty $ map snd $ toList top.defs
          -- TODO try putting mc into TopContext for to see if it gives better terms
          pure $ ("  \{show $ length' matches} Solutions: \{show matches}" :: Nil)
          -- pure $ "  \{show $ length' matches} Solutions:" :: map (("  " ++) ∘ render 90 ∘ pprint (names ctx) ∘ fst) matches

        _ => pure Nil
      info fc $ unlines ((msg :: Nil) ++ msgs ++ sols)
      logMetas rest


-- Used for Class and Record
getSigs : List Decl -> List (FC × String × Raw)
getSigs Nil = Nil
getSigs ((TypeSig _ Nil _) :: xs) = getSigs xs
getSigs ((TypeSig fc (nm :: nms) ty) :: xs) = (fc, nm, ty) :: getSigs xs
getSigs (_ :: xs) = getSigs xs

teleToPi : Telescope -> Raw -> Raw
teleToPi Nil end = end
teleToPi ((info, ty) :: tele) end = RPi (getFC info) info ty (teleToPi tele end)

impTele : Telescope -> Telescope
impTele tele = map foo tele
  where
    foo : BindInfo × Raw → BindInfo × Raw
    foo (BI fc nm _ _ , ty) = (BI fc nm Implicit Zero, ty)



processDecl : List String -> Decl -> M Unit

-- REVIEW I supposed I could have updated top here instead of the dance with the parser...
processDecl ns  (PMixFix _ _ _ _)  = pure MkUnit

processDecl ns  (TypeSig fc names tm) = do
  putStrLn "-----"

  top <- get
  mc <- readIORef top.metaCtx
  -- let mstart = length' mc.metas
  for names $ \nm => do
    let (Nothing) = lookupRaw nm top
      | Just entry => error fc "\{show nm} is already defined at \{show entry.fc}"
    pure MkUnit
  ty <- check (mkCtx fc) tm (VU fc)
  ty <- zonk top 0 Nil ty
  putStrLn "TypeSig \{unwords names} : \{render 90 $ pprint Nil ty}"
  ignore $ for names $ \nm => setDef (QN ns nm) fc ty Axiom

processDecl ns (PType fc nm ty) = do
  top <- get
  ty' <- check (mkCtx fc) (maybe (RU fc) id ty) (VU fc)
  setDef (QN ns nm) fc ty' PrimTCon

processDecl ns  (PFunc fc nm used ty src) = do
  top <- get
  ty <- check (mkCtx fc) ty (VU fc)
  ty' <- nf Nil ty
  putStrLn "pfunc \{nm} : \{render 90 $ pprint Nil ty'} = \{show src}"
  -- TODO wire through fc?
  for used $ \ name => case lookupRaw name top of
    Nothing => error fc "\{name} not in scope"
    _ => pure MkUnit
  setDef (QN ns nm) fc ty' (PrimFn src used)

processDecl ns (Def fc nm clauses) = do
  putStrLn "-----"
  putStrLn "Def \{show nm}"
  top <- get
  mc <- readIORef top.metaCtx
  let (Just entry) = lookupRaw nm top
    | Nothing => throwError $ E fc "No declaration for \{nm}"
  let (MkEntry fc name ty Axiom) = entry
    | _ => throwError $ E fc "\{nm} already defined at \{show entry.fc}"

  putStrLn "check \{nm} at \{render 90 $ pprint Nil ty}"
  vty <- eval Nil CBN ty

  debug $ \ _ => "\{nm} vty is \{show vty}"

  -- I can take LHS apart syntactically or elaborate it with an infer
  clauses' <- traverse makeClause clauses
  tm <- buildTree (mkCtx fc) (MkProb clauses' vty)
  -- putStrLn "Ok \{render 90 $ pprint Nil tm}"

  mc <- readIORef top.metaCtx
  solveAutos
  -- TODO - make nf that expands all metas and drop zonk
  -- Idris2 doesn't expand metas for performance - a lot of these are dropped during erasure.
  -- Day1.newt is a test case
  -- tm' <- nf Nil tm
  tm' <- zonk top 0 Nil tm
  when top.verbose $ \ _ => putStrLn "NF\n\{render 80 $ pprint Nil tm'}"
  -- TODO we want to keep both versions, but this is checking in addition to erasing
  -- currently CompileExp is also doing erasure.
  -- TODO we need erasure info on the lambdas or to fake up an appropriate environment
  -- and erase inside. Currently the checking is imprecise
  tm'' <- erase Nil tm' Nil
  when top.verbose $ \ _ => putStrLn "ERASED\n\{render 80 $ pprint Nil tm'}"
  debug $ \ _ => "Add def \{nm} \{render 90 $ pprint Nil tm'} : \{render 90 $ pprint Nil ty}"
  updateDef (QN ns nm) fc ty (Fn tm')

processDecl ns (DCheck fc tm ty) = do
  putStrLn "----- DCheck"
  top <- get

  putStrLn "INFO at \{show fc}: check \{show tm} at \{show ty}"
  ty' <- check (mkCtx fc) ty (VU fc)
  putStrLn "  got type \{render 90 $ pprint Nil ty'}"
  vty <- eval Nil CBN ty'
  res <- check (mkCtx fc) tm vty
  putStrLn "  got \{render 90 $ pprint Nil res}"
  norm <- nf Nil res
  putStrLn "  NF \{render 90 $ pprint Nil norm}"
  norm <- nfv Nil res
  putStrLn "  NFV \{render 90 $ pprint Nil norm}"

processDecl ns (Class classFC nm tele decls) = do
  -- REVIEW maybe we can leverage Record for this
  -- a couple of catches, we don't want the dotted accessors and
  -- the self argument should be an auto-implicit
  putStrLn "-----"
  putStrLn "Class \{nm}"
  let fields = getSigs decls
  -- We'll need names for the telescope
  let dcName = "Mk\{nm}"
  let tcType = teleToPi tele (RU classFC)
  let tail = foldl mkApp (RVar classFC nm) tele
  let dcType = teleToPi (impTele tele) $ foldr mkPi tail fields

  putStrLn "tcon type \{render 90 $ pretty tcType}"
  putStrLn "dcon type \{render 90 $ pretty dcType}"
  let decl = Data classFC nm tcType (TypeSig classFC (dcName :: Nil) dcType :: Nil)
  putStrLn "Decl:"
  putStrLn $ render 90 $ pretty decl
  processDecl ns decl
  ignore $ for fields $ \case
    (fc,name,ty) => do
      let funType = teleToPi (impTele tele) $ RPi fc (BI fc "_" Auto Many) tail ty
      let autoPat = foldl mkAutoApp (RVar classFC dcName) fields
      let lhs = makeLHS (RVar fc name) tele
      let lhs = RApp classFC lhs autoPat Auto
      let decl = Def fc name ((lhs, (RVar fc name)) :: Nil)

      putStrLn "\{name} : \{render 90 $ pretty funType}"
      putStrLn "\{render 90 $ pretty decl}"
      processDecl ns $ TypeSig fc (name :: Nil) funType
      processDecl ns decl
  where
    makeLHS : Raw → Telescope → Raw
    makeLHS acc ((BI fc' nm icit quant,  _) :: tele) = RApp fc' (makeLHS acc tele) (RVar fc' nm) Implicit
    makeLHS acc Nil = acc

    -- TODO probably should just do the fold ourselves then.
    mkAutoApp : Raw → FC × String × Raw → Raw
    mkAutoApp acc (fc, nm, ty) = RApp fc acc (RVar fc nm) Explicit

    mkPi : FC × String × Raw → Raw → Raw
    mkPi (fc, nm, ty) acc = RPi fc (BI fc nm Explicit Many) ty acc

    mkApp : Raw → BindInfo × Raw → Raw
    mkApp acc (BI fc nm icit _, _) = RApp fc acc (RVar fc nm) icit

-- TODO - these are big, break them out into individual functions
processDecl ns (Instance instfc ty decls) = do

  putStrLn "-----"
  putStrLn "Instance \{render 90 $ pretty ty}"
  top <- get
  let tyFC = getFC ty

  vty <- check (mkCtx instfc) ty (VU instfc)
  -- Here `tele` holds arguments to the instance
  let (codomain, tele) = splitTele vty
  -- env represents the environment of those arguments
  let env = tenv (length tele)
  debug $ \ _ => "codomain \{render 90 $ pprint Nil codomain}"
  debug $ \ _ => "tele is \{show tele}"

  -- ok so we need a name, a hack for now.
  -- Maybe we need to ask the user (e.g. `instance someName : Monad Foo where`)
  -- or use "Monad\{show $ length' defs}"
  let instname = render 90 $ pprint Nil codomain
  let sigDecl = TypeSig instfc (instname :: Nil) ty
  -- This needs to be declared before processing the defs, but the defs need to be
  -- declared before this - side effect is that a duplicate def is noted at the first
  -- member
  case lookupRaw instname top of
    Just _ => pure MkUnit -- TODO check that the types match
    Nothing => processDecl ns sigDecl

  let (Just decls) = collectDecl <$> decls
    | _ => do
      debug $ \ _ => "Forward declaration \{show sigDecl}"

  let (Ref _ tconName, args) = funArgs codomain
    | (tm, _) => error tyFC "\{render 90 $ pprint Nil codomain} doesn't appear to be a TCon application"
  let (Just (MkEntry _ name type (TCon cons))) = lookup tconName top
    | _ => error tyFC "\{show tconName} is not a type constructor"
  let (con :: Nil) = cons
    | _ => error tyFC "\{show tconName} has multiple constructors \{show cons}"
  let (Just (MkEntry _ _ dcty (DCon _ _))) = lookup con top
    | _ => error tyFC "can't find constructor \{show con}"
  vdcty@(VPi _ nm icit rig a b) <- eval Nil CBN dcty
    | x => error (getFC x) "dcty not Pi"
  debug $ \ _ => "dcty \{render 90 $ pprint Nil dcty}"
  let (_,args) = funArgs codomain

  debug $ \ _ => "traverse \{show $ map showTm args}"
  -- This is a little painful because we're reverse engineering the
  -- individual types back out from the composite type
  args' <- traverse (eval env CBN) args
  debug $ \ _ => "args' is \{show args'}"
  appty <- apply vdcty args'
  conTele <- getFields appty env Nil
  -- declare individual functions, collect their defs
  defs <- for conTele $ \case
     (MkBinder fc nm Explicit rig ty) => do
       let ty' = foldr (\ x acc => case the Binder x of (MkBinder fc nm' icit rig ty') => Pi fc nm' icit rig ty' acc) ty tele
       let nm' = "\{instname},\{nm}"
       -- we're working with a Tm, so we define directly instead of processDecl
       let (Just (Def fc name xs)) = find (\x => case the Decl x of
        (Def y name xs) => name == nm
        _ => False) decls
          | _ => error instfc "no definition for \{nm}"

       setDef (QN ns nm') fc ty' Axiom
       let decl = (Def fc nm' xs)
       putStrLn "***"
       putStrLn "«\{nm'}» : \{render 90 $ pprint Nil ty'}"
       putStrLn $ render 80 $ pretty decl
       pure $ Just decl
     _ => pure Nothing

  for (mapMaybe id defs) $ \decl => do
    -- debug because already printed above, but nice to have it near processing
    debug $ \ _ => render 80 $ pretty decl
    processDecl ns decl
  let (QN _ con') = con
  let decl = Def instfc instname ((RVar instfc instname, mkRHS instname conTele (RVar instfc con')) :: Nil)
  putStrLn "SIGDECL"
  putStrLn "\{render 90 $ pretty sigDecl}"
  putStrLn $ render 80 $ pretty decl
  processDecl ns decl
  where
    -- try to extract types of individual fields from the typeclass dcon
    -- We're assuming they don't depend on each other.
    getFields : Val -> Env -> List Binder -> M (List Binder)
    getFields tm@(VPi fc nm Explicit rig ty sc) env bnds = do
      bnd <- MkBinder fc nm Explicit rig <$> quote (length' env) ty
      appsc <- sc $$ VVar fc (length' env) Lin
      getFields appsc env (bnd :: bnds)
    getFields tm@(VPi fc nm _ rig ty sc) env bnds = do
      appsc <- sc $$ VVar fc (length' env) Lin
      getFields appsc env bnds
    getFields tm xs bnds = pure $ reverse bnds

    tenv : Nat -> Env
    tenv Z = Nil
    tenv (S k) = (VVar emptyFC (cast k) Lin :: tenv k)

    mkRHS : String -> List Binder -> Raw -> Raw
    mkRHS instName (MkBinder fc nm Explicit rig ty :: bs) tm = mkRHS instName bs (RApp fc tm (RVar fc "\{instName},\{nm}") Explicit)
    mkRHS instName (b :: bs) tm = mkRHS instName bs tm
    mkRHS instName Nil tm = tm

    apply : Val -> List Val -> M Val
    apply x Nil = pure x
    apply (VPi fc nm icit rig a b) (x :: xs) = do
      bx <- b $$ x
      apply bx xs
    apply x (y :: xs) = error instfc "expected pi type \{show x}"

processDecl ns (ShortData fc lhs sigs) = do
  (nm,args) <- getArgs lhs Nil
  let ty = foldr mkPi (RU fc) args
  cons <- traverse (mkDecl args Nil) sigs
  let dataDecl = Data fc nm ty cons
  putStrLn "SHORTDATA"
  putStrLn "\{render 90 $ pretty dataDecl}"
  processDecl ns dataDecl
  where
    mkPi : FC × Name → Raw → Raw
    mkPi (fc,n) a = RPi fc (BI fc n Explicit Zero) (RU fc) a

    getArgs : Raw -> List (FC × String) -> M (String × List (FC × String))
    getArgs (RVar fc1 nm) acc = pure (nm, acc)
    getArgs (RApp _ t (RVar fc' nm) _) acc = getArgs t ((fc', nm) :: acc)
    getArgs tm _ = error (getFC tm) "Expected contructor application, got: \{show tm}"

    mkDecl : List (FC × Name) -> List Raw -> Raw -> M Decl
    mkDecl args acc (RVar fc' name) = do
      let base = foldr (\ ty acc => RPi (getFC ty) (BI (getFC ty) "_" Explicit Many) ty acc) lhs acc
      let ty = foldr mkPi base args
      pure $ TypeSig fc' (name :: Nil) ty
      where
        mkPi : FC × String → Raw → Raw
        mkPi (fc,nm) acc = RPi fc (BI fc nm Implicit Zero) (RU fc) acc

    mkDecl args acc (RApp fc' t u icit) = mkDecl args (u :: acc) t
    mkDecl args acc tm = error (getFC tm) "Expected contructor application, got: \{show tm}"

processDecl ns (Data fc nm ty cons) = do
  putStrLn "-----"
  putStrLn "Data \{nm}"
  top <- get
  mc <- readIORef top.metaCtx
  tyty <- check (mkCtx fc) ty (VU fc)
  case lookupRaw nm top of
    Just (MkEntry _ name type Axiom) => do
      tyty' <- eval Nil CBN tyty
      type' <- eval Nil CBN type
      unifyCatch fc (mkCtx fc) tyty' type'
    Just (MkEntry _ name type _) => error fc "\{show nm} already declared"
    Nothing => setDef (QN ns nm) fc tyty Axiom
  cnames <- for cons $ \x => case x of
      (TypeSig fc names tm) => do
          debug $ \ _ => "check dcon \{show names} \{show tm}"
          dty <- check (mkCtx fc) tm (VU fc)
          debug $ \ _ => "dty \{show names} is \{render 90 $ pprint Nil dty}"

          -- We only check that codomain used the right type constructor
          -- We know it's in U because it's part of a checked Pi type
          let (codomain, tele) = splitTele dty
          -- for printing
          let tnames = reverse $ map binderName tele
          let (Ref _ hn, args) = funArgs codomain
            | (tm, _) => error (getFC tm) "expected \{nm} got \{render 90 $ pprint tnames tm}"
          when (hn /= QN ns nm) $ \ _ =>
            error (getFC codomain) "Constructor codomain is \{render 90 $ pprint tnames codomain} rather than \{nm}"

          for names $ \ nm' => do
            setDef (QN ns nm') fc dty (DCon (getArity dty) hn)
          pure $ map (QN ns) names
      decl => throwError $ E (getFC decl) "expected constructor declaration"
  putStrLn "setDef \{nm}  TCon \{show $ join cnames}"
  updateDef (QN ns nm) fc tyty (TCon (join cnames))
  where
    binderName : Binder → Name
    binderName (MkBinder _ nm _ _ _) = nm

    checkDeclType : Tm -> M Unit
    checkDeclType (UU _) = pure MkUnit
    checkDeclType (Pi _ str icit rig t u) = checkDeclType u
    checkDeclType _ = error fc "data type doesn't return U"

processDecl ns (Record recordFC nm tele cname decls) = do
  putStrLn "-----"
  putStrLn "Record"
  let fields = getSigs decls
  let dcName = fromMaybe "Mk\{show nm}" cname
  let tcType = teleToPi tele (RU recordFC)
  -- REVIEW - I probably want to stick the telescope in front of the fields
  let tail = foldl (\ acc bi => case the (BindInfo × Raw) bi of (BI fc nm icit _, _) => RApp fc acc (RVar fc nm) icit) (RVar recordFC nm) tele
  let dcType = teleToPi (impTele tele) $
    foldr (\ x acc => case the (FC × String × Raw) x of (fc, nm, ty) =>  RPi fc (BI fc nm Explicit Many) ty acc  ) tail fields

  putStrLn "tcon type \{render 90 $ pretty tcType}"
  putStrLn "dcon type \{render 90 $ pretty dcType}"
  let decl = Data recordFC nm tcType (TypeSig recordFC (dcName :: Nil) dcType :: Nil)
  putStrLn "Decl:"
  putStrLn $ render 90 $ pretty decl
  processDecl ns decl
  ignore $ for fields $ \case
    (fc,name,ty) => do
      -- TODO dependency isn't handled yet
      -- we'll need to replace stuff like `len` with `len self`.
      let funType = teleToPi (impTele tele) $ RPi fc (BI fc "_" Explicit Many) tail ty
      let autoPat = foldl (\acc x => case the (FC × String × Raw) x of (fc,nm,ty) => RApp fc acc (RVar fc nm) Explicit) (RVar recordFC dcName) fields

      -- `fieldName` - consider dropping to keep namespace clean
      -- let lhs = foldl (\acc, (BI fc' nm icit quant, _) => RApp fc' acc (RVar fc' nm) Implicit) (RVar fc name) tele
      -- let lhs = RApp recordFC lhs autoPat Explicit
      -- let decl = Def fc name [(lhs, (RVar fc name))]
      -- putStrLn "\{name} : \{render 90 $ pretty funType}"
      -- putStrLn "\{render 90 $ pretty decl}"
      -- processDecl ns $ TypeSig fc (name :: Nil) funType
      -- processDecl ns decl

      -- `.fieldName`
      let pname = "." ++ name
      let lhs = foldl (\acc x => case the (BindInfo × Raw) x of (BI fc' nm icit quant, _) => RApp fc' acc (RVar fc' nm) Implicit) (RVar fc pname) tele
      let lhs = RApp recordFC lhs autoPat Explicit
      let pdecl = Def fc pname ((lhs, (RVar fc name)) :: Nil)
      putStrLn "\{pname} : \{render 90 $ pretty funType}"
      putStrLn "\{render 90 $ pretty pdecl}"
      processDecl ns $ TypeSig fc (pname :: Nil) funType
      processDecl ns pdecl