module Lib.Syntax

import Prelude
import Lib.Common
import Data.String
import Data.List
import Lib.Prettier
import Lib.Types

Raw : U

-- Maybe this moves to Elab..
record Clause where
  constructor MkClause
  clauseFC : FC
  -- I'm including the type of the left, so we can check pats and get the list of possibilities
  -- But maybe rethink what happens on the left.
  -- It's a VVar k or possibly a pattern.
  -- a pattern either is zipped out, dropped (non-match) or is assigned to rhs
  -- if we can do all three then we can have a VVar here.
  cons : List Constraint
  pats : List Pattern
  -- We'll need some context to typecheck this
  -- it has names from Pats, which will need types in the env
  expr : Maybe Raw

-- could be a pair, but I suspect stuff will be added?

data RCaseAlt = MkAlt Raw (Maybe Raw)

data UpdateClause = AssignField FC String Raw | ModifyField FC String Raw

data DoStmt : U where
  DoExpr : (fc : FC) -> Raw -> DoStmt
  DoLet : (fc : FC) -> String -> Raw -> DoStmt
  DoArrow : (fc : FC) -> Raw -> Raw -> List RCaseAlt -> DoStmt

instance HasFC DoStmt where
  getFC (DoExpr fc _) = fc
  getFC (DoArrow fc _ _ _) = fc
  getFC (DoLet fc _ _) = fc

Decl : U
data Raw : U where
  RVar : (fc : FC) -> (nm : Name) -> Raw
  RLam : (fc : FC) -> BindInfo -> (sc : Raw) -> Raw
  RApp : (fc : FC) -> (t : Raw) -> (u : Raw) -> (icit : Icit) -> Raw
  RU   : (fc : FC) -> Raw
  RImpossible   : (fc : FC) -> Raw
  RPi  : (fc : FC) -> BindInfo -> (ty : Raw) -> (sc : Raw) -> Raw
  RLet : (fc : FC) -> (nm : Name) -> (ty : Raw) -> (v : Raw) -> (sc : Raw) -> Raw
  RLit : (fc : FC) -> Literal -> Raw
  RCase : (fc : FC) -> (scrut : Raw) -> (mty : Maybe Raw) -> (alts : List RCaseAlt) -> Raw
  RImplicit : (fc : FC) -> Raw
  RHole : (fc : FC) -> Raw
  RDo : (fc : FC) -> List DoStmt -> Raw
  RIf : (fc : FC) -> Raw -> Raw -> Raw -> Raw
  RWhere : (fc : FC) -> (List Decl) -> Raw -> Raw
  RAs : (fc : FC) -> Name -> Raw -> Raw
  -- has to be applied or we have to know its type as Foo → Foo to elaborate.
  -- I can bake the arg in here, or require an app in elab.
  RUpdateRec : (fc : FC) → List UpdateClause → Maybe Raw → Raw

instance HasFC Raw where
  getFC (RVar fc nm) = fc
  getFC (RLam fc _ ty) = fc
  getFC (RApp fc t u icit) = fc
  getFC (RU fc) = fc
  getFC (RPi fc _ ty sc) = fc
  getFC (RLet fc nm ty v sc) = fc
  getFC (RLit fc y) = fc
  getFC (RCase fc scrut mty alts) = fc
  getFC (RImplicit fc) = fc
  getFC (RHole fc) = fc
  getFC (RDo fc stmts) = fc
  getFC (RIf fc _ _ _) = fc
  getFC (RWhere fc _ _) = fc
  getFC (RAs fc _ _) = fc
  getFC (RUpdateRec fc _ _) = fc
  getFC (RImpossible fc) = fc

data Import = MkImport FC (FC × Name)

Telescope : U
Telescope = List (BindInfo × Raw)

