content: module Agda.Compiler.JS.Pretty where
  
  import GHC.Generics (Generic)
  
  import Data.Char ( isAsciiLower, isAsciiUpper, isDigit )
  import Data.List ( intercalate )
  import Data.String ( IsString (fromString) )
  import Data.Set ( Set, toList, insert, member )
  import qualified Data.Set as Set
  import Data.Map ( Map, toAscList )
  import qualified Data.Text as T
  
  import Agda.Syntax.Common ( Nat )
  
  import Agda.Utils.Function ( applyWhen )
  import Agda.Utils.Hash
  import Agda.Utils.List ( indexWithDefault )
  import Agda.Utils.List1 ( List1, pattern (:|), (<|) )
  import qualified Agda.Utils.List1 as List1
  
  import Agda.Utils.Impossible
  
  import Agda.Compiler.JS.Syntax hiding (exports)
  
  -- Pretty-print a lambda-calculus expression as ECMAScript.
  
  --- The indentation combinators of the pretty library does not fit C-like languages
  --- like ECMAScript.
  --- A simple pretty printer is implemented with a better `indent` and punctuation compaction.
  ---
  --- More explanation:
  ---
  --- I have struggled with different pretty printers, and at the end it was much easier
  --- to implement and use this ~100 SLOC code pretty printer library.
  --- It produces really better quality indentation than I could achieve with the
  --  standard pretty printers.
  --- This library code is only used in this module, and it is specialized to pretty
  --- print JavaScript code for the Agda backend, so I think its best place is in this module.
  data JSModuleStyle = JSES6 | JSCJS | JSAMD
    deriving Generic
  
  data Doc
      = Doc String
      | Indent Int Doc
      | Group Doc
      | Beside Doc Doc
      | Above Doc Doc
      | Enclose Doc Doc Doc
      | Space
      | Empty
  
  minifiedCodeLinesLength :: Int
  minifiedCodeLinesLength = 500
  
  render :: Bool -> Doc -> String
  render minify = intercalate "\n" . joinLines . map (uncurry mkIndent) . go 0
    where
      joinLines :: [String] -> [String]
      joinLines = applyWhen minify $ chunks 0 []
        where
          chunks len acc [] = [concat (reverse acc)]
          chunks len acc (s: ss)
              | len + n <= minifiedCodeLinesLength = chunks (len + n) (s: acc) ss
              | otherwise = concat (reverse acc): chunks n [s] ss
            where
              n = length s
  
      joinBy f [x] (y: ys) = f x y ++ ys
      joinBy f (x:xs) ys = x: joinBy f xs ys
      joinBy f xs ys = xs ++ ys
  
      mkIndent n s | minify = s
      mkIndent n "" = ""
      mkIndent n s = replicate n ' ' ++ s
  
      overlay (i, s) (j, s') | all punctuation (s ++ s') && n > 0 = [(i, s ++ mkIndent n s')]
        where n = j - (i + length s)
      overlay (j, s') (i, s) | all punctuation (s ++ s') && n > 0 = [(i, s' ++ mkIndent n s)]
        where n = j - (i + length s)
      overlay a b = [a, b]
  
      punctuation = (`elem` ("(){}[];:, " :: String))
  
      go i Space = [(i, " ") | not minify]
      go i Empty = []
      go i (Doc s) = [(i, s)]
      go i (Beside d d') = joinBy (\(i, s) (_, s') -> [(i, s ++ s')]) (go i d) (go i d')
      go i (Above d d') = joinBy overlay (go i d) (go i d')
      go i (Indent j d) = go (i + j) d
      go i (Enclose open close d) = go i $ Group $ Above open $ Above d close
      go i (Group d)
          | size ss < 40 = compact ss
          | otherwise    = ss
        where
          ss = go i d
          size = sum . map (length . snd)
          compact [] = []
          compact ((i, x): xs) = [(i, x ++ concatMap snd xs)]
  
  instance IsString Doc where
      fromString = Doc
  
  instance Semigroup Doc where
      Empty <> d = d
      d <> Empty = d
      d <> d' = Beside d d'
  
  instance Monoid Doc where
      mempty = Empty
      mappend = (<>)
  
  infixr 5 $+$
  infixr 5 $++$
  infixr 6 <+>  -- fixity has to match the one of Semigroup.(<>)
  
  ($+$) :: Doc -> Doc -> Doc
  Empty $+$ d = d
  d $+$ Empty = d
  d $+$ d' = Above d d'
  
  -- | Separate by blank line.
  
  ($++$) :: Doc -> Doc -> Doc
  Empty $++$ d = d
  d $++$ Empty = d
  d $++$ d' = d `Above` "" `Above` d'
  
  -- | Separate by space that will be removed by minify.
  --
  -- For non-removable space, use @d <> " " <> d'@.
  
  (<+>) :: Doc -> Doc -> Doc
  Empty <+> d = d
  d <+> Empty = d
  d <+> d' = d `Beside` Space `Beside` d'
  
  text :: String -> Doc
  text = Doc
  
  group :: Doc -> Doc
  group = Group
  
  indentBy :: Int -> Doc -> Doc
  indentBy i Empty = Empty
  indentBy i (Indent j d) = Indent (i + j) d
  indentBy i d = Indent i d
  
  enclose :: Doc -> Doc -> Doc -> Doc
  enclose open close (Enclose o c d) = Enclose (open <> o) (c <> close) d
  enclose open close (Indent _ (Enclose o c d)) = Enclose (open <> o) (c <> close) d
  enclose open close d = Enclose open close d
  
  ----------------------------------------------------------------------------------------------
  
  space :: Doc
  space = Space
  
  indent :: Doc -> Doc
  indent = indentBy 2
  
  hcat :: [Doc] -> Doc
  hcat = foldr (<>) mempty
  
  vcat :: [Doc] -> Doc
  vcat = foldr ($+$) mempty
  
  -- | Concatenate vertically, separated by blank lines.
  
  vsep :: [Doc] -> Doc
  vsep = foldr ($++$) mempty
  
  punctuate :: Doc -> [Doc] -> Doc
  punctuate _ []     = mempty
  punctuate p (x:xs) = indent $ vcat $ go x xs
                     where go y []     = [y]
                           go y (z:zs) = (y <> p) : go z zs
  
  parens, brackets, braces :: Doc -> Doc
  parens   = enclose "(" ")"
  brackets = enclose "[" "]"
  braces   = enclose "{" "}"
  
  -- | Apply 'parens' to 'Doc' if boolean is true.
  mparens :: Bool -> Doc -> Doc
  mparens True  d = parens d
  mparens False d = d
  
  ----------------------------------------------------------------------------------------------
  
  unescape :: Char -> String
  unescape '"'      = "\\\""
  unescape '\\'     = "\\\\"
  unescape '\n'     = "\\n"
  unescape '\r'     = "\\r"
  unescape '\x2028' = "\\u2028"
  unescape '\x2029' = "\\u2029"
  unescape c        = [c]
  
  unescapes :: String -> Doc
  unescapes s = text $ concatMap unescape s
  
  -- pretty (n,b) i e pretty-prints e, under n levels of de Bruijn binding
  --   if b is true then the output is minified
  
  class Pretty a where
      pretty :: (Nat, Bool, JSModuleStyle) -> a -> Doc
      prettyPrec :: Int -> (Nat, Bool, JSModuleStyle) -> a -> Doc
  
      pretty = prettyPrec 0
      prettyPrec = const pretty
  
      {-# MINIMAL pretty | prettyPrec #-}
  
  prettyShow :: Pretty a => Bool -> JSModuleStyle -> a -> String
  prettyShow minify ms = render minify . pretty (0, minify, ms)
  
  instance Pretty a => Pretty (Maybe a) where
    prettyPrec p n = maybe mempty (prettyPrec p n)
  
  instance (Pretty a, Pretty b) => Pretty (a,b) where
    pretty n (x,y) = pretty n x <> ":" <+> pretty n y
  
  -- Pretty-print collections
  
  class Pretties a where
      pretties :: (Nat, Bool, JSModuleStyle) -> a -> [Doc]
  
  instance Pretty a => Pretties [a] where
    pretties n = map (pretty n)
  
  instance Pretty a => Pretties (List1 a) where
    pretties n = pretties n . List1.toList
  
  instance (Pretty a, Pretty b) => Pretties (Map a b) where
    pretties n = pretties n . toAscList
  
  -- Pretty print identifiers
  
  instance Pretty LocalId where
    pretty (n, _, _) (LocalId x) = text $ indexWithDefault __IMPOSSIBLE__ vars (n - x - 1)
      where
        vars = ("": map show [0..]) >>= \s -> map (:s) ['a'..'z']
  
  instance Pretty GlobalId where
    pretty n (GlobalId m) = text $ variableName $ intercalate "_" m
  
  instance Pretty MemberId where
    pretty _ (MemberId s) = "\"" <> unescapes s <> "\""
    pretty n (MemberIndex i comment) = text (show i) <> pretty n comment
  
  instance Pretty Comment where
    pretty _ (Comment "") = mempty
    pretty (_, True, _) _ = mempty
    pretty _ (Comment s) = text $ "/* " ++ s ++ " */"
  
  -- Pretty print expressions
  
  instance Pretty Exp where
    prettyPrec p n (Self)         = "exports"
    prettyPrec p n (Local x)      = pretty n x
    prettyPrec p n (Global m)     = pretty n m
    prettyPrec p n (Undefined)    = "undefined"
    prettyPrec p n (Null)         = "null"
    prettyPrec p n (String s)     = "\"" <> unescapes (T.unpack s) <> "\""
    prettyPrec p n (Char c)       = "\"" <> unescapes [c] <> "\""
    prettyPrec p n (Integer x)    = "agdaRTS.primIntegerFromString(\"" <> text (show x) <> "\")"
    prettyPrec p n (Double x)     = text $ show x
    prettyPrec p (n, min, ms) (Lambda x e) = mparens (p > 2) $
      mparens (x /= 1) (punctuate "," (pretties (n + x, min, ms) (map LocalId [x-1, x-2 .. 0])))
      <+> "=>" <+> block (n + x, min, ms) e
    prettyPrec p n (Object o)     = braces $ punctuate "," $ pretties n o
    prettyPrec p n (Array es)     = brackets $ punctuate "," [pretty n c <> pretty n e | (c, e) <- es]
    prettyPrec p n (Apply f es)   = prettyPrec 17 n f <> parens (punctuate "," $ pretties n es)
    prettyPrec p n (Lookup e l)   = prettyPrec 17 n e <> brackets (pretty n l)
    prettyPrec p n (If e f g)     = mparens (p > 2) $
      prettyPrec 3 n e <> "?" <+> prettyPrec 3 n f <> ":" <+> prettyPrec 2 n g
    prettyPrec p n (PreOp op e)   = parens $ text op <> " " <> prettyPrec 17 n e
    prettyPrec p n (BinOp e op f) = parens $ prettyPrec 17 n e <> " " <> text op <> " " <> prettyPrec 17 n f
    prettyPrec p n (Const c)      = text c
    prettyPrec p n (PlainJS js)   = text js
  
  block :: (Nat, Bool, JSModuleStyle) -> Exp -> Doc
  block n e = mparens (doNest e) $ prettyPrec 2 n e
    where
      doNest Object{} = True
      doNest _ = False
  
  modname :: GlobalId -> Doc
  modname (GlobalId ms) = text $ intercalate "." ms
  
  exports :: (Nat, Bool, JSModuleStyle) -> Set JSQName -> [Export] -> Doc
  exports n lss [] = Empty
  exports n lss es0@(Export ls e : es)
    -- If the parent of @ls@ is already defined (or no parent exists), @ls@ can be defined
    | maybe True (`member` lss) parent =
        "exports" <> hcat (map brackets (pretties n ls)) <+> "=" <+> indent (pretty n e) <> ";" $+$
        exports n (insert ls lss) es
    -- If the parent is not yet defined, first define it as empty object, and then continue with @ls@.
    | otherwise =
        exports n lss $ maybe es0 (\ ls' -> Export ls' (Object mempty) : es0) parent
    where
    parent = List1.nonEmpty $ List1.init ls
  
  instance Pretty [(GlobalId, Export)] where
    pretty n es
      = vcat [ pretty n g <> hcat (map brackets (pretties n ls)) <+> "=" <+> indent (pretty n e) <> ";"
             | (g, Export ls e) <- es ]
  
  instance Pretty Module where
    pretty opt@(n, min, JSES6) (Module m is es callMain) = vsep
      [ "import agdaRTS from \"./agda-rts.mjs\";"
      , imports
      , "const exports = {};"
      , exports opt Set.empty es
      , pretty opt callMain
      , ";export default exports;"
      ]
      $+$ ""
      where
        imports = vcat [
          "import " <> indent (pretty opt e) <+> " from \"./" <> modname e <> ".mjs\";"
          | e <- toList (globals es <> Set.fromList is)
          ]
        les = toList (globals es <> Set.fromList is)
    pretty opt@(n, min, JSCJS) (Module m is es callMain) = vsep
      [ "var agdaRTS" <+> "=" <+> "require(\"agda-rts\");"
      , imports
      , exports opt Set.empty es
      , pretty opt callMain
      ]
      $+$ ""
      where
        imports = vcat [
          "var " <> indent (pretty opt e) <+> "=" <+> "require(\"" <> modname e <> "\");"
          | e <- toList (globals es <> Set.fromList is)
          ]
        les = toList (globals es <> Set.fromList is)
    pretty opt@(n, min, JSAMD) (Module m is es callMain) = vsep
      [ "define(['agda-rts'"
        <+> hcat [ ", " <+> ("\"" <> modname e <> "\"") | e <- les ]
        <+> "],"
      , "function(agdaRTS"
        <+> hcat [ ", " <+> pretty opt e | e <- les ]
        <+> ") {"
      , "var exports = {};"
      , exports opt Set.empty es
      , pretty opt callMain
      , "; return exports; });"
      ]
      $+$ "" -- Final newline
      where
        les = toList (globals es <> Set.fromList is)
  
  
  variableName :: String -> String
  variableName s = if isValidJSIdent s then "z_" ++ s else "h_" ++ show (hashString s)
  
  -- | Check if a string is a valid JS identifier. The check ignores keywords
  -- as we prepend z_ to our identifiers. The check
  -- is conservative and may not admit all valid JS identifiers.
  
  isValidJSIdent :: String -> Bool
  isValidJSIdent []     = False
  isValidJSIdent (c:cs) = validFirst c && all validOther cs
    where
    validFirst :: Char -> Bool
    validFirst c = isAsciiUpper c || isAsciiLower c || c == '_' || c == '$'
  
    validOther :: Char -> Bool
    validOther c = validFirst c || isDigit c
  