-- Pretty printing of an ART expression with a sub-expression highlit.
-- First convert a value (of class HatRep) into an S-expression,
-- then convert the S-expression to a Doc for pretty printing.
module SExp
  ( SExp(..)
  , fileNode2SExp
  , sExp2Doc
  , arity, child, label, children, rebuild, parent
  , funId, funLabel
  ) where

import LowLevel hiding    (nil)
import qualified LowLevel (nil)
import HighlightStyle     (Highlight(..),Colour(..))
import PrettyLibHighlight as Pretty
                          (Doc,text,(<>),delimiter,fdelimiter,nil,group,parens
                          ,groupNest,pretty,highlight)
import Char               (isAlpha)
import List               (unzip3)
import Numeric            (showHex)
import IO                 (hPutStrLn,stderr)
import IOExtras           (unsafePerformIO)


bold = highlight [Bold, Foreground Blue]

errorT  :: String -> a
errorT s = unsafePerformIO (do hPutStrLn stderr s; return (error ""))

data SExp a
  = SApp a [SExp a]	-- n-ary application of at least 2 expressions
  | SId a String SFixity
  | SLiteral a String
  | SCons a (SExp a) (SExp a)	-- character strings are treated specially
  | SNil a
-- | SLambda a		-- no longer used, lambda appears as an SLiteral
  | SWithin a [SExp a]	-- chains of if/case/guard inside an expression
  | SIf a (SExp a)
  | SCase a (SExp a)
  | SGuard a (SExp a)
  | SCut a		-- subexpression that was cut off (to limit depth)
  | SUnevaluated a
  | SBottom a
  | SCycle a String (SExp a)
			-- cyclic expression to be shown as `id where id = ..'
  | SEquation a (SExp a) (SExp a)
			-- an equation only makes sense as the root of an SExp
  | SProj a (SExp a)	-- a projection

data SFixity = 
  SInfix Int | SInfixL Int | SInfixR Int | SAssoc Int String | SInfixDefault
  -- need own type for some hardcoded operators that are known to be
  -- semantically associative

-- translate fixity from the file representation to the structured fixity type
transFixity :: Int -> SFixity
transFixity f = case f `divMod` 4 of
                  (p,0) -> SInfix p
                  (p,1) -> SInfixL p
                  (p,2) -> SInfixR p
                  (p,3) -> SInfixDefault

-- arity of an S-expression
arity :: SExp a -> Int
arity (SApp _ exps)     = length exps
arity (SWithin _ exps)  = length exps
arity (SCons _ _ _)     = 2
arity (SIf _ _)         = 1
arity (SCase _ _)       = 1
arity (SGuard _ _)      = 1
arity (SEquation _ _ _) = 2
arity (SProj _ e)       = arity e
arity _                 = 0

-- get child of an S-expression
-- precondition: i < arity exp
child :: Int -> SExp a -> SExp a
child i (SApp _ exps)    = skipCaseIfGuard (exps!!i)
child i (SWithin _ exps) = skipCaseIfGuard (exps!!i)
child 0 (SCons _ hd tl)  = hd
child 1 (SCons _ hd tl)  = tl
--child 0 (SIf _ exp)      = exp
--child 0 (SCase _ exp)    = exp
--child 0 (SGuard _ exp)   = exp
--child n (SProj _ exp)    = child n exp
child n (SProj _ exp)      = exp
child 0 (SEquation _ l r) = l
child 1 (SEquation _ l r) = r
child i exp               = errorT ("SExp.child: "++show i++" is too large.")

skipCaseIfGuard (SIf _ exp) = exp
skipCaseIfGuard (SCase _ exp) = exp
skipCaseIfGuard (SGuard _ exp) = exp
skipCaseIfGuard exp = exp

label :: SExp a -> a
label (SApp l _)      = l
label (SId l _ _)     = l
label (SLiteral l _)  = l
label (SCons l _ _)   = l
label (SNil l)        = l
--label (SLambda l)     = l
label (SWithin l _)   = l
label (SIf l _)       = l
label (SCase l _)     = l
label (SGuard l _)    = l
label (SCut l)        = l
label (SUnevaluated l)= l
label (SBottom l)     = l
label (SCycle l _ _)  = l
label (SEquation l _ _) = l
--label (SProj _ e)     = label e
label (SProj l _)     = l

children :: SExp a -> [SExp a]
children (SApp _ es)      = es
children (SWithin l es)   = es
children (SCons l hd tl)  = [hd,tl]
children (SIf _ e)        = [e]
children (SCase _ e)      = [e]
children (SGuard _ e)     = [e]
children (SCycle _ _ e)   = [e]
children (SEquation _ e r)= [e,r]
children (SProj _ e)      = [e]
children _                = []

rebuild :: SExp a -> [SExp a] -> SExp a
rebuild (SApp l _)       es = SApp l es
rebuild (SWithin l _)    es = SWithin l es
rebuild (SCons l _ _) [h,t] = SCons l h t
rebuild (SIf l _)       [e] = SIf l e
rebuild (SCase l _)     [e] = SCase l e
rebuild (SGuard l _)    [e] = SGuard l e
rebuild (SCycle l v _)  [e] = SCycle l v e
rebuild (SEquation l _ _) [e,r] = SEquation l e r
rebuild (SProj l _)     [e] = SProj l e
rebuild sexp             _  = sexp

parent :: SExp FileNode -> FileNode
parent (SWithin _ (x:_)) = parent x
parent (SEquation _ x _) = parent x
parent x                 = getParentNode (label x)

funId :: SExp a -> String
funId (SApp _ es) = funId (head es)
funId (SWithin _ es) = funId (head es)
funId (SCons _ _ _) = ":"
funId (SNil _) = "[]"
funId (SEquation _ e _) = funId e
funId (SId _ s _) = s
funId _ = ""

funLabel :: SExp FileNode -> FileNode
funLabel (SApp _ es) = funLabel (head es)
funLabel (SWithin _ es) = funLabel (head es)
funLabel (SEquation _ e _) = funLabel e
funLabel (SId n _ _) = n
funLabel (SNil n) = n
funLabel _ = LowLevel.nil

fst3 (x,_,_) = x
snd3 (_,x,_) = x
thd3 (_,_,x) = x

-- conversion function
-- If `cyc' boolean is True, then cycles are located and expressed as SCycles,
-- if False, there is the danger of obtaining an infinite expression.
-- If `verbose' boolean is True, then unevaluated arguments appear in the
-- result, otherwise they are represented by SUnevaluated.
fileNode2SExp :: Int -> Bool -> Bool -> FileNode -> SExp FileNode
fileNode2SExp cutoff cyc verbose node =
  case go cutoff cyc verbose [] node of (e,_,_) -> e  
  where
  simple e = (e,[],[])
  go :: Int			-- cutoff depth
       -> Bool			-- look for cycles?
       -> Bool			-- transform unevaluated args?
       -> [(FileNode,String)]	-- enclosing nodes w/ variable name for `where'
       -> FileNode 		-- root node of expression
       -> ( SExp FileNode	-- expression 
          , [FileNode]		-- nodes that start cycle
          , [String] )		-- variable names occurring (except for cycles)
  go 0 cyc verbose nodesAbove node = simple (SCut node)
  go depth cyc verbose nodesAbove node =
    case nodeType node of
      Trace -> case traceType node of
        TAp -> 
          let partCycles :: [FileNode]
              partCycles = funCycles ++ concat argsCycles
              partVars :: [String]
              partVars = funVars ++ concat argsVars
              isCycle = cyc && node `elem` partCycles
              var :: String
              var = head . filter (not . (`elem` partVars)) .
                    map (("cyc"++) . show) $ [1..] 
              newNodesAbove :: [(FileNode,String)]
              newNodesAbove = (node,var) : nodesAbove
              apArgs :: [FileNode]
              apArgs = getApArgs node
              -- (fun,funCycles,funVars) = ... not accepted by nhc98
              fun = fst3 z
              funCycles = snd3 z
              funVars = thd3 z
              z = let f = head apArgs in
                  if f==LowLevel.nil then simple (SCut f)
                  else go depth cyc verbose newNodesAbove f
              args = fst3 zs
              argsCycles = snd3 zs
              argsVars = thd3 zs
              -- (args,argsCycles,argsVars) = ... not accepted by nhc98
              zs = unzip3 $ map (go (depth-1) cyc verbose newNodesAbove)
                                (tail apArgs)
              sexp = case fun of
                SLiteral n "if" ->
                  if length args == 2	-- small weirdness in the .hat file
                    then within node
                                (fileNode2SExp depth cyc verbose
                                               (getParentNode n))
                                (SIf n (head args))
                    else errorT ("SExp: bad If at 0x"++showHex (int n) "")
                SLiteral n "case" ->
                  if length args == 1
                    then within node
                                (fileNode2SExp depth cyc verbose
                                               (getParentNode n))
                                (SCase n (head args))
                    else errorT ("SExp: bad Case at 0x"++showHex (int n) "")
                SLiteral n "|" ->
                  if length args == 2	-- small weirdness in the .hat file
                    then within node
                                (fileNode2SExp depth cyc verbose
                                               (getParentNode n))
                                (SGuard n (head args))
                    else errorT ("SExp: bad Guard at 0x"++showHex (int n) "")
                SLiteral n "[]" | not verbose -> SNil n
                SLiteral n ":"  | length args == 2 && not verbose ->
                  case args!!1 of
                    SLiteral m "[]" -> SCons node (args!!0) (SNil m)
                    SCons _ _ _     -> SCons node (args!!0) (args!!1)
                    SNil m          -> SCons node (args!!0) (SNil m)
                    _ -> SApp node (fun:args)
                SApp n args1 -> SApp n (args1++args)  -- combine applications
                _ -> SApp node (fun:args)
          in case lookup node nodesAbove of
               Just var -> (SId node var SInfixDefault,[node],[]) 
                           -- `lower' end of cycle
               Nothing ->
                 case sexp of
                   SApp n [SId n' "IO" fixity,_] -> 
                     -- get argument of IO, otherwise we see IO <hidden>
                     let arg = head (tail apArgs)
                         par = getParentNode arg
                         newArg = if nodeType arg == Trace
                                     && traceType arg == THidden
                                  then par else arg
                         z = go depth cyc verbose newNodesAbove newArg
                     in (SApp n [SId n' "IO" fixity,fst3 z],snd3 z,thd3 z)
                   _ -> -- normal case 
                     (if isCycle then SCycle node var sexp else sexp
                      ,partCycles,partVars)
        TNm ->
          let i = getIdent node in
          if isLiteral node
          then simple (SLiteral node i)
          else (SId node i
                  (case i of
                     "."  | getIdentMod node == "Prelude" -> SAssoc 9 i
                     "++" | getIdentMod node == "Prelude" -> SAssoc 5 i
                     "&&" | getIdentMod node == "Prelude" -> SAssoc 3 i
                     "||" | getIdentMod node == "Prelude" -> SAssoc 2 i
                     "*"  | getIdentMod node == "Prelude" -> SAssoc 7 i
                     "+"  | getIdentMod node == "Prelude" -> SAssoc 6 i
                     ">>" | getIdentMod node == "Prelude" -> SAssoc 1 i
                     ">>=" | getIdentMod node == "Prelude" -> SAssoc 1 i
                     _ -> transFixity (getFixity node))
               ,[]
               ,if isConstructor node then [] else [i])
        TInd -> simple $ SProj node
                  (fst3 (go depth cyc verbose nodesAbove ((getApArgs node)!!1)))
        THidden -> simple $ SLiteral node "{_}"
        TSatA -> if verbose
                 then go depth cyc verbose nodesAbove (head (getApArgs node))
                 else simple $ SUnevaluated node
        TSatB -> simple $ SBottom node
        TSatC -> go depth cyc verbose nodesAbove (head (getApArgs node))
        TSatAL -> if verbose
                  then go depth cyc verbose nodesAbove (head (getApArgs node))
                  else simple $ SUnevaluated node
        TSatBL -> simple $ SBottom node
        TSatCL -> go depth cyc verbose nodesAbove (head (getApArgs node))

      ModuleInfo -> errorT ("SExp: got a ModInfo at 0x"++showHex (int node) "")
      NmType -> errorT ("SExp: got a NmType at 0x"++showHex (int node) "")
      SR -> errorT ("SExp: got an SR at 0x"++showHex (int node) "")
  within :: FileNode -> SExp FileNode -> SExp FileNode -> SExp FileNode
  within node parent exp =
    case parent of SWithin n ps -> SWithin node (ps++[exp])
                   _            -> SWithin node [parent,exp]


-- useful document combinators:
-- non-breaking space
(<->) :: Doc -> Doc -> Doc
d1 <-> d2 = d1 <> delimiter " " <> d2 

-- breakable space
(<+>) :: Doc -> Doc -> Doc
d1 <+> d2 = d1 <> fdelimiter " " <> d2

-- breakable non-space
(<|>) :: Doc -> Doc -> Doc
d1 <|> d2 = d1 <> fdelimiter "" <> d2

-- breakable before a comma
(<*>) :: Doc -> Doc -> Doc
d1 <*> d2 = d1 <|> text "," <> d2

-- breakable before or after a cons
(<:>) :: Doc -> Doc -> Doc
d1 <:> d2 = d1 <|> text ":" <|> d2

indentation :: Int
indentation = 2


isOpSym :: String -> Bool
isOpSym sym = let c = head sym in not (isAlpha c || c == '[') 

funDoc :: String -> Doc
funDoc var = (if isOpSym var then parens else id) $ text var 

opDoc :: String -> Doc
opDoc var =  text (if isAlpha (head var) then ('`' :var++ "`") else var)


