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

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 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
  | Data FC Name Raw (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 Name Telescope (List Decl)
  | Instance FC Raw (Maybe (List Decl))
  | Record FC Name Telescope (Maybe Name) (List Decl)


instance HasFC Decl where
  getFC (TypeSig x strs tm) = 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


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

foo : List String -> String
foo ts = "(" ++ unwords ts ++ ")"

instance Show Raw


instance Show Clause where
  show (MkClause fc cons pats expr) = show (fc, cons, pats, expr)

instance Show Import where
  show (MkImport _ str) = foo ("MkImport" :: show str :: Nil)

instance Show BindInfo where
  show (BI _ nm icit quant) = foo ("BI" :: show nm :: show icit :: show quant :: Nil)

-- this is for debugging, use pretty when possible

instance Show Decl where
  show (TypeSig _ str x) = foo ("TypeSig" :: show str :: show x :: Nil)
  show (FunDef _ str clauses) = foo ("FunDef" :: show str :: show clauses :: Nil)
  show (Data _ str xs ys) = foo ("Data" :: show str :: show xs :: show ys :: Nil)
  show (DCheck _ x y) = foo ("DCheck" :: show x :: show y :: Nil)
  show (PType _ name ty) = foo ("PType" :: name :: show ty :: Nil)
  show (ShortData _ lhs sigs) = foo ("ShortData" :: show lhs :: show sigs :: Nil)
  show (PFunc _ nm used ty src) = foo ("PFunc" :: nm :: show used :: show ty :: show src :: Nil)
  show (PMixFix _ nms prec fix) = foo ("PMixFix" :: show nms :: show prec :: show fix :: Nil)
  show (Class _ nm tele decls) = foo ("Class" ::  nm :: "..." :: (show $ map show decls) :: Nil)
  show (Instance _ nm decls) = foo ("Instance" ::  show nm :: (show $ map show decls) :: Nil)
  show (Record _ nm tele nm1 decls) = foo ("Record" :: show nm :: show tele :: show nm1 :: show decls :: Nil)


instance Show Module where
  show (MkModule name imports decls) = foo ("MkModule" :: show name :: show imports :: show decls :: Nil)


instance Show RCaseAlt where
  show (MkAlt x y)= foo ("MkAlt" :: show x :: show y :: Nil)

instance Show UpdateClause where
  show (ModifyField _ nm tm) = foo ("ModifyField" :: nm :: show tm :: Nil)
  show (AssignField _ nm tm) = foo ("AssignField" :: nm :: show tm :: Nil)

instance Show Raw where
  show (RImplicit _) = "_"
  show (RImpossible _) = "()"
  show (RHole _) = "?"
  show (RUpdateRec _ clauses tm) = foo ("RUpdateRec" :: show clauses :: show tm :: Nil)
  show (RVar _ name) = foo ("RVar" :: show name :: Nil)
  show (RLit _ x) = foo ( "RLit" :: show x :: Nil)
  show (RLet _ x ty v scope) = foo ( "Let" :: show x :: " : " :: show ty :: " = " :: show v :: " in " :: show scope :: Nil)
  show (RPi _ bi y z) = foo ( "Pi" :: show bi :: show y :: show z :: Nil)
  show (RApp _ x y z) = foo ( "App" :: show x :: show y :: show z :: Nil)
  show (RLam _ bi y) = foo ( "Lam" :: show bi :: show y :: Nil)
  show (RCase _ x Nothing xs) = foo ( "Case" :: show x :: " of " :: show xs :: Nil)
  show (RCase _ x (Just ty) xs) = foo ( "Case" :: show x :: " : " :: show ty :: " of " :: show xs :: Nil)
  show (RDo _ stmts) = foo ( "DO" :: "FIXME" :: Nil)
  show (RU _) = "U"
  show (RIf _ x y z) = foo ( "If" :: show x :: show y :: show z :: Nil)
  show (RWhere _ _ _) = foo ( "Where" :: "FIXME" :: Nil)
  show (RAs _ nm x) = foo ( "RAs" :: nm :: show x :: Nil)


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 (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 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 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 _ _ _) = text "TODO pretty Instance"
  pretty (ShortData _ lhs sigs) = text "data" <+> pretty lhs <+> text "=" <+> pipeSep (map pretty sigs)

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