module Haskell.AST where

import Data.List
import Pretty


type Name = String
type TyVar = String

data AST = AST [Toplevel]
  deriving (Show, Eq)

data Toplevel = TopDef Def
              | TopDecl Decl
              | TopData Data
              | TopClass Class
              | TopInst Inst
  deriving (Show, Eq)

data Def = Def Name Expr
  deriving (Show, Eq)

data Expr = App Expr [Expr]
          | Ref Name
          | Num Integer
          | Tup [Expr]
          | Lam [Name] Expr
          | Case Expr [(Pat, Expr)]
  deriving (Show, Eq)

data Pat = PatAny
         | PatVar Name
         | PatCon Name [Pat]
         | PatTup [Pat]
  deriving (Show, Eq)

data Decl = Decl Name Type
  deriving (Show, Eq)

data Type = TyTup [Type]
          | TyInt
          | TyFun Type Type
          | TyRef Name [Type]
          | TyVar Name
          | TyVoid
  deriving (Show, Eq)

data Data = Data Name [TyVar] [(Name, [Type])]
  deriving (Show, Eq)

data Class = Class Name [TyVar] [Decl]
  deriving (Show, Eq)

data Inst = Inst Name Type [Def]
  deriving (Show, Eq)


instance Pretty AST where
    pretty (AST tops) = Node "" (map pretty tops)

instance Pretty Toplevel where
    pretty (TopDef x) = pretty x
    pretty (TopDecl x) = pretty x
    pretty (TopData x) = pretty x
    pretty (TopClass x) = pretty x
    pretty (TopInst x) = pretty x

instance Pretty Def where
    pretty (Def n e) = Node (n ++ " =") [pretty e]

instance Pretty Expr where
    pretty (App e as) = Bracket "(" ")" "" (map pretty (e:as))
    pretty (Ref n) = Leaf n
    pretty (Num n) = Leaf (show n)
    pretty (Tup es) = Bracket "(" ")" "," (map pretty es)
    pretty (Lam as e) = Bracket "(" ")" "" [Node ("\\" ++ intercalate " " as ++ " ->") [pretty e]]
    pretty (Case e arms) = Bracket "(" ")" "" [Node ("case " ++ pprintOneline e ++ " of") [Bracket "{" "}" ";" (map go arms)]]
      where go (p, e') = Node (pprintOneline p ++ " ->") [pretty e']

instance Pretty Pat where
    pretty PatAny = Leaf "_"
    pretty (PatVar n) = Leaf n
    pretty (PatCon n ps) = Bracket "(" ")" "" (Leaf n : map pretty ps)
    pretty (PatTup ps) = Bracket "(" ")" "," (map pretty ps)

instance Pretty Decl where
    pretty (Decl n t) = Node (n ++ " :: ") [pretty t]

instance Pretty Type where
    pretty (TyTup ts) = Bracket "(" ")" "," (map pretty ts)
    pretty TyInt = Leaf "Int"
    pretty (TyFun t u) = Leaf $ pprintOneline t ++ " -> " ++ pprintOneline u
    pretty (TyRef n as) = Bracket "(" ")" "" (Leaf n : map pretty as)
    pretty (TyVar n) = Leaf n
    pretty TyVoid = Leaf "#Void"

instance Pretty Data where
    pretty (Data n as cs) = Node ("data " ++ n ++ " " ++ intercalate " " as ++ " =") [Bracket "" "" "|" (map go cs)]
      where go (m, ts) = Node m (map pretty ts)

instance Pretty Class where
    pretty (Class n as ds) = Node ("class " ++ n ++ " " ++ intercalate " " as ++ " where") [Bracket "{" "}" ";" (map pretty ds)]

instance Pretty Inst where
    pretty (Inst n t ds) = Node ("instance " ++ n ++ " " ++ pprintOneline t ++ " where") [Bracket "{" "}" ";" (map pretty ds)]


class AllRefs a where
    allRefs :: a -> [Name]

instance AllRefs AST where
    allRefs (AST tops) = nub $ concatMap allRefs tops

instance AllRefs Toplevel where
    allRefs (TopDef def) = allRefs def
    allRefs (TopDecl _) = []
    allRefs (TopData _) = []
    allRefs (TopClass _) = []
    allRefs (TopInst inst) = allRefs inst

instance AllRefs Def where
    allRefs (Def _ e) = allRefs e

instance AllRefs Expr where
    allRefs (App e es) = nub $ concatMap allRefs (e : es)
    allRefs (Ref n) = [n]
    allRefs (Num _) = []
    allRefs (Tup es) = nub $ concatMap allRefs es
    allRefs (Lam ns e) = allRefs e \\ ns
    allRefs (Case e pairs) = nub $ allRefs e ++ concatMap (allRefs . snd) pairs

instance AllRefs Inst where
    allRefs (Inst _ _ ds) = nub $ concatMap allRefs ds