||| 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
0 Def : Type
Def = Definition Three

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

public export
0 Defs : Type
Defs = Definitions Three

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