module Quox.Untyped.Syntax

import Quox.Var
import Quox.Context
import Quox.Name
import Quox.Pretty
import Quox.Syntax.Subst

import Data.Vect
import Data.DPair
import Data.SortedMap
import Data.SnocVect
import Derive.Prelude
%hide TT.Name

%default total
%language ElabReflection


public export
data Term : Nat -> Type where
  F : (x : Name) -> Loc -> Term n
  B : (i : Var n) -> Loc -> Term n

  Lam : (x : BindName) -> (body : Term (S n)) -> Loc -> Term n
  App : (fun, arg : Term n) -> Loc -> Term n

  Pair : (fst, snd : Term n) -> Loc -> Term n
  Fst : (pair : Term n) -> Loc -> Term n
  Snd : (pair : Term n) -> Loc -> Term n

  Tag : (tag : String) -> Loc -> Term n
  CaseEnum : (tag : Term n) -> (cases : List (String, Term n)) -> Loc -> Term n
  ||| empty match with an erased head
  Absurd : Loc -> Term n

  Zero : Loc -> Term n
  Succ : (nat : Term n) -> Loc -> Term n
  CaseNat : (nat : Term n) ->
            (zer : Term n) ->
            (x, ih : BindName) -> (suc : Term (2 + n)) ->
            Loc ->
            Term n

  Erased : Loc -> Term n
%name Term s, t, u
%runElab deriveIndexed "Term" [Eq, Ord, Show]


export
Located (Term n) where
  (F x loc).loc                      = loc
  (B i loc).loc                      = loc
  (Lam x body loc).loc               = loc
  (App fun arg loc).loc              = loc
  (Pair fst snd loc).loc             = loc
  (Fst pair loc).loc                 = loc
  (Snd pair loc).loc                 = loc
  (Tag tag loc).loc                  = loc
  (CaseEnum tag cases loc).loc       = loc
  (Absurd loc).loc                   = loc
  (Zero loc).loc                     = loc
  (Succ nat loc).loc                 = loc
  (CaseNat nat zer x ih suc loc).loc = loc
  (Erased loc).loc                   = loc


public export
data Definition = ErasedDef | KeptDef (Term 0)

public export
0 Definitions : Type
Definitions = SortedMap Name Definition


export
prettyTerm : {opts : LayoutOpts} -> BContext n ->
             Term n -> Eff Pretty (Doc opts)

export
prettyArg : {opts : LayoutOpts} -> BContext n -> Term n -> Eff Pretty (Doc opts)
prettyArg xs arg = withPrec Arg $ prettyTerm xs arg

export
prettyApp' : {opts : LayoutOpts} -> BContext n -> Doc opts ->
             Term n -> Eff Pretty (Doc opts)
prettyApp' xs fun arg =
  parensIfM App =<< do
    arg <- prettyArg xs arg
    pure $ sep [fun, arg]

export
prettyApp : {opts : LayoutOpts} -> BContext n ->
            Term n -> Term n -> Eff Pretty (Doc opts)
prettyApp xs fun arg =
  prettyApp' xs !(withPrec App $ prettyTerm xs fun) arg

public export
record PrettyCaseArm a n where
  constructor MkPrettyCaseArm
  lhs : a
  {len : Nat}
  vars : Vect len BindName
  rhs : Term (len + n)

export
prettyCase : {opts : LayoutOpts} -> BContext n ->
             (a -> Eff Pretty (Doc opts)) ->
             Term n -> List (PrettyCaseArm a n) ->
             Eff Pretty (Doc opts)