data ArgPos = ALeft | ARight
isRight ARight = True
isRight ALeft = False
isLeft = not . isRight

-- surround by parentheses if necessary
-- first fixity of surrounding expression, then if left or right argument,
-- then fixity of expression itself
optParens :: SFixity -> ArgPos -> SFixity -> Doc -> Doc
optParens surFixity aPos ownFixity =
  case (priority surFixity) `compare` (priority ownFixity) of
    LT -> if priority surFixity == (-1) then groupNest indentation else id
    GT -> groupNest indentation . parens
    EQ -> if (isInfixR surFixity && isInfixR ownFixity && isRight aPos)
            || (isInfixL surFixity && isInfixL ownFixity && isLeft aPos)
            || sameAssoc surFixity ownFixity 
            then id
            else groupNest indentation . parens

sameAssoc :: SFixity -> SFixity -> Bool
sameAssoc (SAssoc _ var1) (SAssoc _ var2) = True
sameAssoc _ _ = False

foldr0 :: (a -> a -> a) -> a -> [a] -> a
foldr0 f c [] = c
foldr0 f c xs = foldr1 f xs

newlistDoc :: Eq a => (a->Doc->Doc) -> SFixity -> ArgPos -> SExp a -> Doc
newlistDoc high surFixity aPos e@(SCons _ hd tl) =
    if all isCharLiteral (elems e) then
         unsafePerformIO (hPutStrLn stderr "newListDoc: got a char list\n")
         `seq`
         groupNest 1 (text "\"" <> string e)
    else group (text "[" <> commas e)
  where
    elems (SNil _) = []
    elems (SCons _ hd tl) = hd: elems tl

    isCharLiteral (SLiteral _ ('\'':_)) = True
    isCharLiteral (SProj _ e) = isCharLiteral e
    isCharLiteral (SCut _)    = True
    isCharLiteral _           = False

    char :: Eq a => SExp a -> Doc
    char (SLiteral v ('\'':cs)) = high v $ text (init cs)
    char (SProj v e) = char e
    char (SCut v) = high v $ highlight [ReverseVideo] $ text " "

    colons = sExpFold False (sExp2Doc False high) (text ":") (text "[]")
    commas = sExpFold True  (sExp2Doc False high) (text ",") (text "]")
    string = sExpFold False char                  Pretty.nil (text "\"")

    -- first arg is boolean:  True = separators;  False = terminators.
    sExpFold sepr head cons nil (SApp v [SId c ":" _, hd, SId n "[]" _])
      | sepr =
        high v (head hd <|> high n nil)
    sExpFold sepr head cons nil (SApp v [SId c ":" _, hd, tl]) =
        high v (head hd <|> high c (cons) <|> sExpFold sepr head cons nil tl)
    sExpFold sepr head cons nil (SId n "[]" _) = high n nil
    sExpFold sepr head cons nil (SProj _ e)    = sExpFold sepr head cons nil e
    sExpFold sepr head cons nil e              = sExp2Doc False high e

