||| take freshly-parsed input, scope check, type check, add to env
module Quox.Parser.FromParser

import public Quox.Parser.FromParser.Error as Quox.Parser.FromParser

import Quox.Pretty
import Quox.Parser.Syntax
import Quox.Parser.Parser
import public Quox.Parser.LoadFile
import Quox.Typechecker

import Data.List
import Data.Maybe
import Data.SnocVect
import Quox.EffExtra
import Control.Monad.ST.Extra

import System.File
import System.Path
import Data.IORef


%hide Typing.Error
%hide Lexer.Error
%hide Parser.Error

%default total


public export
NDefinition : Type
NDefinition = (Name, Definition)


public export
data StateTag = NS | SEEN

public export
FromParserPure : List (Type -> Type)
FromParserPure = [Except Error, DefsState, StateL NS Mods, NameGen]

public export
FromParserIO : List (Type -> Type)
FromParserIO = FromParserPure ++ [LoadFile]


export
fromParserPure : NameSuf -> Definitions ->
                 Eff FromParserPure a ->
                 Either Error (a, NameSuf, Definitions)
fromParserPure suf defs act = runSTErr $ do
  suf  <- liftST $ newSTRef suf
  defs <- liftST $ newSTRef defs
  res  <- runEff act $ with Union.(::)
    [handleExcept (\e => stLeft e),
     handleStateSTRef defs,
     handleStateSTRef !(liftST $ newSTRef [<]),
     handleStateSTRef suf]
  pure (res, !(liftST $ readSTRef suf), !(liftST $ readSTRef defs))


export covering
fromParserIO : (MonadRec io, HasIO io) =>
               IncludePath -> IORef SeenSet ->
               IORef NameSuf -> IORef Definitions ->
               Eff FromParserIO a -> io (Either Error a)
fromParserIO inc seen suf defs act =
  liftIO $ fromIOErr $ runEff act $ with Union.(::)
    [handleExcept (\e => ioLeft e),
     handleStateIORef defs,
     handleStateIORef !(newIORef [<]),
     handleStateIORef suf,
     handleLoadFileIOE LoadError WrapParseError seen inc]


