||| take freshly-parsed input, translate it to core, and check it
module Quox.Parser.FromParser

import Quox.Parser.Syntax
import Quox.Parser.Parser
import Quox.Typechecker

import Data.List
import Data.SnocVect
import Quox.EffExtra

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

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

%default total

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


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

public export
IncludePath : Type
IncludePath = List String

public export
SeenFiles : Type
SeenFiles = SortedSet String


public export
data StateTag = NS | SEEN

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

public export
FromParserEff : List (Type -> Type)
FromParserEff =
  [Except Error, StateL DEFS Definitions, StateL NS Mods,
   Reader IncludePath, StateL SEEN SeenFiles, IO]

public export
FromParser : Type -> Type
FromParser = Eff FromParserEff

-- [todo] put the locs in the core ast, obv


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 _) = pure $ K e
fromPDimWith ds (V i _) =
  fromBaseName (pure . B) (const $ throw $ DimNotInScope i) ds i

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

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

export
fromPatVar : PatVar -> BaseName
fromPatVar (Unused _) = Unused
fromPatVar (PV x   _) = UN x

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

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


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

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

    Lam x s _ =>
      Lam <$> fromPTermTScope ds ns [< x] s

    App s t _ =>
      map E $ (:@) <$> fromPTermElim ds ns s <*> fromPTermWith ds ns t

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

    Pair s t _ =>
      Pair <$> fromPTermWith ds ns s <*> fromPTermWith ds ns t

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

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

    Nat _ => pure Nat
    Zero _ => pure Zero
    Succ n _ => [|Succ $ fromPTermWith ds ns n|]

    Case pi nat (r, ret) (CaseNat zer (s, pi', ih, suc) _) _ =>
      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

    Enum strs _ =>
      let set = SortedSet.fromList strs in
      if length strs == length (SortedSet.toList set) then
        pure $ Enum set
      else
        throw $ DuplicatesInEnum strs

    Tag str _ => pure $ Tag str

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

    DLam i s _ =>
      DLam <$> fromPTermDScope ds ns [< i] s

    DApp s p _ =>
      map E $ (:%) <$> fromPTermElim ds ns s <*> fromPDimWith ds p

    BOX q ty _ => BOX (fromPQty q) <$> fromPTermWith ds ns ty

    Box val _ => Box <$> fromPTermWith ds ns val

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

    V x _ =>
      fromName (pure . E . B) (resolveName !(getAt NS) <=< avoidDim ds) ns x

    Ann s a _ =>
      map E $ (:#) <$> fromPTermWith ds ns s <*> fromPTermWith ds ns a

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

    Comp (i, ty) p q val r (j0, val0) (j1, val1) _ =>
      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

  private
  fromPTermEnumArms : Context' PatVar d -> Context' PatVar n ->
                      List (PTagVal, PTerm) ->
                      Eff FromParserPure (CaseEnumArms d n)
  fromPTermEnumArms ds ns =
    map SortedMap.fromList .
    traverse (bitraverse (pure . fromPTagVal) (fromPTermWith ds ns))

  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 ctx t

  -- [todo] use SN if the var is named but still unused
  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 : (q : Qty) -> Eff [Except Error] (So $ isGlobal q)
globalPQty pi = case choose $ isGlobal pi of
  Left  y => pure y
  Right _ => throw $ QtyNotGlobal pi


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

private
injTC : TC a -> Eff FromParserPure a
injTC act = rethrow $ mapFst WrapTypeError $ runTC !(getAt DEFS) act

export covering
fromPDef : PDefinition -> Eff FromParserPure NDefinition
fromPDef (MkPDef qty pname ptype pterm _) = do
  name <- lift $ fromPBaseNameNS pname
  let qty = fromPQty qty
  qtyGlobal <- lift $ globalPQty qty
  let gqty = Element qty qtyGlobal
  let sqty = globalToSubj gqty
  type <- lift $ traverse fromPTerm ptype
  term <- lift $ fromPTerm pterm
  case type of
    Just type => do
      injTC $ checkTypeC empty type Nothing
      injTC $ ignore $ checkC empty sqty term type
      let def = mkDef gqty type term
      modifyAt DEFS $ insert name def
      pure (name, def)
    Nothing => do
      let E elim = term | t => throw $ AnnotationNeeded empty t
      res <- injTC $ inferC empty sqty elim
      let def = mkDef gqty res.type term
      modifyAt DEFS $ insert name def
      pure (name, def)

export covering
fromPDecl : PDecl -> Eff FromParserPure (List NDefinition)
fromPDecl (PDef def) = singleton <$> fromPDef def
fromPDecl (PNs  ns)  =
  localAt NS (<+> ns.name) $ concat <$> traverse fromPDecl ns.decls


public export
LoadFile : List (Type -> Type)
LoadFile = [IO, StateL SEEN SeenFiles, Reader IncludePath, Except Error]

export covering
loadFile : String -> Eff LoadFile (Maybe String)
loadFile file =
  if contains file !(getAt SEEN) then
    pure Nothing
  else do
    Just ifile <- firstExists (map (</> file) !ask)
      | Nothing => throw $ LoadError file FileNotFound
    case !(readFile ifile) of
      Right res => modifyAt SEEN (insert file) $> Just res
      Left  err => throw $ LoadError ifile err

mutual
  export covering
  loadProcessFile : String -> FromParser (List NDefinition)
  loadProcessFile file =
    case !(lift $ loadFile file) of
      Just inp => do
        tl <- either (throw . WrapParseError file) pure $ lexParseInput file inp
        concat <$> traverse fromPTopLevel tl
      Nothing => pure []

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

export
fromParserPure : Definitions ->
                 Eff FromParserPure a ->
                 Either Error (a, Definitions)
fromParserPure defs act =
  extract $
  runExcept $
  evalStateAt NS [<] $
  runStateAt DEFS defs act

export
fromParserIO : (MonadRec io, HasIO io) =>
               IncludePath -> IORef SeenFiles -> IORef Definitions ->
               FromParser a -> io (Either Error a)
fromParserIO inc seen defs act =
  runIO $
  runExcept $
  evalStateAt NS [<] $
  runStateIORefAt SEEN seen $
  runStateIORefAt DEFS defs $
  runReader inc act