{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}

module BNFC.Backend.Haskell.Printer where

import BNFC.Prelude

import Control.Monad.State

import qualified Data.Map                  as Map
import           Data.List                 (intersperse, sortBy, (\\))
import           Data.String               (fromString)

import Prettyprinter

import System.FilePath (takeBaseName)

import BNFC.Backend.Common.Utils as Utils
import BNFC.Backend.CommonInterface.Backend

import BNFC.Backend.Haskell.Utilities.Printer
import BNFC.Backend.Haskell.Utilities.Utils
import BNFC.Backend.Haskell.Options
import BNFC.Backend.Haskell.State

import BNFC.CF
import BNFC.Options.GlobalOptions


haskellPrinter :: LBNF -> State HaskellBackendState Result
haskellPrinter lbnf = do
  st <- get
  let cfName        = takeBaseName $ optInput $ globalOpt st
      inDirectory   = inDir $ haskellOpts st
      nSpace        = nameSpace $ haskellOpts st
      useGadt       = gadt $ haskellOpts st
      rules         = filterRules $ astRules st
      tks           = tokens st
      funct         = functor $ haskellOpts st
      tt            = tokenText $ haskellOpts st
      prettyPrinter = cf2printer lbnf cfName inDirectory nSpace useGadt rules tks funct tt
  return [(mkFilePath inDirectory nSpace cfName "Print" "hs", prettyPrinter)]
  where
    filterRules :: [(Type, [(Label, ([Type], (Integer, ARHS)))])]
                -> [(Type, [(Label, ([Type], (Integer, ARHS)))])]
    filterRules rules =
      filter
      (\(_,l) -> not (null l))
      ((\(f,s) -> (f, filterLabelsPrinter fNames s)) <$> rules)

    -- Functions names.
    fNames :: [String]
    fNames = toList <$> Map.keys (_lbnfFunctions lbnf)

cf2printer :: LBNF
           -> String
           -> Bool
           -> Maybe String
           -> Bool
           -> [(Type, [(Label, ([Type], (Integer, ARHS)))])]
           -> [(CatName,TokenDef)]
           -> Bool
           -> TokenText
           -> String
cf2printer lbnf cfName inDir nameSpace gadt astRules tks funct tokenText =
  docToString defaultLayoutOptions $ cf2doc lbnf cfName inDir nameSpace gadt astRules tks funct tokenText


cf2doc :: LBNF
       -> String
       -> Bool
       -> Maybe String
       -> Bool
       -> [(Type, [(Label, ([Type], (Integer, ARHS)))])]
       -> [(CatName,TokenDef)]
       -> Bool
       -> TokenText
       -> Doc ()
cf2doc lbnf cfName inDir nameSpace gadt astRules tokens functor tokenText =
  (vsep . intersperse emptyDoc) $
    printPrologue lbnf cfName inDir nameSpace gadt absName astRules
    :
    toBePrinted
  where

    absName :: ModuleName
    absName = mkModule inDir nameSpace cfName "Abs"

    tokenPrintInstances =
      if null tokens
      then Nothing
      else Just $ printTokenInstances absName tokenText tokens

    catPrintInstances =
      if null astRules
      then Nothing
      else Just $ printCatInstances absName functor astRules

    toBePrinted = catMaybes [tokenPrintInstances, catPrintInstances]

printPrologue :: LBNF
              -> String
              -> Bool
              -> Maybe String
              -> Bool
              -> ModuleName
              -> [(Type, [(Label, ([Type], (Integer, ARHS)))])]
              -> Doc ()