prettyCase xs f head arms =
  parensIfM Outer =<< do
    header <- hsep <$> sequence [caseD, prettyTerm xs head, ofD]
    cases  <- for arms $ \(MkPrettyCaseArm lhs ys rhs) => do
      lhs <- hsep <$> sequence [f lhs, darrowD]
      rhs <- withPrec Outer $ prettyTerm (xs <>< ys) rhs
      hangDSingle lhs rhs
    lb <- hl Delim "{"; sc <- semiD; rb <- hl Delim "}"; d <- askAt INDENT
    pure $ ifMultiline
      (hsep [header, lb, separateTight sc cases, rb])
      (vsep [hsep [header, lb], indent d $ vsep (map (<+> sc) cases), rb])

prettyTerm _  (F x _) = prettyFree x
prettyTerm xs (B i _) = prettyTBind $ xs !!! i
prettyTerm xs (Lam x body _) =
  parensIfM Outer =<< do
    header <- hsep <$> sequence [lamD, prettyTBind x, darrowD]
    body   <- withPrec Outer $ prettyTerm (xs :< x) body
    hangDSingle header body
prettyTerm xs (App fun arg _) = prettyApp xs fun arg
prettyTerm xs (Pair fst snd _) =
  parens =<< separateTight !commaD <$>
    sequence {t = List} [prettyTerm xs fst, prettyTerm xs snd]
prettyTerm xs (Fst pair _) = prettyApp' xs !fstD pair
prettyTerm xs (Snd pair _) = prettyApp' xs !sndD pair
prettyTerm xs (Tag tag _) = prettyTag tag
prettyTerm xs (CaseEnum tag cases _) =
  assert_total
  prettyCase xs prettyTag tag $
    map (\(t, rhs) => MkPrettyCaseArm t [] rhs) cases
prettyTerm xs (Absurd _) = hl Syntax "absurd"
prettyTerm xs (Zero _) = zeroD
prettyTerm xs (Succ nat _) = prettyApp' xs !succD nat
prettyTerm xs (CaseNat nat zer x ih suc _) =
  assert_total
  prettyCase xs pure nat
    [MkPrettyCaseArm !zeroD  []      zer,
     MkPrettyCaseArm !sucPat [x, ih] suc]
where
  sucPat = pure $
    hsep [!succD, !(prettyTBind x) <+> !commaD, !(prettyTBind ih)]
prettyTerm _ (Erased _) =
  hl Syntax =<< ifUnicode "⌷" "[]"

export
prettyDef : {opts : LayoutOpts} -> Name ->
            Definition -> Eff Pretty (Maybe (Doc opts))
prettyDef _    ErasedDef     = [|Nothing|]
prettyDef name (KeptDef rhs) = map Just $ do
  name <- prettyFree name
  eq   <- cstD
  rhs  <- prettyTerm [<] rhs
  hangDSingle (name <++> eq) rhs


public export
USubst : Nat -> Nat -> Type
USubst = Subst Term


public export FromVar Term where fromVarLoc = B


public export
CanSubstSelf Term where
  s // th = case s of
    F x loc =>
      F x loc
    B i loc =>
      getLoc th i loc
    Lam x body loc =>
      Lam x (assert_total $ body // push th) loc
    App fun arg loc =>
      App (fun // th) (arg // th) loc
    Pair fst snd loc =>
      Pair (fst // th) (snd // th) loc
    Fst pair loc =>
      Fst (pair // th) loc
    Snd pair loc =>
      Snd (pair // th) loc
    Tag tag loc =>
      Tag tag loc
    CaseEnum tag cases loc =>
      CaseEnum (tag // th) (map (assert_total mapSnd (// th)) cases) loc
    Absurd loc =>
      Absurd loc
    Zero loc =>
      Zero loc
    Succ nat loc =>
      Succ (nat // th) loc
    CaseNat nat zer x ih suc loc =>
      CaseNat (nat // th) (zer // th)
              x ih (assert_total $ suc // pushN 2 th) loc
    Erased loc =>
      Erased loc

public export
subN : SnocVect s (Term n) -> Term (s + n) -> Term n
subN th t = t // fromSnocVect th

public export
sub1 : Term n -> Term (S n) -> Term n
sub1 e = subN [< e]