parameters {auto _ : Functor m} (b : Var n -> m a) (f : PName -> m a)
           (xs : Context' PatVar n)
  private
  fromBaseName : PBaseName -> m a
  fromBaseName x = maybe (f $ MakePName [<] x) b $
                   Context.find (\y => y.name == Just x) xs

  private
  fromName : PName -> m a
  fromName x = if null x.mods then fromBaseName x.base else f x

export
fromPDimWith : Has (Except Error) fs =>
               Context' PatVar d -> PDim -> Eff fs (Dim d)
fromPDimWith ds (K e loc) = pure $ K e loc
fromPDimWith ds (V i loc) =
  fromBaseName (\i => pure $ B i loc)
               (const $ throw $ DimNotInScope loc i) ds i

private
avoidDim : Has (Except Error) fs =>
           Context' PatVar d -> Loc -> PName -> Eff fs Name
avoidDim ds loc x =
  fromName (const $ throw $ DimNameInTerm loc x.base) (pure . fromPName) ds x

private
resolveName : Mods -> Loc -> Name -> Maybe Universe ->
              Eff FromParserPure (Term d n)
resolveName ns loc x u =
  let here = addMods ns x in
  if isJust $ lookup here !(getAt DEFS) then
    pure $ FT here (fromMaybe 0 u) loc
  else do
    let ns :< _ = ns
      | _ => throw $ TermNotInScope loc x
    resolveName ns loc x u

export
fromPatVar : PatVar -> BindName
fromPatVar (Unused loc) = BN Unused loc
fromPatVar (PV x   loc) = BN (UN x) loc

export
fromPQty : PQty -> Qty
fromPQty (PQ q _) = q

export
fromPTagVal : PTagVal -> TagVal
fromPTagVal (PT t _) = t


private
fromV : Context' PatVar d -> Context' PatVar n ->
        PName -> Maybe Universe -> Loc -> Eff FromParserPure (Term d n)
fromV ds ns x u loc = fromName bound free ns x where
  bound : Var n -> Eff FromParserPure (Term d n)
  bound i = unless (isNothing u) (throw $ DisplacedBoundVar loc x) $> BT i loc

  free : PName -> Eff FromParserPure (Term d n)
  free x = resolveName !(getAt NS) loc !(avoidDim ds loc x) u

mutual
  export
  fromPTermWith : Context' PatVar d -> Context' PatVar n ->
                  PTerm -> Eff FromParserPure (Term d n)
  fromPTermWith ds ns t0 = case t0 of
    TYPE k loc =>
      pure $ TYPE k loc

    IOState loc =>
      pure $ IOState loc

    Pi pi x s t loc =>
      Pi (fromPQty pi)
        <$> fromPTermWith ds ns s
        <*> fromPTermTScope ds ns [< x] t
        <*> pure loc

    Lam x s loc =>
      Lam <$> fromPTermTScope ds ns [< x] s <*> pure loc

    App s t loc =>
      map E $ App
        <$> fromPTermElim ds ns s
        <*> fromPTermWith ds ns t
        <*> pure loc

    Sig x s t loc =>
      Sig <$> fromPTermWith ds ns s
          <*> fromPTermTScope ds ns [< x] t
          <*> pure loc

    Pair s t loc =>
      Pair <$> fromPTermWith ds ns s <*> fromPTermWith ds ns t <*> pure loc

    Case pi pair (r, ret) (CasePair (x, y) body _) loc =>
      map E $ CasePair (fromPQty pi)
        <$> fromPTermElim ds ns pair
        <*> fromPTermTScope ds ns [< r] ret
        <*> fromPTermTScope ds ns [< x, y] body
        <*> pure loc

    Fst pair loc =>
      map E $ Fst <$> fromPTermElim ds ns pair <*> pure loc

    Snd pair loc =>
      map E $ Snd <$> fromPTermElim ds ns pair <*> pure loc

    Case pi tag (r, ret) (CaseEnum arms _) loc =>
      map E $ CaseEnum (fromPQty pi)
        <$> fromPTermElim ds ns tag
        <*> fromPTermTScope ds ns [< r] ret
        <*> assert_total fromPTermEnumArms loc ds ns arms
        <*> pure loc

    NAT loc => pure $ NAT loc
    Nat n loc => pure $ Nat n loc
    Succ n loc => [|Succ (fromPTermWith ds ns n) (pure loc)|]

    STRING loc => pure $ STRING loc
    Str str loc => pure $ Str str loc

    Case pi nat (r, ret) (CaseNat zer (s, pi', ih, suc) _) loc =>
      map E $ CaseNat (fromPQty pi) (fromPQty pi')
        <$> fromPTermElim ds ns nat
        <*> fromPTermTScope ds ns [< r] ret
        <*> fromPTermWith ds ns zer
        <*> fromPTermTScope ds ns [< s, ih] suc
        <*> pure loc

    Enum strs loc => do
      let set = SortedSet.fromList strs
      unless (length strs == length (SortedSet.toList set)) $
        throw $ DuplicatesInEnumType loc strs
      pure $ Enum set loc

    Tag str loc => pure $ Tag str loc

    Eq (i, ty) s t loc =>
      Eq <$> fromPTermDScope ds ns [< i] ty
         <*> fromPTermWith ds ns s
         <*> fromPTermWith ds ns t
         <*> pure loc

    DLam i s loc =>
      DLam <$> fromPTermDScope ds ns [< i] s <*> pure loc

    DApp s p loc =>
      map E $ DApp
        <$> fromPTermElim ds ns s
        <*> fromPDimWith ds p
        <*> pure loc

    BOX q ty loc => BOX (fromPQty q) <$> fromPTermWith ds ns ty <*> pure loc

    Box val loc => Box <$> fromPTermWith ds ns val <*> pure loc

    Case pi box (r, ret) (CaseBox b body _) loc =>
      map E $ CaseBox (fromPQty pi)
        <$> fromPTermElim ds ns box
        <*> fromPTermTScope ds ns [< r] ret
        <*> fromPTermTScope ds ns [< b] body
        <*> pure loc

    V x u loc => fromV ds ns x u loc

    Ann s a loc =>
      map E $ Ann
        <$> fromPTermWith ds ns s
        <*> fromPTermWith ds ns a
        <*> pure loc

    Coe (i, ty) p q val loc =>
      map E $ Coe
        <$> fromPTermDScope ds ns [< i] ty
        <*> fromPDimWith ds p
        <*> fromPDimWith ds q
        <*> fromPTermWith ds ns val
        <*> pure loc

    Comp (i, ty) p q val r (j0, val0) (j1, val1) loc =>
      map E $ CompH'
        <$> fromPTermDScope ds ns [< i] ty
        <*> fromPDimWith ds p
        <*> fromPDimWith ds q
        <*> fromPTermWith ds ns val
        <*> fromPDimWith ds r
        <*> fromPTermDScope ds ns [< j0] val0
        <*> fromPTermDScope ds ns [< j1] val1
        <*> pure loc

  private
  fromPTermEnumArms : Loc -> Context' PatVar d -> Context' PatVar n ->
                      List (PTagVal, PTerm) ->
                      Eff FromParserPure (CaseEnumArms d n)
  fromPTermEnumArms loc ds ns arms = do
    res <- SortedMap.fromList <$>
           traverse (bitraverse (pure . fromPTagVal) (fromPTermWith ds ns)) arms
    unless (length (keys res) == length arms) $
      throw $ DuplicatesInEnumCase loc (map (fromPTagVal . fst) arms)
    pure res

  private
  fromPTermElim : Context' PatVar d -> Context' PatVar n ->
                  PTerm -> Eff FromParserPure (Elim d n)
  fromPTermElim ds ns e =
    case !(fromPTermWith ds ns e) of
      E e => pure e
      t   => let ctx = MkNameContexts (map fromPatVar ds) (map fromPatVar ns) in
             throw $ AnnotationNeeded t.loc ctx t

  private
  fromPTermTScope : {s : Nat} -> Context' PatVar d -> Context' PatVar n ->
                    SnocVect s PatVar -> PTerm ->
                    Eff FromParserPure (ScopeTermN s d n)
  fromPTermTScope ds ns xs t =
    if all isUnused xs then
      SN <$> fromPTermWith ds ns t
    else
      ST (fromSnocVect $ map fromPatVar xs) <$> fromPTermWith ds (ns ++ xs) t

  private
  fromPTermDScope : {s : Nat} -> Context' PatVar d -> Context' PatVar n ->
                    SnocVect s PatVar -> PTerm ->
                    Eff FromParserPure (DScopeTermN s d n)
  fromPTermDScope ds ns xs t =
    if all isUnused xs then
      SN <$> fromPTermWith ds ns t
    else
      DST (fromSnocVect $ map fromPatVar xs) <$> fromPTermWith (ds ++ xs) ns t


export %inline
fromPTerm : PTerm -> Eff FromParserPure (Term 0 0)
fromPTerm = fromPTermWith [<] [<]


export
globalPQty : Has (Except Error) fs => (q : Qty) -> Loc -> Eff fs GQty
globalPQty pi loc = case toGlobal pi of
  Just g  => pure g
  Nothing => throw $ QtyNotGlobal loc pi

export
fromPBaseNameNS : Has (StateL NS Mods) fs => PBaseName -> Eff fs Name
fromPBaseNameNS name = pure $ addMods !(getAt NS) $ fromPBaseName name


private
liftTC : Eff TC a -> Eff FromParserPure a
liftTC tc = runEff tc $ with Union.(::)
  [handleExcept $ \e => throw $ WrapTypeError e,
   handleReaderConst !(getAt DEFS),
   \g => send g]

private
addDef : Has DefsState fs => Name -> Definition -> Eff fs NDefinition
addDef name def = do
  modifyAt DEFS $ insert name def
  pure (name, def)


export covering
fromPDef : PDefinition -> Maybe String -> Bool ->
           Eff FromParserPure NDefinition
fromPDef (MkPDef qty pname pbody defLoc) scheme isMain = do
  name <- fromPBaseNameNS pname
  when !(getsAt DEFS $ isJust . lookup name) $ do
    throw $ AlreadyExists defLoc name
  gqty <- globalPQty qty.val qty.loc
  let sqty = globalToSubj gqty
  case pbody of
    PConcrete ptype pterm => do
      type <- traverse fromPTerm ptype
      term <- fromPTerm pterm
      case type of
        Just type => do
          ignore $ liftTC $ do
            checkTypeC empty type Nothing
            checkC empty sqty term type
          addDef name $ mkDef gqty type term scheme isMain defLoc
        Nothing => do
          let E elim = term
            | _ => throw $ AnnotationNeeded term.loc empty term
          res <- liftTC $ inferC empty sqty elim
          addDef name $ mkDef gqty res.type term scheme isMain defLoc
    PPostulate ptype => do
      type <- fromPTerm ptype
      addDef name $ mkPostulate gqty type scheme isMain defLoc


public export
data HasFail = NoFail | AnyFail | FailWith String

export
hasFail : List PDeclMod -> HasFail
hasFail []               = NoFail
hasFail (PFail str :: _) = maybe AnyFail FailWith str
hasFail (_ :: rest)      = hasFail rest

export
getScheme : List PDeclMod -> Maybe String
getScheme []                        = Nothing
getScheme (PCompileScheme str :: _) = Just str
getScheme (_ :: rest)               = getScheme rest

export
isMain : List PDeclMod -> Bool
isMain []           = False
isMain (PMain :: _) = True
isMain (_ :: rest)  = isMain rest

export covering
fromPDecl : PDecl -> Eff FromParserPure (List NDefinition)

export covering
fromPDeclBody : PDeclBody -> Maybe String -> Bool -> Loc ->
                Eff FromParserPure (List NDefinition)
fromPDeclBody (PDef def) scheme isMain loc =
  singleton <$> fromPDef def scheme isMain
fromPDeclBody (PNs  ns)  scheme isMain loc = do
  when (isJust scheme) $ throw $ SchemeOnNamespace loc ns.name
  when isMain $ throw $ MainOnNamespace loc ns.name
  localAt NS (<+> ns.name) $ concat <$> traverse fromPDecl ns.decls

export covering
expectFail : PDeclBody -> Loc -> Eff FromParserPure Error
expectFail body loc =
  let res = fromParserPure !(getAt GEN) !(getAt DEFS) $
            fromPDeclBody body Nothing False loc in
  case res of
    Left  err => pure err
    Right _   => throw $ ExpectedFail body.loc


fromPDecl (MkPDecl mods decl loc) = case hasFail mods of
  NoFail       => fromPDeclBody decl (getScheme mods) (isMain mods) loc
  AnyFail      => expectFail decl loc $> []
  FailWith str => do
    err <- expectFail decl loc
    let msg = runPretty $ prettyError False err {opts = Opts 10_000} -- w/e
    if str `isInfixOf` renderInfinite msg
       then pure []
       else throw $ WrongFail str err loc


mutual
  export covering
  loadProcessFile : Loc -> String -> Eff FromParserIO (List NDefinition)
  loadProcessFile loc file =
    case !(loadFile loc file) of
      Just tl => concat <$> traverse fromPTopLevel tl
      Nothing => pure []

  ||| populates the `defs` field of the state
  export covering
  fromPTopLevel : PTopLevel -> Eff FromParserIO (List NDefinition)
  fromPTopLevel (PD    decl)     = lift $ fromPDecl decl
  fromPTopLevel (PLoad file loc) = loadProcessFile loc file