data Decl
  = TypeSig FC (List Name) Raw
  | FunDef FC Name (List (Raw × Maybe Raw))
  | DCheck FC Raw Raw
  | DDerive FC (FC × String) (FC × String)
  -- TODO maybe add Telescope (before `:`) and auto-add to constructors...
  | Data FC (FC × Name) Raw (Maybe $ List Decl)
  | ShortData FC Raw (List Raw)
  | PType FC Name (Maybe Raw)
  | PFunc FC Name (List String) Raw String
  | PMixFix FC (List Name) Int Fixity
  | Class FC (FC × Name) Telescope (List Decl)
  | Instance FC Raw (Maybe (List Decl))
  | Record FC (FC × Name) Telescope (Maybe $ FC × Name) (List Decl)
  | Exports FC (List $ FC × Name)


instance HasFC Decl where
  getFC (TypeSig x strs tm) = x
  getFC (Exports x _) = x
  getFC (FunDef x str xs) = x
  getFC (DCheck x tm tm1) = x
  getFC (Data x str tm xs) = x
  getFC (ShortData x _ _) = x
  getFC (PType x str mtm) = x
  getFC (PFunc x str _ tm str1) = x
  getFC (PMixFix x strs k y) = x
  getFC (Class x str xs ys) = x
  getFC (Instance x tm xs) = x
  getFC (Record x str tm str1 xs) = x
  getFC (DDerive x _ _) = x


record Module where
  constructor MkModule
  modname : Name
  imports : List Import
  decls : List Decl

instance Show Raw

derive Show Clause
derive Show Import
derive Show BindInfo
derive Show DoStmt
derive Show Decl
derive Show Module
derive Show RCaseAlt
derive Show UpdateClause
derive Show Raw

instance Pretty Literal where
  pretty (LString t) = text t
  pretty (LBool b) = if b then text "true" else text "false"
  pretty (LInt i) = text $ show i
  pretty (LChar c) = text $ show c

wrap : Icit -> Doc -> Doc
wrap Explicit x = text "(" ++ x ++ text ")"
wrap Implicit x = text "{" ++ x ++ text "}"
wrap Auto x = text "{{" ++ x ++ text "}}"

instance Pretty Pattern where
  pretty (PatVar _ Implicit str) = text str
  pretty (PatImpossible _) = text "()"
  pretty (PatVar _ icit str) = wrap icit $ text str
  pretty (PatCon _ icit nm args Nothing) = text (show nm) <+> spread (map pretty args)
  pretty (PatCon _ icit nm args (Just as)) = text as ++ text "@(" ++ text (show nm) <+> spread (map pretty args) ++ text ")"
  pretty (PatWild _ icit) = text "_"
  pretty (PatLit _ lit) = pretty lit

instance Pretty Raw where
  pretty = asDoc 0
    where
    bindDoc : BindInfo -> Doc
    bindDoc (BI _ nm icit quant) = text "BINDDOC"

    par : Int -> Int -> Doc -> Doc
    par p p' d = if p' < p then text "(" ++ d ++ text ")" else d

    asDoc : Int -> Raw -> Doc
    asDoc p (RVar _ str) = text str
    asDoc p (RLam _ (BI _ nm icit q) x) = par p 0 $ text "\\" ++ wrap icit (text nm) <+> text "=>" <+> asDoc 0 x
    -- This needs precedence and operators...
    asDoc p (RApp _ x y Explicit) = par p 2 $ asDoc 2 x <+> asDoc 3 y
    asDoc p (RApp _ x y Implicit) = par p 2 $ asDoc 2 x <+> text "{" ++ asDoc 0 y ++ text "}"
    asDoc p (RApp _ x y Auto) = par p 2 $ asDoc 2 x <+> text "{{" ++ asDoc 0 y ++ text "}}"
    asDoc p (RU _) = text "U"
    asDoc p (RPi _ (BI _ "_" Explicit Many) ty scope) = par p 1 $ asDoc p ty <+> text "->" <+/> asDoc p scope
    asDoc p (RPi _ (BI _ nm icit quant) ty scope) =
      par p 1 $ wrap icit (text (show quant ++ nm) <+> text ":" <+> asDoc p ty ) <+> text "->" <+/> asDoc 1 scope
    asDoc p (RLet _ x v ty scope) =
      par p 0 $ text "let" <+> text x <+> text ":" <+> asDoc p ty
          <+> text "=" <+> asDoc p v
          <+/> text "in" <+> asDoc p scope
    asDoc p (RLit _ lit) = pretty lit
    asDoc p (RCase _ x _ xs) = text "TODO - Pretty RCase"
    asDoc p (RImplicit _) = text "_"
    asDoc p (RImpossible _) = text "()"
    asDoc p (RHole _) = text "?"
    asDoc p (RDo _ stmts) = text "TODO - Pretty RDo"
    asDoc p (RIf _ x y z) = par p 0 $ text "if" <+> asDoc 0 x <+/> text "then" <+> asDoc 0 y <+/> text "else" <+> asDoc 0 z
    asDoc p (RWhere _ dd b) = text "TODO pretty RWhere"
    asDoc p (RAs _ nm x) = text nm ++ text "@(" ++ asDoc 0 x ++ text ")"
    asDoc p (RUpdateRec _ clauses tm) = text "{" <+> text "TODO RUpdateRec" <+> text "}"

