module Quox.Untyped.Scheme

import Quox.Name
import Quox.Context
import Quox.Untyped.Syntax
import Quox.Pretty

import Quox.EffExtra
import Quox.CharExtra
import Data.DPair
import Data.List1
import Data.String
import Data.SortedSet
import Data.Vect
import Derive.Prelude

%default total
%language ElabReflection

%hide TT.Name



export
isSchemeInitial : Char -> Bool
isSchemeInitial c =
  let gc = genCat c in
  isLetter gc || isSymbol gc && c /= '|' ||
  gc == Number Letter ||
  gc == Number Other ||
  gc == Mark NonSpacing ||
  gc == Punctuation Dash ||
  gc == Punctuation Connector ||
  gc == Punctuation Other && c /= '\'' && c /= '\\' ||
  gc == Other PrivateUse ||
  (c `elem` unpack "!$%&*/:<=>?~_^")

export
isSchemeSubsequent : Char -> Bool
isSchemeSubsequent c =
  let gc = genCat c in
  isSchemeInitial c ||
  isNumber gc ||
  isMark gc ||
  (c `elem` unpack ".+-@")

export
isSchemeId : String -> Bool
isSchemeId str =
  str == "1+" || str == "1-" ||
  case unpack str of
    []      => False
    c :: cs => isSchemeInitial c && all isSchemeSubsequent cs

export
escId : String -> String
escId str =
  let str' = concatMap doEsc $ unpack str in
  if isSchemeId str' then str' else "|\{str}|"
where
  doEsc : Char -> String
  doEsc '\\' = "\\\\"
  doEsc '|'  = "\\|"
  doEsc '\'' = "^"
  doEsc c    = singleton c


public export
data Id = I String Nat
%runElab derive "Id" [Eq, Ord]

private
prettyId' : {opts : LayoutOpts} -> Id -> Doc opts
prettyId' (I str 0) = text $ escId str
prettyId' (I str k) = text $ escId "\{str}:\{show k}"

export
prettyId : {opts : LayoutOpts} -> Id -> Eff Pretty (Doc opts)
prettyId x = hl TVar $ prettyId' x


public export
data StateTag = AVOID | MAIN

public export
Scheme : List (Type -> Type)
Scheme = [StateL AVOID (SortedSet Id), StateL MAIN (List Id)]
  -- names to avoid, and functions with #[main] (should only be one)


public export
data Sexp =
  V Id
| L (List Sexp)
| Q Sexp
| N Nat
| S String
| Lambda (List Id) Sexp
| LambdaC (List Id) Sexp -- curried lambda
| Let Id Sexp Sexp
| Case Sexp (List1 (List Sexp, Sexp))
| Define Id Sexp
| Literal String

export
FromString Sexp where fromString s = V $ I s 0


private
makeIdBase : Mods -> String -> String
makeIdBase mods str = joinBy "." $ toList $ mods :< str

export
makeId : Name -> Id
makeId (MakeName mods (UN str))   = I (makeIdBase mods str) 0
makeId (MakeName mods (MN str k)) = I (makeIdBase mods str) 0
makeId (MakeName mods Unused)     = I (makeIdBase mods "_") 0

export
makeIdB : BindName -> Id
makeIdB (BN name _) = makeId $ MakeName [<] name

private
bump : Id -> Id
bump (I x i) = I x (S i)

export covering
getFresh : SortedSet Id -> Id -> Id
getFresh used x =
  if contains x used then getFresh used (bump x) else x

export covering
freshIn : Id -> (Id -> Eff Scheme a) -> Eff Scheme a
freshIn x k =
  let x = getFresh !(getAt AVOID) x in
  localAt AVOID (insert x) $ k x

export covering
freshInB : BindName -> (Id -> Eff Scheme a) -> Eff Scheme a
freshInB x = freshIn (makeIdB x)