listDoc :: Eq a => (a->Doc->Doc) -> SFixity -> ArgPos -> SExp a -> Doc
listDoc high surFixity aPos e =
--unsafePerformIO (hPutStrLn stderr "oldListDoc: got a list\n")
--  `seq`
  case maybeRest of
    Nothing -> if all isCharLiteral elems 
                 then groupNest 1 $ 
                        text "\"" <> string e <> text "\""
                 else group $ text "[" <> commas e
    Just eRest -> groupNest indentation
                  . optParens surFixity aPos (SInfixR 5)
                  . colons $ e
  where
  (elems,maybeRest) = getListElems e
  getListElems :: SExp a -> ([SExp a],Maybe (SExp a))
  getListElems (SApp va [SId vt ":" _,ee,er]) = (ee:ees,mr)
                                            where (ees,mr) = getListElems er
  getListElems (SProj _ e)    = getListElems e
  getListElems (SId _ "[]" _) = ([],Nothing)
  getListElems (SCut _)       = ([],Nothing)
  getListElems e              = ([],Just e)
  isCharLiteral :: SExp a -> Bool
  isCharLiteral (SLiteral _ ('\'':_)) = True
  isCharLiteral (SProj _ e) = isCharLiteral e
  isCharLiteral (SCut _)    = True
  isCharLiteral _           = False

  char :: Eq a => SExp a -> Doc
  char (SLiteral v ('\'':cs)) = high v $ text (init cs)
  char (SProj v e) = char e
  char (SCut v) = high v $ highlight [ReverseVideo] $ text " "

  colons = sExpFold False (sExp2Doc False high) (text ":") (text "[]")
  commas = sExpFold True  (sExp2Doc False high) (text ",") (text "]")
  string = sExpFold False char                  Pretty.nil Pretty.nil

  -- first arg is boolean:  True = separators;  False = terminators.
  sExpFold sepr head cons nil (SApp v [SId c ":" _, hd, SId n "[]" _]) | sepr =
      high v (head hd <|> high n nil)
  sExpFold sepr head cons nil (SApp v [SId c ":" _, hd, tl]) =
      high v (head hd <|> high c (cons) <|> sExpFold sepr head cons nil tl)
  sExpFold sepr head cons nil (SId n "[]" _) = high n nil
  sExpFold sepr head cons nil (SProj _ e)    = sExpFold sepr head cons nil e
  sExpFold sepr head cons nil e              = sExp2Doc False high e