prettyBind : (BindInfo × Raw) -> Doc
prettyBind (BI _ nm icit quant, ty) = wrap icit (text (show quant ++ nm) <+> text ":" <+> pretty ty)

pipeSep : List Doc -> Doc
pipeSep = folddoc (\a b => a <+/> text "|" <+> b)


instance Pretty Decl where
  pretty (TypeSig _ nm ty) = spread (map text nm) <+> text ":" <+> nest 2 (pretty ty)
  pretty (DDerive _ x y) = text "derive" <+> text (snd x) <+> text (snd y)
  pretty (FunDef _ nm clauses) = stack $ map prettyPair clauses
    where
      prettyPair : Raw × Maybe Raw → Doc
      prettyPair (a, Nothing) = pretty a
      prettyPair (a, Just b) = pretty a <+> text "=" <+> pretty b
  pretty (Data _ (_,nm) x Nothing) = text "data" <+> text nm <+> text ":" <+>  pretty x
  pretty (Data _ (_,nm) x (Just xs)) = text "data" <+> text nm <+> text ":" <+>  pretty x <+> (nest 2 $ text "where" </> stack (map pretty xs))
  pretty (DCheck _ x y) = text "#check" <+> pretty x <+> text ":" <+> pretty y
  pretty (PType _ nm ty) = text "ptype" <+> text nm <+> (maybe empty (\ty => text ":" <+> pretty ty) ty)
  pretty (PFunc _ nm Nil ty src) = text "pfunc" <+> text nm <+> text ":" <+> nest 2 (pretty ty <+> text ":=" <+/> text (show src))
  pretty (PFunc _ nm used ty src) = text "pfunc" <+> text nm <+> text "uses" <+> spread (map text used) <+> text ":" <+> nest 2 (pretty ty <+> text ":=" <+/> text (show src))
  pretty (PMixFix _ names prec fix) = text (show fix) <+> text (show prec) <+> spread (map text names)
  pretty (Record _ (_,nm) tele (cname) decls) = text "record" <+> text nm <+> text ":" <+> spread (map prettyBind tele)
    <+> (nest 2 $ text "where" </> stack (maybe empty (\ nm' => text "constructor" <+> text (snd nm')) cname :: map pretty decls))
  pretty (Class _ (_,nm) tele decls) = text "class" <+> text nm <+> text ":" <+> spread (map prettyBind tele)
    <+> (nest 2 $ text "where" </> stack (map pretty decls))
  pretty (Instance fc top Nothing) = text "instance" <+> pretty top
  pretty (Instance fc top (Just decls)) = text "instance" <+> pretty top <+> nest 2 (text "where" </> stack (map pretty decls))
  pretty (ShortData _ lhs sigs) = text "data" <+> pretty lhs <+> text "=" <+> pipeSep (map pretty sigs)
  pretty (Exports _ nms) = text "#export" <+> spread (map (text ∘ show ∘ snd) nms)

