{-# LANGUAGE OverloadedStrings #-}

module BNFC.Backend.Agda.AbstractSyntax where

import BNFC.CF
import BNFC.Options.GlobalOptions
import BNFC.Prelude

import BNFC.Backend.Agda.IOLib
import BNFC.Backend.Agda.Main
import BNFC.Backend.Agda.Options
import BNFC.Backend.Agda.State
import BNFC.Backend.Agda.Utilities.ReservedWords
import BNFC.Backend.Agda.Utilities.Utils

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

import BNFC.Backend.Haskell.AbstractSyntax (cf2abs)
import BNFC.Backend.Haskell.Options (TokenText (..))
import BNFC.Backend.Haskell.Template (cf2template)
import BNFC.Backend.Haskell.Utilities.Utils

import Control.Monad.State

import           Data.Bifunctor  (second)
import qualified Data.Map as Map
import           Data.List       (intersperse)
import qualified BNFC.Utils.List1 as List1
import           Data.String     (fromString)

import Prettyprinter

import System.FilePath ( takeBaseName )


agdaAbstractSyntax :: LBNF -> State AgdaBackendState Result
agdaAbstractSyntax lbnf = do
  st    <- get
  ioLib <- agdaIOLib
  main  <- agdaMain lbnf
  let cfName      = takeBaseName $ optInput $ globalOpt st
      hsRules     = filterRules $ hsAstRules st
      agdaRules   = agdaAstRules st
      hasData     = not $ null hsRules
      hsFuns      = hsFunctions st
      agdaFuns    = agdaFunctions st
      toks        = tokens st
      tokensNames = fst <$> toks
      inDirectory = inDir $ agdaOpts st
      nSpace      = nameSpace $ agdaOpts st
      funct       = functor $ agdaOpts st
      gen         = generic $ agdaOpts st
      absSyntaxHaskell =
        cf2abs lbnf cfName inDirectory nSpace hsRules hsFuns toks funct gen hasData TextToken
      templateHaskell  =
        cf2template hsRules tokensNames cfName inDirectory nSpace funct
      astAgda = cf2AST lbnf cfName inDirectory nSpace agdaRules agdaFuns toks tokensNames funct
  return $
    [ (mkFilePath inDirectory nSpace cfName "Abs" "hs", absSyntaxHaskell)
    , (mkFilePath inDirectory nSpace cfName "Skel" "hs", templateHaskell)
    , (mkFilePath inDirectory nSpace cfName "AST" "agda", astAgda)
    ]
    ++
    ioLib
    ++
    main

  where

    filterRules :: [(Type, [(Label, ([Type], (Integer, ARHS)))])]
                -> [(Type, [(Label, ([Type], (Integer, ARHS)))])]
    filterRules rules =
      filter
      (\(_,l) -> not (null l))
      (Data.Bifunctor.second (filterLabelsAST fNames) <$> rules)

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

cf2AST :: LBNF
       -> String
       -> Bool
       -> Maybe String
       -> [(Type, [(Label, [Type])])]
       -> [(LabelName,Function)]
       -> [(CatName,TokenDef)]
       -> [CatName]
       -> Bool
       -> String
cf2AST lbnf cfName inDir nameSpace astRules functions toks toksName functor =
  docToString defaultLayoutOptions $
    cf2doc lbnf cfName inDir nameSpace astRules functions toks toksName functor

cf2doc :: LBNF
       -> String
       -> Bool
       -> Maybe String
       -> [(Type, [(Label, [Type])])]
       -> [(LabelName,Function)]
       -> [(CatName,TokenDef)]
       -> [CatName]
       -> Bool
       -> Doc ()
cf2doc lbnf cfName inDir nameSpace astRules functions toks toksNames functor =
  (vsep . intersperse emptyDoc) $ concat
  [ [ "-- File generated by the BNF Converter."
    , "-- Agda bindings for the Haskell abstract syntax data types."
    , "module" <+> fromString agdaASTName <+> "where"
    , imports YesImportNumeric False usesPos functor
    ]
    , Utils.when usesString [ "String = #List Char" ]
    , [ importPragmas usesPos
        [ "qualified " ++ hsAbsName, hsPrinterName ++ " (printTree)" ] ]
    , Utils.when usesPos [ intAndPair ]
    , Utils.when functor [ bnfcPosition ]
    , Utils.when hasTokens [ prTokens hsAbsName toks ]
    , Utils.when hasData [ prDatas functor hsAbsName NamedArg astRules ]
    , Utils.when hasFunctions [ prFunctions functor functions ]
    , [ printerBindings hsAbsName astCats usedBuiltins toksNames ]
  ]

  where

    intAndPair :: Doc ()
    intAndPair = vsep
      [ "postulate"
      , indent 2 "#Int      : Set"
      , indent 2 "#intToNat : #Int → #Nat"
      , indent 2 "#natToInt : #Nat → #Int"
      , emptyDoc
      , "{-# COMPILE GHC #Int      = type Prelude.Int    #-}"
      , "{-# COMPILE GHC #intToNat = Prelude.toInteger   #-}"
      , "{-# COMPILE GHC #natToInt = Prelude.fromInteger #-}"
      , emptyDoc
      , "data #Pair (A B : Set) : Set where"
      , indent 2 "#pair : A → B → #Pair A B"
      , emptyDoc
      , "{-# COMPILE GHC #Pair = data (,) ((,)) #-}"
      ]

    bnfcPosition :: Doc ()
    bnfcPosition = "#BNFCPosition = #Maybe (#Pair #Int #Int)"

    agdaASTName :: ModuleName
    agdaASTName = mkModule inDir nameSpace cfName "AST"

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

    hsPrinterName :: ModuleName
    hsPrinterName = mkModule inDir nameSpace cfName "Print"

    usesPos :: Bool
    usesPos = hasPosTokens || (hasData && functor)

    hasPosTokens :: Bool
    hasPosTokens = any isPositionToken (_lbnfTokenDefs lbnf)

    usesString :: Bool
    usesString = BString `elem` Map.keys (_lbnfASTBuiltins lbnf)

    astCats :: [Type]
    astCats = fst <$> astRules

    usedBuiltins :: [BuiltinCat]
    usedBuiltins = Map.keys $ _lbnfASTBuiltins lbnf

    hasData :: Bool
    hasData = not $ null astRules

    hasTokens :: Bool
    hasTokens = not $ null toks

    hasFunctions :: Bool
    hasFunctions = not $ null functions


-- | Print Agda types for user defined tokens.

prTokens :: ModuleName -> [(CatName,TokenDef)] -> Doc ()
prTokens hsAbsName = vsep . intersperse emptyDoc . fmap (prToken hsAbsName)

prToken :: ModuleName -> (CatName,TokenDef) -> Doc ()
prToken hsAbsName (cName, def) = vsep
    [ "data" <+> tokenNameAgda <+> ": Set where"
    , if isPosToken def
      then indent 2 $ tokenConstructorName <+> ": #Pair (#Pair #Int #Int) #String →" <+>
        tokenNameAgda
      else indent 2 $ tokenConstructorName <+> ": #String →" <+> tokenNameAgda
    , emptyDoc
    , "{-# COMPILE GHC" <+> tokenNameHs <+> "=" <+> "data" <+> qualifiedToken <+>
       parens qualifiedToken <+> "#-}"
    ]

  where

    tokenNameAgda :: Doc ()
    tokenNameAgda = fromString $ toList $ avoidAgdaReservedWords cName

    tokenNameHs :: Doc ()
    tokenNameHs = fromString $ toList cName

    tokenConstructorName :: Doc ()
    tokenConstructorName = fromString $ toList $ agdaLower cName

    qualifiedToken :: Doc ()
    qualifiedToken = fromString hsAbsName <> dot <> tokenNameHs

-- | Print Agda types for categories in AST rules.

prDatas :: Bool
        -> ModuleName
        -> ConstructorStyle
        -> [(Type, [(Label, [Type])])]
        -> Doc ()
prDatas functor absName style astRules = vsep
  [ "mutual"
  , emptyDoc
  , indent 2 . vsep . intersperse emptyDoc $
    fmap (prData functor absName style) astRules
  ]

-- | Pretty-print Agda data types and pragmas for AST.

prData :: Bool
       -> ModuleName
       -> ConstructorStyle
       -> (Type, [(Label, [Type])])
       -> Doc ()
prData True absName style (t, labelsItems) = vsep
  [ tName <+> "=" <+> hsDataName <+> "#BNFCPosition"
  , prData' absName style agdaDataName hsDataName labelsItemsWithParam
  ]
  where
    -- Add param to rhs non terminals.
    labelsItemsWithParam :: [(Label, [Type])]
    labelsItemsWithParam = second (addParam <$>) <$> labelsItems

    agdaDataName :: Doc ()
    agdaDataName = fromString (printTypeName $ addParam $ avoidResWordsType t )

    hsDataName :: Doc ()
    hsDataName = fromString (printTypeName t ) <> "'"

    tName :: Doc ()
    tName = fromString ( sanitize $ printTypeName t )

    addParam :: Type -> Type
    addParam (BaseType b)  = BaseType $ addP b
    addParam (ListType tt) = addParam tt

    addP :: BaseCat -> BaseCat
    addP = \case
      BuiltinCat b -> BuiltinCat b
      IdentCat   i -> IdentCat i
      TokenCat   x -> TokenCat x
      BaseCat    x -> BaseCat $ x <> "' Pos#"

prData False absName style (t, labelsItems) =
  prData' absName style agdaDataName hsDataName labelsItems

  where

    hsDataName :: Doc ()
    hsDataName = fromString $ printTypeName t

    agdaDataName :: Doc ()
    agdaDataName = fromString $ printTypeName $ avoidResWordsType t

-- | Pretty-print Agda data types and pragmas.

prData' :: ModuleName
        -> ConstructorStyle
        -> Doc () -- Agda data name
        -> Doc () -- Haskell data name
        -> [(Label, [Type])]
        -> Doc ()
prData' absName style agdaDataName hsDataName labelsItems = vsep
  [ prettyData style agdaDataName labelsItems
  , emptyDoc
  , compilePragma absName agdaDataName hsDataName labels
  ]

  where

    labels :: [Doc ()]
    labels = fromString . printLabelName . fst <$> labelsItems


-- | Pretty-print AST definition in Agda syntax.

prettyData :: ConstructorStyle -> Doc () -> [(Label, [Type])] -> Doc ()
prettyData style typeName labelsItems = hang 2 $ vsep
  [ "data" <+> typeName <+> colon <+> "Set where"
  ,  vsep $ printCase style typeName <$> labelsItems
  ]

printCase :: ConstructorStyle -> Doc () -> (Label, [Type]) -> Doc ()
printCase style typeName (label, args) =
  if null args
  then constructor <+> colon <+> typeName
  else constructor <+> colon <+> printConstructorArgs style args <+>
    fromString uArrow <+> typeName
  where
    constructor :: Doc ()
    constructor = fromString $ printLabelName $ lowerLabel label


printConstructorArgs :: ConstructorStyle -> [Type] -> Doc ()
printConstructorArgs style args =
  case style of
    UnnamedArg -> hsep $ map fromString $ intersperse uArrow ts
    NamedArg   -> hsep $ map (\ (x :| xs, t) -> parens (hsep [x, hsep xs, colon, t])) tel
  where
  -- Arguments types.
  ts :: [String]
  ts  = map prettyType args
  -- Arguments names.
  ns :: [String]
  ns  = map subscript $ numberUniquely $ map nameSuggestion args
  -- Arguments names grouped by type.
  tel :: [(List1 (Doc ()), Doc ())]
  tel = bimap (fromString <$>) fromString <$> aggregateOn snd (zip ns ts)
  deltaSubscript = ord '₀' - ord '0' -- exploiting that '0' comes before '₀' in character table
  subscript (m, s) = maybe s (\ n -> s ++ map (chr . (deltaSubscript +) . ord) (show n)) m
  -- Aggregate consecutive arguments of the same type.
  aggregateOn :: Eq c => ((a,b) -> c) -> [(a,b)] -> [(List1 a,b)]
  aggregateOn f
    = map (\ p -> (List1.map fst p, snd (List1.head p)))
    . List1.groupBy ((==) `on` f)

prettyType :: Type -> String
prettyType = \case
  BaseType bc -> printBaseCatName bc
  ListType ty -> "#List " ++ prettyType ty

compilePragma :: ModuleName -> Doc () -> Doc () -> [Doc ()] -> Doc ()
compilePragma absName agdaTypeName hsTypeName labels = vsep
  [ "{-# COMPILE GHC" <+> agdaTypeName <+> "= data" <+>
      fromString absName <> dot <> hsTypeName
  , vsep $ zipWith
      (\ a b -> a <+> fromString absName <> dot <> b)
      ("(" : repeat "|")
      labels
  , ") #-}"
  ]

-- | Generate Haskell/Agda code for the @define@d constructors.
prFunctions :: Bool -> [(LabelName,Function)] -> Doc ()
prFunctions functor = vsep . intersperse emptyDoc . map (prFunction functor)

prFunction :: Bool -> (LabelName,Function) -> Doc ()
prFunction functor (lName, fun) = vsep [ header, withBody ]

  where

    funName :: Doc ()
    funName = fromString $ toList lName
    arrow :: Doc ()
    arrow = fromString uArrow
    header :: Doc ()
    header = funName <+> ":" <+> paramsTypes <+> arrow <+> returnType
    paramsTypes :: Doc ()
    paramsTypes =
      if functor
      then "{A : Set} a" <+> arrow <+> hsep (intersperse arrow (paramT <$> types))
      else hsep (intersperse arrow (paramT <$> types))
    paramT :: Type -> Doc ()
    paramT t =
      if functor && not (isBuiltinType t)
      then (fromString . printTypeName) t <> "' a"
      else (fromString . printTypeName) t
    types :: [Type]
    types = paramType <$> funPars fun
    returnType :: Doc ()
    returnType =
      if functor
      then (fromString . printTypeName . funType) fun <> "' a"
      else (fromString . printTypeName . funType) fun
    paramsNames :: [Doc ()]
    paramsNames = fromString . toList . paramName <$> funPars fun
    withBody :: Doc ()
    withBody =
      funName <+> args <+> "=" <+>
      (fromString . printExp functor functorParam . funBody) fun
    args :: Doc ()
    args =
      if functor
      then hsep (fromString functorParam : paramsNames)
      else hsep paramsNames
    functorParam :: String
    functorParam = mkFunctorParam "a"
    mkFunctorParam :: String -> String
    mkFunctorParam a =
      if a `notElem` (l : paramNames)
      then a
      else mkFunctorParam (a ++ "'")
      where
        l :: String
        l = toList lName
        paramNames :: [String]
        paramNames = toList . paramName <$> funPars fun

printerBindings :: ModuleName   -- Abs module name.
                -> [Type]       -- AST category names.
                -> [BuiltinCat] -- Used builtins.
                -> [CatName]    -- User defined tokens names.
                -> Doc ()
printerBindings absName astCats builtins toks = vsep
  [ "-- Binding the pretty printers."
  , prPrinterBindings $ zip agdaPrinters agdaTypes
  , emptyDoc
  , prPrinterPragmas $ zip agdaPrinters hsTypes
  ]
  where

    agdaPrinters :: [Doc ()]
    agdaPrinters =
      map (\t -> "print" <> fromString (identType $ avoidResWordsType t)) astCats
      ++
      map (\t -> "print" <> fromString (toList $ avoidAgdaReservedWords t)) toks
      ++
      map (\b -> "print" <> builtinDoc b) builtins

    agdaTypes :: [Doc ()]
    agdaTypes =
      map (fromString . prettyType . avoidResWordsType) astCats
      ++
      map (fromString . toList . avoidAgdaReservedWords) toks
      ++
      map builtinDoc builtins

    -- hs types with names suggestions
    hsTypes :: [(Doc (), Doc())]
    hsTypes =
      map
      (\t ->
        if isListType t
        then
          ( brackets $ fromString $ (++) (absName ++ ".") $ printTypeName t
          , fromString $ nameSuggestion t
          )
        else
          ( fromString $ (++) (absName ++ ".") $ printTypeName t
          , fromString $ nameSuggestion t
          )
      )
      astCats
      ++
      map (\t -> (fromString $ (++) (absName ++ ".") $ toList t, "x")) toks
      ++
      map (\b -> (builtinDoc b, "x")) builtins

    builtinDoc :: BuiltinCat -> Doc ()
    builtinDoc b = fromString $ toList $ printBuiltinCat b

prPrinterBindings :: [(Doc (), Doc ())] -> Doc ()
prPrinterBindings names = vsep
  [ "postulate"
  , indent 2 $ vsep $ map prPrinterBinding names
  ]

prPrinterBinding :: (Doc (), Doc ()) -> Doc ()
prPrinterBinding (agdaPrinter, agdaType) =
  agdaPrinter <+> colon <+> agdaType <+> fromString uArrow <+> "#String"

prPrinterPragmas :: [(Doc (), (Doc (), Doc()))] -> Doc ()
prPrinterPragmas names = vsep $ map prPrinterPragma names

prPrinterPragma :: (Doc (), (Doc (), Doc())) -> Doc ()
prPrinterPragma (agdaPrinter, (hsType, varName)) = hsep
  [ "{-# COMPILE GHC"
  , agdaPrinter <+> "="
  , backslash <+> varName <+> "->" <+> "Data.Text.pack"
  , parens $ "printTree" <+> parens (varName <+> "::" <+> hsType)
  , "#-}"]