printPrologue lbnf cfFileName inDir nameSpace gadt absName astRules =
  (vsep . intersperse emptyDoc)
  [ printPragmas, printModule cfFileName, printImports
  , annUtilities, printTree, streamTree, printDocTree
  , printRenderTree, renderFunction, prtPrec, printClass, printClassOverlappable
  , prtChar, prtDouble, prtInteger, prtString, printString, mkEsc
  ]
  where
    processedCats = fst <$> astRules
    lstcts = fst <$> filter (isListType . fst) astRules
    annUtilities = printAnn
      (toks lbnf)
      (cats (processedCats \\ lstcts))
      (listcats lstcts)

    printPragmas :: Doc ()
    printPragmas = vsep $
      [ "{-# LANGUAGE CPP #-}"
      , "{-# LANGUAGE FlexibleInstances #-}"
      ]
      ++
      Utils.when gadt [ "{-# LANGUAGE GADTs #-}" ]
      ++
      [ "{-# LANGUAGE OverloadedStrings #-}"
      , "#if __GLASGOW_HASKELL__ <= 708"
      , "{-# LANGUAGE OverlappingInstances #-}"
      , "#endif"
      , emptyDoc
      , "{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}"
      ]

    printModule :: String -> Doc ()
    printModule cfName = vsep
      [ "-- | Pretty-printer for language" <+> fromString cfName
      , "--   Generated by the BNF converter."
      , emptyDoc
      ,"module" <+> fromString (mkModule inDir nameSpace cfName "Print")
      , indent 2 $ "(" <+> "printTree"
      , indent 2 $ "," <+> "streamTree"
      , indent 2 $ "," <+> "renderTree"
      , indent 2 $ "," <+> "render"
      , indent 2 $ "," <+> "annToAnsiStyle"
      , indent 2 $ "," <+> "Print"
      , indent 2 $ ")" <+> "where" ]

    printImports :: Doc ()
    printImports = vsep
      [ "import qualified" <+> fromString absName
      , emptyDoc
      , "import           Data.String    (fromString)"
      , "import qualified Data.Text      as T"
      , "import           Data.Text      (Text)"
      , "import           Data.Text.Lazy (unpack)"
      , emptyDoc
      , "import Prettyprinter"
      , "import Prettyprinter.Render.Util.SimpleDocTree"
      , "import Prettyprinter.Render.Terminal" ]

    printTree :: Doc ()
    printTree = vsep
      [ "-- | The top-level printing method."
      , "printTree :: Print a => a -> String"
      , "printTree = renderTree . streamTree annToAnsiStyle" ]

    streamTree :: Doc ()
    streamTree = vsep
      [ "streamTree :: Print a => (Doc Ann -> Doc AnsiStyle) -> a -> SimpleDocStream AnsiStyle"
      , "streamTree f a = layoutSmart defaultLayoutOptions $ f (docTree 0 a)" ]

    printDocTree :: Doc ()
    printDocTree = vsep
      [ "docTree :: Print a => Int -> a -> Doc Ann"
      , "docTree = prt" ]

    printClass :: Doc ()
    printClass = vsep
      [ "-- | The printer class does the job."
        , hang 2 $ vsep [ "class Print a where"
        , "prt :: Int -> a -> Doc Ann" ]
      ]

    printClassOverlappable :: Doc ()
    printClassOverlappable = hang 2 $ vsep
      [ "instance {-# OVERLAPPABLE #-} Print a => Print [a] where"
      , "prt i as = hsep $ map (prt i) as" ]

    printRenderTree :: Doc ()
    printRenderTree = vsep
      [ "renderTree :: SimpleDocStream AnsiStyle -> String"
      , "renderTree = unpack . renderLazy . render 0 False" ]

    prtPrec :: Doc ()
    prtPrec = vsep
      [ "prPrec :: Int -> Int -> Doc Ann -> Doc Ann"
      , "prPrec i j d = if i > j then parens d else d" ]

    prtInteger :: Doc ()
    prtInteger = vsep
      [ "instance Print Integer where"
      , indent 2 "prt _ x = annotate (Literal LitInteger) (fromString $ show x)" ]

    prtDouble :: Doc ()
    prtDouble = vsep
      [ "instance Print Double where"
      , indent 2 "prt _ x = annotate (Literal LitDouble) (fromString $ show x)" ]

    prtChar :: Doc ()
    prtChar = vsep
      [ "instance Print Char where"
      , indent 2 "prt _ x = annotate (Literal LitChar) (pretty '\\'' <> mkEsc '\\'' x <> pretty '\\'')"
      ]

    prtString :: Doc ()
    prtString = vsep
      [ "instance Print String where"
      , indent 2 "prt _ x = printString x" ]

    printString :: Doc ()
    printString = vsep
      [ "printString :: String -> Doc Ann"
      , "printString s = annotate (Literal LitString) (pretty '\"' <> hcat (map (mkEsc '\"') s) <> pretty '\"')" ]

    mkEsc :: Doc ()
    mkEsc = vsep
      [ "mkEsc :: Char -> Char -> Doc Ann"
      , "mkEsc q s = case s of"
      , "  s | s == q -> pretty '\\\\' <> pretty s"
      , "  '\\\\' -> fromString \"\\\\\\\\\""
      , "  '\\n' -> fromString \"\\\\n\""
      , "  '\\t' -> fromString \"\\\\t\""
      , "  s -> pretty s"
      ]


-- | Print tokens instances for the printer.

printTokenInstances :: ModuleName -> TokenText -> [(CatName,TokenDef)] -> Doc ()
printTokenInstances absName tokenText =
  vsep . intersperse emptyDoc . fmap (printTokenInstance absName tokenText)