lhsNames : Raw → List String
lhsNames tm = case tm of
  RVar fc n => n :: Nil
  RAs _ n tm => n :: lhsNames tm
  RApp _ t u _ => lhsNames t ++ lhsNames u
  -- the rest have no names or don't occur on LHS
  _ => Nil

-- used for the projection type in dependent records, probably overkill, but maybe it will be useful elsewhere
-- TODO unit tests on substRaw
substRaw : List (String × Raw) → Raw → Raw
substRaw ss t = case t of
  RVar fc n => fromMaybe t (lookup n ss)
  (RUpdateRec fc uc target) => RUpdateRec fc (map substUC uc) (map (substRaw ss) target)
  -- LHS only
  (RAs fc nm t) => RAs fc nm (substRaw ss t)
  (RIf fc c t e) => RIf fc (substRaw ss c) (substRaw ss t) (substRaw ss e)
  (RLet fc nm ty v sc) => RLet fc nm (substRaw ss ty) (substRaw ss v) (substRaw ss sc)
  (RPi fc info a b) => RPi fc info (substRaw ss a) (substRaw (filterBind info ss) b)
  (RApp fc t u icit) => RApp fc (substRaw ss t) (substRaw ss u) icit
  (RLam fc info sc) => RLam fc info (substRaw (filterBind info ss) sc)
  -- FIXME shadowing
  (RWhere fc ds body) => RWhere fc (map substDecl ds) (substRaw ss body)
  (RDo fc stmts) => RDo fc (substStmts ss stmts)
  (RCase fc scrut mdef alts) => RCase fc (substRaw ss scrut) (map (substRaw ss) mdef) (map substAlt alts)
  -- Enumerate to force consideration of new cases
  t@(RImpossible _) => t
  t@(RU _) => t
  t@(RHole fc) => t
  t@(RImplicit fc) => t
  t@(RLit _ _) => t
    where
      -- Need to handle shadowing!
      filter : ∀ a. List String → List (String × a) → List (String × a)
      filter nms Nil = Nil
      filter nms (x@(a,b) :: xs) = if elem a nms then filter nms xs else x :: filter nms xs

      filterBind : ∀ a. BindInfo → List (String × a) → List (String × a)
      filterBind (BI fc nm _ _) xs = filter (nm :: Nil) xs

      substUC : UpdateClause → UpdateClause
      substUC (AssignField fc nm t) = AssignField fc nm (substRaw ss t)
      substUC (ModifyField fc nm t) = ModifyField fc nm (substRaw ss t)

      substClause : Raw × Maybe Raw → Raw × Maybe Raw
      substClause (a,b) = (substRaw ss a, map (substRaw ss) b)

      substDecl : Decl → Decl
      substDecl (TypeSig fc nms ty) = TypeSig fc nms (substRaw ss ty)
      substDecl (FunDef fc nm clauses) = FunDef fc nm $ map substClause clauses
      substDecl d = d -- shouldn't happen

      substAlt : RCaseAlt → RCaseAlt
      substAlt (MkAlt a b) = MkAlt (substRaw ss a) (map (substRaw (filter (lhsNames a) ss)) b)

      substStmts : List (String × Raw) → List DoStmt → List DoStmt
      substStmts ss Nil = Nil
      substStmts ss (DoExpr fc t :: rest) = DoExpr fc (substRaw ss t) :: substStmts ss rest
      substStmts ss (DoArrow fc pat sc alts :: rest) =
        DoArrow fc (substRaw ss pat) (substRaw ss sc) (map (substAlt) alts) :: substStmts (filter (lhsNames pat) ss) rest
      substStmts ss (DoLet fc nm t :: rest) = DoLet fc nm (substRaw ss t) :: substStmts (filter (nm :: Nil) ss) rest

instance Pretty Module where
  pretty (MkModule name imports decls) =
    text "module" <+> text name
      </> stack (map doImport imports)
      </> stack (map pretty decls)
      where
        doImport : Import -> Doc
        doImport (MkImport _ (_,nm)) = text "import" <+> text nm ++ line