priority (SInfix p) = p
priority (SInfixL p) = p
priority (SInfixR p) = p
priority (SAssoc p _) = p
priority SInfixDefault = 9

isInfixL (SInfixL _) = True
isInfixL SInfixDefault = True
isInfixL _ = False

isInfixR (SInfixR _) = True
isInfixR _ = False

isNotInfixDefault SInfixDefault = False
isNotInfixDefault _ = False

considerAsOperator :: String -> SFixity -> Bool
considerAsOperator var fixity =  isOpSym var || isNotInfixDefault fixity 

-- A central function.  Convert an SExpression to a Document using
-- the pretty-printing combinators.  In verbose mode, strings and lists
-- are shown in full with cons applications, rather than sugared.
sExp2Doc :: Eq a => Bool -> (a->Doc->Doc) -> SExp a -> Doc
sExp2Doc verbose high = goDoc (SInfix (-1)) ARight 
  where
  -- fixity of surrounding expression and which sort of argument
  goDoc :: SFixity -> ArgPos -> SExp a -> Doc
  goDoc surFixity aPos (SApp va ((SId vt (',':xs) _):args)) =
    if length xs + 2 == length args 
    then high va $
         group (text "(" <> foldr1 comma (map (sExp2Doc verbose high) args))
                         <> text ")"
           -- print tuple properly
    else high va $
         optParens surFixity aPos ownFixity
         . (high vt (text ("(,"++xs++")")) <+>)
         . foldr1 (<+>) . map (goDoc ownFixity ARight) $ args
           -- partial application of tuple constructor
    where
    ownFixity = SInfix 10
    comma l r = l <|> high vt (text ",") <> r

  goDoc surFixity aPos (e@(SApp v [SId _ ":" _,_,_])) | not verbose =
    high v $ listDoc high surFixity aPos e
  goDoc surFixity aPos (SApp va [SId vf var ownFixity,e1,e2]) 
    | considerAsOperator var ownFixity = 
      high va $
      optParens surFixity aPos ownFixity
        (goDoc ownFixity ALeft e1
        <+> high vf (opDoc var)
        <+> goDoc ownFixity ARight e2)
  goDoc surFixity aPos (SApp va [SId vf var ownFixity,e]) 
    | considerAsOperator var ownFixity =
    -- show infix operator with single argument as section
    groupNest indentation . high va . parens $
      goDoc ownFixity ALeft e <-> high vf (opDoc var)
  goDoc surFixity aPos (SApp va (fun:args)) =
    high va
    . optParens surFixity aPos ownFixity
    . (goDoc ownFixity ALeft fun <+>)
    .  foldr1 (<+>) . map (goDoc ownFixity ARight) $ args
    where
    ownFixity = SInfix 10
  goDoc surFixity aPos e@(SCons v _ _) | not verbose =
    high v $ newlistDoc high surFixity aPos e
  goDoc surFixity aPos e@(SNil v) =
    high v $ text "[]"
  goDoc _ _ (SId v var fixity) = high v $ funDoc var
  goDoc _ _ (SLiteral v lit) = high v $ text lit
