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

public export
data CaseNatSuc : Nat -> Type

data Term 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 : List1 (String, Term n)) -> Loc -> Term n
  ||| empty match
  Absurd : Loc -> Term n

  Nat : (val : Nat) -> Loc -> Term n
  Succ : (nat : Term n) -> Loc -> Term n
  CaseNat : (nat : Term n) ->
            (zer : Term n) ->
            (suc : CaseNatSuc n) ->
            Loc ->
            Term n

  Str : (str : String) -> Loc -> Term n

  Let : (x : BindName) -> (rhs : Term n) -> (body : Term (S n)) -> Loc ->
        Term n

  Erased : Loc -> Term n
%name Term s, t, u

data CaseNatSuc where
  NSRec    : (x, ih : BindName) -> Term (2 + n) -> CaseNatSuc n
  NSNonrec : (x     : BindName) -> Term (S n)   -> CaseNatSuc n
%name CaseNatSuc suc

%runElab deriveParam $
  map (\ty => PI ty allIndices [Eq, Ord, Show]) ["Term", "CaseNatSuc"]


export
Located (Term n) where
  (F _ loc).loc           = loc
  (B _ loc).loc           = loc
  (Lam _ _ loc).loc       = loc
  (App _ _ loc).loc       = loc
  (Pair _ _ loc).loc      = loc
  (Fst _ loc).loc         = loc
  (Snd _ loc).loc         = loc
  (Tag _ loc).loc         = loc
  (CaseEnum _ _ loc).loc  = loc
  (Absurd loc).loc        = loc
  (Nat _ loc).loc         = loc
  (Succ _ loc).loc        = loc
  (CaseNat _ _ _ loc).loc = loc
  (Str _ loc).loc         = loc
  (Let _ _ _ loc).loc     = loc
  (Erased loc).loc        = loc


public export
data Definition =
  ErasedDef
| KeptDef   Bool (Term 0)
| SchemeDef Bool String
  -- bools are presence of #[main] flag

public export
0 Definitions : Type
Definitions = SortedMap Name Definition


export
letD, inD : {opts : LayoutOpts} -> Eff Pretty (Doc opts)
letD = hl Syntax "let"
inD  = hl Syntax "in"

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

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

export covering
prettyAppHead : {opts : LayoutOpts} -> BContext n ->
                Term n -> Eff Pretty (Doc opts)
prettyAppHead xs fun = withPrec App $ prettyTerm xs fun

export
prettyApp' : {opts : LayoutOpts} ->
             Doc opts -> SnocList (Doc opts) -> Eff Pretty (Doc opts)
prettyApp' fun args = do
  d <- askAt INDENT
  let args = toList args
  parensIfM App $
        hsep (fun :: args)
    <|> hsep [fun, vsep args]
    <|> vsep (fun :: map (indent d) args)

export covering
prettyApp : {opts : LayoutOpts} -> BContext n ->
            Doc opts -> SnocList (Term n) -> Eff Pretty (Doc opts)
prettyApp xs fun args = prettyApp' fun =<< traverse (prettyArg xs) args

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

export covering
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])

private
sucPat : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
sucPat x = pure $ !succD <++> !(prettyTBind x)

export
splitApp : Term n -> (Term n, SnocList (Term n))
splitApp (App f x _) = mapSnd (:< x) $ splitApp f
splitApp f           = (f, [<])

export
splitPair : Term n -> List (Term n)
splitPair (Pair s t _) = s :: splitPair t
splitPair t            = [t]