export covering
freshInBT : Telescope' BindName m n ->
            (Telescope' Id m n -> Eff Scheme a) ->
            Eff Scheme a
freshInBT xs act = do
  let (xs', used') = go (map makeIdB xs) !(getAt AVOID)
  localAt_ AVOID used' $ act xs'
where
  go : forall n. Telescope' Id m n ->
       SortedSet Id -> (Telescope' Id m n, SortedSet Id)
  go [<]       used = ([<], used)
  go (xs :< x) used =
    let x          = getFresh used x
        (xs, used) = go xs (insert x used)
    in
    (xs :< x, used)

export covering
freshInBC : Context' BindName n -> (Context' Id n -> Eff Scheme a) ->
            Eff Scheme a
freshInBC = freshInBT

export covering
toScheme : Context' Id n -> Term n -> Eff Scheme Sexp
toScheme xs (F x _) = pure $ V $ makeId x

toScheme xs (B i _) = pure $ V $ xs !!! i

toScheme xs (Lam x body _) =
  let Evidence n' (ys, body) = splitLam [< x] body in
  freshInBT ys $ \ys => do
    pure $ LambdaC (toList' ys) !(toScheme (xs . ys) body)

toScheme xs (App fun arg _) = do
  let (fun, args) = splitApp fun
  fun  <- toScheme xs fun
  args <- traverse (toScheme xs) args
  arg  <- toScheme xs arg
  pure $ if null args
            then L [fun, arg]
            else L $ "%" :: fun :: toList (args :< arg)

toScheme xs (Pair fst snd _) =
  pure $ L ["cons", !(toScheme xs fst), !(toScheme xs snd)]

toScheme xs (Fst pair _) =
  pure $ L ["car", !(toScheme xs pair)]

toScheme xs (Snd pair _) =
  pure $ L ["cdr", !(toScheme xs pair)]

toScheme xs (Tag tag _) =
  pure $ Q $ fromString tag

toScheme xs (CaseEnum tag cases _) =
  Case <$> toScheme xs tag
       <*> for cases (\(t, rhs) => ([fromString t],) <$> toScheme xs rhs)

toScheme xs (Absurd _) =
  pure $ Q "absurd"

toScheme xs (Zero _) =
  pure $ N 0

toScheme xs (Succ nat _) =
  case !(toScheme xs nat) of
    N n => pure $ N $ S n
    s   => pure $ L ["+", s, N 1]

toScheme xs (CaseNat nat zer (NSRec p ih suc) _) =
  freshInBC [< p, ih] $ \[< p, ih] =>
  pure $
    L ["case-nat-rec",
       Lambda []      !(toScheme xs zer),
       Lambda [p, ih] !(toScheme (xs :< p :< ih) suc),
       !(toScheme xs nat)]

toScheme xs (Str s _) = pure $ S s

toScheme xs (CaseNat nat zer (NSNonrec p suc) _) =
  freshInB p $ \p =>
  pure $
    L ["case-nat-nonrec",
       Lambda []  !(toScheme xs zer),
       Lambda [p] !(toScheme (xs :< p) suc),
       !(toScheme xs nat)]

toScheme xs (Let x rhs body _) =
  freshInB x $ \x =>
  pure $ Let x !(toScheme xs rhs) !(toScheme (xs :< x) body)

toScheme xs (Erased _) =
  pure $ Q "erased"


export
prelude : String
prelude = """
#!r6rs

; curried lambda
(define-syntax lambda%
  (syntax-rules ()
    [(_ (x0 x1 ...) body ...)
     (lambda (x0) (lambda% (x1 ...) body ...))]
    [(_ () body ...)
     (begin body ...)]))

; curried application
(define-syntax %
  (syntax-rules ()
    [(_ e0 e1 e2 ...)
     (% (e0 e1) e2 ...)]
    [(_ e) e]))

; curried function definition
(define-syntax define%
  (syntax-rules ()
    [(_ (f x ...) body ...)
     (define f (lambda% (x ...) body ...))]
    [(_ x body ...)
     (define x body ...)]))

(define-syntax builtin-io
  (syntax-rules ()
    [(_ body ...)
     (lambda (s)
       (let [(res (begin body ...))]
         (cons res s)))]))

(define (case-nat-rec z s n)
  (let go [(acc (z)) (i 0)]
    (if (= i n) acc (go (s i acc) (+ i 1)))))

(define (case-nat-nonrec z s n)
  (if (= n 0) (z) (s (- n 1))))

(define (run-main f) (f 'io-state) (void))
;;;;;;
"""

export covering
defToScheme : Name -> Definition -> Eff Scheme (Maybe Sexp)
defToScheme x ErasedDef = pure Nothing
defToScheme x (KeptDef isMain def) = do
  let x = makeId x
  when isMain $ modifyAt MAIN (x ::)
  modifyAt AVOID $ insert x
  pure $ Just $ Define x !(toScheme [<] def)
defToScheme x (SchemeDef isMain str) = do
  let x = makeId x
  when isMain $ modifyAt MAIN (x ::)
  modifyAt AVOID $ insert x
  pure $ Just $ Define x $ Literal str

orIndent : {opts : LayoutOpts} -> Doc opts -> Doc opts -> Doc opts
orIndent a b =
  parens $ ifMultiline (a <++> b) (a `vappend` indent 2 b)

export covering
prettySexp : {opts : LayoutOpts} -> Sexp -> Eff Pretty (Doc opts)

private covering
prettyLambda : {opts : LayoutOpts} ->
               String -> List Id -> Sexp -> Eff Pretty (Doc opts)
prettyLambda lam xs e =
  pure $ orIndent
    (hsep [!(hl Syntax $ text lam), !(prettySexp $ L $ map V xs)])
    !(prettySexp e)

private covering
prettyBind : {opts : LayoutOpts} -> (Id, Sexp) -> Eff Pretty (Doc opts)
prettyBind (x, e) = parens $ sep [!(prettyId x), !(prettySexp e)]

private covering
prettyLet : {opts : LayoutOpts} ->
            SnocList (Id, Sexp) -> Sexp -> Eff Pretty (Doc opts)
prettyLet ps (Let x rhs body) = prettyLet (ps :< (x, rhs)) body
prettyLet ps e =
  pure $ orIndent
    (hsep [!(hl Syntax "let*"),
           !(bracks . vsep . toList =<< traverse prettyBind ps)])
    !(prettySexp e)

private covering
prettyDefine : {opts : LayoutOpts} ->
               String -> Either Id (List Id) -> Sexp -> Eff Pretty (Doc opts)
prettyDefine def xs body =
  parens $ vappend
    (hsep [!(hl Syntax $ text def),
           !(either prettyId (prettySexp . L . map V) xs)])
    (indent 2 !(prettySexp body))

prettySexp (V x) = prettyId x
prettySexp (L []) = hl Delim "()"
prettySexp (L (x :: xs)) = do
  d  <- prettySexp x
  ds <- traverse prettySexp xs
  parens $ (hsep $ d :: ds) <|> (hsep [d, vsep ds]) <|>
           (vsep $ d :: map (indent 2) ds)
prettySexp (Q (V x)) = hl Tag $ "'" <+> prettyId' x
prettySexp (Q x) = pure $ hcat [!(hl Tag "'"), !(prettySexp x)]
prettySexp (N n) = hl Tag $ pshow n
prettySexp (S s) = prettyStrLit s
prettySexp (Lambda xs e) = prettyLambda "lambda" xs e
prettySexp (LambdaC xs e) = prettyLambda "lambda%" xs e
prettySexp (Let x rhs e) = prettyLet [< (x, rhs)] e
prettySexp (Case h as) = do
  header' <- prettySexp h
  case_   <- caseD
  let header = ifMultiline (case_ <++> header')
                           (case_ `vappend` indent 2 header')
  arms <- traverse prettyCase $ toList as
  pure $ ifMultiline
    (parens $ header <++> hsep arms)
    (parens $ vsep $ header :: map (indent 2) arms)
where
  prettyCase : (List Sexp, Sexp) -> Eff Pretty (Doc opts)
  prettyCase (ps, e) = bracks $
    ifMultiline
      (hsep [!(parens . hsep =<< traverse prettySexp ps), !(prettySexp e)])
      (vsep [!(parens .  sep =<< traverse prettySexp ps), !(prettySexp e)])
prettySexp (Define x e) = case e of
  LambdaC xs e => prettyDefine "define%" (Right $ x :: xs) e
  Lambda  xs e => prettyDefine "define"  (Right $ x :: xs) e
  _            => prettyDefine "define"  (Left x)          e
prettySexp (Literal sexp) =
  pure $ text sexp

export covering
makeRunMain : {opts : LayoutOpts} -> Id -> Eff Pretty (Doc opts)
makeRunMain x = prettySexp $ L ["run-main", V x]