--goDoc _ _ (SLambda v) = high v $ text "(\\..)"
  goDoc _ _ (SWithin v es) = 
    high v $ foldr1 (\a b-> a <+> group (bold (text "|") <+> b))
                    (map (sExp2Doc verbose high) es)
  goDoc _ _ (SIf v exp) = 
    (groupNest indentation . high v) (bold (text "if")
                                     <+> sExp2Doc verbose high exp)
  goDoc _ _ (SCase v exp) = 
    (groupNest indentation . high v) (bold (text "case")
                                     <+> sExp2Doc verbose high exp)
  goDoc _ _ (SGuard v exp) = 
    (groupNest indentation . high v) (sExp2Doc verbose high exp)
  goDoc _ _ (SCut v) = high v $ highlight [ReverseVideo] (text " ")
  goDoc _ _ (SUnevaluated v) = high v $ text "_"
  goDoc _ _ (SBottom v) = high v $ text "_|_"
  goDoc _ _ (SCycle v var exp) = 
    groupNest indentation . high v . parens $
      text var <+> group (text "where" <+> sExp2Doc verbose high exp)
  goDoc _ _ (SEquation v lhs rhs) =
    group (high v (sExp2Doc verbose high lhs
                  <+> text "=" <+> sExp2Doc verbose high rhs))
  goDoc surFixity aPos (SProj v exp) = high v $ goDoc surFixity aPos exp


{-
-- only for testing:
test1 =
  SApp (mkHatNode 0)
    [ SId (mkHatNode 1) "fun" SInfixDefault
    , SLiteral (mkHatNode 2) "24"
    , SLiteral (mkHatNode 3) "True"
    , SApp (mkHatNode 4)
        [ SId (mkHatNode 5) "+" (SInfixL 5)
        , SLiteral (mkHatNode 6) "3"
        , SLiteral (mkHatNode 7) "4"
        ]
    ]

test2 =
  SApp (mkHatNode 10)
    [ SId (mkHatNode 11) "*" (SInfixL 6)
    , SApp (mkHatNode 12)
        [ SId (mkHatNode 13) "+" (SInfixL 5)
        , SApp (mkHatNode 14)
            [ SId (mkHatNode 15) "-" (SInfixL 5)
            , SLiteral (mkHatNode 16) "3"
            , SLiteral (mkHatNode 17) "6"
            ]
        , SApp (mkHatNode 18)
            [ SId (mkHatNode 19) "-" (SInfixL 5)
            , SLiteral (mkHatNode 20) "3"
            , SLiteral (mkHatNode 21) "6"
            ]
        ]
    , test1
    ]
-}