export
splitLam : Telescope' BindName a b -> Term b ->
           Exists $ \c => (Telescope' BindName a c, Term c)
splitLam ys (Lam x body _) = splitLam (ys :< x) body
splitLam ys t              = Evidence _ (ys, t)

export
splitLet : Telescope (\i => (BindName, Term i)) a b -> Term b ->
           Exists $ \c => (Telescope (\i => (BindName, Term i)) a c, Term c)
splitLet ys (Let x rhs body _) = splitLet (ys :< (x, rhs)) body
splitLet ys t                  = Evidence _ (ys, t)

private covering
prettyLets : {opts : LayoutOpts} ->
             BContext a -> Telescope (\i => (BindName, Term i)) a b ->
             Eff Pretty (SnocList (Doc opts))
prettyLets xs lets = sequence $ snd $ go lets where
  go : forall b. Telescope (\i => (BindName, Term i)) a b ->
       (BContext b, SnocList (Eff Pretty (Doc opts)))
  go [<]                = (xs, [<])
  go (lets :< (x, rhs)) =
    let (ys, docs) = go lets
        doc = do
          x   <- prettyTBind x
          rhs <- withPrec Outer $ prettyTerm ys rhs
          hangDSingle (hsep [!letD, x, !cstD]) (hsep [rhs, !inD]) in
    (ys :< x, docs :< doc)

private
sucCaseArm : {opts : LayoutOpts} ->
             CaseNatSuc n -> Eff Pretty (PrettyCaseArm (Doc opts) n)
sucCaseArm (NSRec x ih s) = pure $
  MkPrettyCaseArm (!(sucPat x) <+> !commaD <++> !(prettyTBind ih)) [x, ih] s
sucCaseArm (NSNonrec x s) = pure $
  MkPrettyCaseArm !(sucPat x) [x] s

prettyTerm _  (F x _) = prettyFree x
prettyTerm xs (B i _) = prettyTBind $ xs !!! i
prettyTerm xs (Lam x body _) =
  parensIfM Outer =<< do
    let Evidence n' (ys, body) = splitLam [< x] body
    vars <- hsep . toList' <$> traverse prettyTBind ys
    body <- withPrec Outer $ prettyTerm (xs . ys) body
    hangDSingle (hsep [!lamD, vars, !darrowD]) body
prettyTerm xs (App fun arg _) =
  let (fun, args) = splitApp fun in
  prettyApp xs !(prettyAppHead xs fun) (args :< arg)
prettyTerm xs (Pair fst snd _) =
  parens . separateTight !commaD =<<
    traverse (withPrec Outer . prettyTerm xs) (fst :: splitPair 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 _) =
  prettyCase xs prettyTag tag $
    map (\(t, rhs) => MkPrettyCaseArm t [] rhs) $ toList cases
prettyTerm xs (Absurd _) = hl Syntax "absurd"
prettyTerm xs (Nat n _) = hl Constant $ pshow n
prettyTerm xs (Succ nat _) = prettyApp xs !succD [< nat]
prettyTerm xs (CaseNat nat zer suc _) =
  prettyCase xs pure nat [MkPrettyCaseArm !zeroD [] zer, !(sucCaseArm suc)]
prettyTerm xs (Str s _) =
  prettyStrLit s
prettyTerm xs (Let x rhs body _) =
  parensIfM Outer =<< do
    let Evidence n' (lets, body) = splitLet [< (x, rhs)] body
    heads <- prettyLets xs lets
    body  <- withPrec Outer $ prettyTerm (xs . map fst lets) body
    let lines = toList $ heads :< body
    pure $ ifMultiline (hsep lines) (vsep lines)
prettyTerm _ (Erased _) =
  hl Syntax =<< ifUnicode "⌷" "[]"

export covering
prettyDef : {opts : LayoutOpts} -> Name ->
            Definition -> Eff Pretty (Doc opts)
prettyDef name ErasedDef =
  pure $ hsep [!(prettyFree name), !cstD, !(prettyTerm [<] $ Erased noLoc)]
prettyDef name (KeptDef isMain rhs) = do
  name <- prettyFree name {opts}
  eq   <- cstD
  rhs  <- withPrec Outer $ prettyTerm [<] rhs
  let header = if isMain then text "#[main]" <++> name else name
  hangDSingle (header <++> eq) rhs
prettyDef name (SchemeDef isMain str) = do
  name <- prettyFree name {opts}
  eq   <- cstD
  let rhs = text $ "scheme:" ++ str
  let header = if isMain then text "#[main]" <++> name else name
  hangDSingle (header <++> 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
    Nat n loc =>
      Nat n loc
    Succ nat loc =>
      Succ (nat // th) loc
    CaseNat nat zer suc loc =>
      CaseNat (nat // th) (zer // th)
              (assert_total substSuc suc th) loc
    Str s loc =>
      Str s loc
    Let x rhs body loc =>
      Let x (rhs // th) (assert_total $ body // push th) loc
    Erased loc =>
      Erased loc
  where
    substSuc : forall from, to.
               CaseNatSuc from -> USubst from to -> CaseNatSuc to
    substSuc (NSRec    x ih t) th = NSRec    x ih $ t // pushN 2 th
    substSuc (NSNonrec x    t) th = NSNonrec x    $ t // push    th

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]