printTokenInstance :: ModuleName -> TokenText -> (CatName,TokenDef) -> Doc ()
printTokenInstance absName tokenText (cName, tokenDef) =
  case tokenDef of
    (TokenDef PositionToken _ _)   ->
      hang 2 $ vsep
        [ "instance Print" <+> absModule <> cat'
            <+> "where"
        , "prt _" <+> parens (absModule <> cat' <+> "(_,i)")
          <+> rhs ]
    (TokenDef NoPositionToken _ _) ->
      hang 2 $ vsep
        [ "instance Print" <+> absModule <> cat'
            <+> "where"
        , "prt _" <+> parens (absModule <> cat' <+> "i")
          <+> rhs ]
  where
    cat' = parseTokenName cName
    absModule = fromString absName <> dot
    rhs = "=" <+> "token" <+> "Tok" <> cat' <+>
      parens
       ( if isStringToken tokenText
         then "fromString i"
         else "fromString $ T.unpack i"
        )

-- | Print cateries instances for the printer.

printCatInstances :: ModuleName
                  -> Bool
                  -> [(Type, [(Label, ([Type], (Integer, ARHS)))])]
                  -> Doc ()
printCatInstances absName functor =
  vsep .
  intersperse emptyDoc .
  fmap (uncurry (printCatInstance absName functor))

printCatInstance :: ModuleName
                 -> Bool
                 -> Type
                 -> [(Label, ([Type], (Integer, ARHS)))]
                 -> Doc ()
printCatInstance absName functor t labelsRhs =
  hang 2 $ vsep
    [ "instance Print" <+> name t <+> "where"
    , instances ]
  where
    absModule =  fromString absName <> "."
    name :: Type -> Doc ()
    name tt
      | isListType tt
      = if isBuiltinType tt
        then brackets $ parseType t
        else
          if isTokenType tt
          then brackets $ absModule <> parseType t
          else
            if functor then
              if isIdentType tt
              then brackets (absModule <> parseType t)
              else brackets (absModule <> parseType t <> "'" <+> "a")
            else
                brackets (absModule <> parseType t)
      | functor = parens $ absModule <> parseType t <> "'" <+> "a"
      | otherwise = absModule <> parseType t
    instances = vsep $ fmap snd $
      sortBy (compare `on` fst) $
      fmap
      (printCase absModule functor)
      ((\(l,(_,tup)) -> (l,tup)) <$> labelsRhs)

printCase :: Doc ()
          -> Bool
          -> (Label, (Integer, ARHS))
          -> (Integer, Doc ())
printCase absModule functor (label, (p, arhs)) = case label of
  LId _ ->
    (0, leftRhs <+> "=" <+> "prPrec i" <+> (fromString . show) p
    <+> if null arhs
        then "emptyDoc"
        else
          if length arhs == 1
          then "$" <+> hsep (rhsToPrint arhs)
          else "$ hsep" <+> brackets (hsep $ rhsToPrint arhs))
    where
      leftRhs
        | null (printArgs arhs)
        = if functor then
            "prt i" <+> parens (absModule <> (fromString . printLabelName) label <+> "_")
          else
            "prt i" <+> absModule <> (fromString . printLabelName) label
        | functor
        = "prt i" <+> "(" <> absModule <> (fromString . printLabelName) label
            <+> "_" <+> hsep (printArgs arhs) <> ")"
        | otherwise
        = "prt i" <+> "(" <> absModule <> (fromString . printLabelName) label
            <+> hsep (printArgs arhs) <> ")"
  LDef _ -> panic "LDef labels should have been filtered out"
  LWild  -> panic "LWild labels should have been filtered out"
  LNil -> (p, "prt" <+> "_" <+> "[]" <+> "=" <+> "emptyDoc")
  LCons -> (p, "prt" <+> "_" <+> lcons <+> "="
      <+> "hsep" <+> brackets (hsep (rhsToPrint arhs)))
    where
      lcons = parens $ head (printArgs arhs) <> ":" <> head (tail (printArgs arhs))
  LSg -> (p, "prt" <+> "_" <+> brackets (hsep (printArgs arhs))
      <+> "=" <+> "hsep" <+> brackets (hsep (rhsToPrint arhs)))

rhsToPrint :: ARHS -> [Doc ()]
rhsToPrint items =
  if null items
  then [ "emptyDoc" ]
  else punctuate comma $ prtItem <$> zip anns itemsWithPrec
  where
    prtItem :: (String,((String, Integer), Bool)) -> Doc ()
    prtItem (a, ((s, p), b))
      | null a = if b then prtNT else prtT
      | b = fromString a <+> parens prtNT
      | otherwise = fromString a <+> parens prtT
      where
        prtNT = "prt" <+> (fromString . show) p <+> fromString s
        prtT  = "fromString" <+> fromString s
    itemsWithPrec = zip (indexVars items) (isNTerminal <$> items)
    anns = annotations items