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


%default total

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


public export
0 Def : Type
Def = Definition Three

public export
0 NDef : Type
NDef = (Name, Def)

public export
0 Defs : Type
Defs = Definitions Three

public export
data FromParserError =
    AnnotationNeeded PTerm
  | DuplicatesInEnum (List TagVal)
  | DimNotInScope PBaseName
  | QtyNotGlobal PQty
  | DimNameInTerm PBaseName
  | TypeError (Typing.Error Three)
  | LoadError String FileError
  | ParseError Parser.Error

public export
0 IncludePath : Type
IncludePath = List String

public export
0 SeenFiles : Type
SeenFiles = SortedSet String


public export
data StateTag = DEFS | NS | SEEN

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

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


parameters {auto _ : Functor m} (b : Var n -> m a) (f : PName -> m a)
           (xs : Context' BName n)
  private
  fromBaseName : PBaseName -> m a
  fromBaseName x = maybe (f $ MakePName [<] x) b $ Context.find (== 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' BName 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' BName d -> PName -> Eff fs (Term q d n)
avoidDim ds x =
  fromName (const $ throw $ DimNameInTerm x.base) (pure . FT . fromPName) ds x


mutual
  export
  fromPTermWith : Has (Except Error) fs =>
                  Context' BName d -> Context' BName n ->
                  PTerm -> Eff fs (Term Three d n)
  fromPTermWith ds ns t0 = case t0 of
    TYPE k =>
      pure $ TYPE $ k

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

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

    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 $ Base.CasePair 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 $ Base.CaseEnum 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)) =>
      Prelude.map E $ Base.CaseNat pi 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

    BOX q ty => BOX q <$> fromPTermWith ds ns ty

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

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

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

    V x =>
      fromName (pure . E . B) (avoidDim ds) ns x

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

  private
  fromPTermEnumArms : Has (Except Error) fs =>
                      Context' BName d -> Context' BName n ->
                      List (TagVal, PTerm) -> Eff fs (CaseEnumArms Three d n)
  fromPTermEnumArms ds ns =
    map SortedMap.fromList . traverse (traverse $ fromPTermWith ds ns)

  private
  fromPTermElim : Has (Except Error) fs =>
                  Context' BName d -> Context' BName n ->
                  PTerm -> Eff fs (Elim Three d n)
  fromPTermElim ds ns e =
    case !(fromPTermWith ds ns e) of
      E e => pure e
      _   => throw $ AnnotationNeeded e

  -- [todo] use SN if the var is named but still unused
  private
  fromPTermTScope : {s : Nat} -> Has (Except Error) fs =>
                    Context' BName d -> Context' BName n ->
                    SnocVect s BName ->
                    PTerm -> Eff fs (ScopeTermN s Three d n)
  fromPTermTScope ds ns xs t =
    if all isNothing xs then
      SN <$> fromPTermWith ds ns t
    else
      SY (fromSnocVect $ map (maybe Unused UN) xs)
        <$> fromPTermWith ds (ns ++ xs) t

  private
  fromPTermDScope : {s : Nat} -> Has (Except Error) fs =>
                    Context' BName d -> Context' BName n ->
                    SnocVect s BName ->
                    PTerm -> Eff fs (DScopeTermN s Three d n)
  fromPTermDScope ds ns xs t =
    if all isNothing xs then
      SN <$> fromPTermWith ds ns t
    else
      SY (fromSnocVect $ map (maybe Unused UN) xs)
        <$> fromPTermWith (ds ++ xs) ns t


export %inline
fromPTerm : Has (Except Error) fs => PTerm -> Eff fs (Term Three 0 0)
fromPTerm = fromPTermWith [<] [<]


export
globalPQty : Has (Except Error) fs =>
             (q : PQty) -> Eff fs (IsGlobal q)
globalPQty pi = case isGlobal pi of
  Yes y => pure y
  No  n => throw $ QtyNotGlobal pi


export
fromPNameNS : Has (StateL NS Mods) fs => PName -> Eff fs Name
fromPNameNS name = pure $ addMods !(getAt NS) $ fromPName name

private
injTC : (Has (StateL DEFS Defs) fs, Has (Except Error) fs) =>
        TC Three a -> Eff fs a
injTC act = rethrow $ mapFst TypeError $ runTC !(getAt DEFS) act

export covering
fromPDef : (Has (StateL DEFS Defs) fs,
            Has (StateL NS Mods) fs,
            Has (Except Error) fs) =>
           PDefinition -> Eff fs NDef
fromPDef (MkPDef qty pname ptype pterm) = do
  name <- fromPNameNS pname
  qtyGlobal <- globalPQty qty
  let sqty = globalToSubj $ Element qty qtyGlobal
  type <- traverse fromPTerm ptype
  term <- fromPTerm pterm
  case type of
    Just type => do
      injTC $ checkTypeC empty type Nothing
      injTC $ ignore $ checkC empty sqty term type
      let def = mkDef qty type term
      modifyAt DEFS $ insert name def
      pure (name, def)
    Nothing => do
      let E elim = term | _ => throw $ AnnotationNeeded pterm
      res <- injTC $ inferC empty sqty elim
      let def = mkDef qty res.type term
      modifyAt DEFS $ insert name def
      pure (name, def)

export covering
fromPDecl : (Has (StateL DEFS Defs) fs,
             Has (StateL NS Mods) fs,
             Has (Except Error) fs) =>
            PDecl -> Eff fs (List NDef)
fromPDecl (PDef def) = singleton <$> fromPDef def
fromPDecl (PNs  ns)  =
  localAt NS (<+> ns.name) $ concat <$> traverse fromPDecl ns.decls


export covering
loadFile : (Has IO fs,
            Has (StateL SEEN SeenFiles) fs,
            Has (Reader IncludePath) fs,
            Has (Except Error) fs) =>
           String -> Eff fs (Maybe String)
loadFile file =
  if contains file !(getAt SEEN) then
    pure Nothing
  else do
    Just file <- firstExists (map (</> file) !ask)
      | Nothing => throw $ LoadError file FileNotFound
    case !(readFile file) of
      Right res => modifyAt SEEN (insert file) $> Just res
      Left  err => throw $ LoadError file err

parameters {auto _ : (Has IO fs,
                      Has (StateL SEEN SeenFiles) fs,
                      Has (Reader IncludePath) fs,
                      Has (Except Error) fs,
                      Has (StateL DEFS Defs) fs,
                      Has (StateL NS Mods) fs)}
 mutual
  export covering
  loadProcessFile : String -> Eff fs (List NDef)
  loadProcessFile file =
    case !(loadFile file) of
      Just inp => do
        tl <- either (throw . ParseError) pure $ lexParseInput inp
        concat <$> traverse fromPTopLevel tl
      Nothing => pure []

  ||| populates the `defs` field of the state
  export covering
  fromPTopLevel : PTopLevel -> Eff fs (List NDef)
  fromPTopLevel (PD    decl) = fromPDecl decl
  fromPTopLevel (PLoad file) = loadProcessFile file


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