{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module Trans.HaskellToText (moduleToText) where

import Data.List (intersperse)
import Data.String (fromString)
import qualified Data.Text as T
import qualified Data.Text.Lazy as LT
import qualified IR.Common as C
import IR.Haskell
import qualified IR.Name as Name
import Text.PrettyPrint.Leijen.Text (hcat, vcat)
import qualified Text.PrettyPrint.Leijen.Text as PP
import Trans.ToHaskellCommon (eStd_)

moduleToText :: Module -> LT.Text
moduleToText mod = PP.displayT $ PP.renderPretty 0.4 80 (format mod)

-- | Types which can be rendered as haskell source code.
--
-- Note well: while format returns a Doc, it is not safe to render it using
-- wl-pprint's "compact" output; we rely on newline significance in some ways
-- without enforcing it.
--
-- It would be nice to fix this, but given that we don't currently expose this
-- via the library and only ever render it in the one place (in main), it isn't
-- a huge priority.
class Format a where
  format :: a -> PP.Doc

instance Format Module where
  format Module {modName, modLangPragmas, modDecls, modExports, modImports} =
    vcat
      [ vcat ["{-# LANGUAGE " <> PP.textStrict ext <> " #-}" | ext <- modLangPragmas],
        -- The generated code sometimes triggers these warnings, but they're nothing to
        -- worry about.
        --
        -- Dodgy-exports comes up when we don't actaully generate any definitions to
        -- export, in which case GHC complains in the machine readable module alias that
        -- we're trying to re-export everything from a module with no-exports. This
        -- happens with e.g. c++.capnp, which only defines annotations.
        "{-# OPTIONS_GHC -Wno-unused-imports #-}",
        "{-# OPTIONS_GHC -Wno-dodgy-exports #-}",
        "{-# OPTIONS_GHC -Wno-unused-matches #-}",
        "{-# OPTIONS_GHC -Wno-orphans #-}",
        "{-# OPTIONS_GHC -Wno-unticked-promoted-constructors #-}",
        "{-# OPTIONS_GHC -Wno-name-shadowing #-}",
        hcat
          [ "module ",
            PP.textStrict $ mconcat $ intersperse "." $ map Name.renderUnQ modName,
            case modExports of
              Nothing -> ""
              Just exports ->
                PP.tupled (map format exports),
            " where"
          ],
        vcat $ map format modImports,
        -- We import many things, including the prelude, qualified under the
        -- "Std_" namespace, so that they don't collide with names in the
        -- generated code; see issue #58.
        vcat $
          map
            format
            [ ImportAs {importAs = "Std_", parts = ["Prelude"]},
              ImportAs {importAs = "Std_", parts = ["Data", "Word"]},
              ImportAs {importAs = "Std_", parts = ["Data", "Int"]}
            ],
        -- ...but there are a couple operators we still want unqaulified:
        "import Prelude ((<$>), (<*>), (>>=))",
        vcat $ map format modDecls
      ]

instance Format Export where
  format (ExportMod parts) =
    "module " <> PP.textStrict (mconcat $ intersperse "." $ map Name.renderUnQ parts)
  format (ExportLCtors name) = format name <> "(..)"
  format (ExportGCtors name) = format name <> "(..)"
  format (ExportGName name) = format name
  format (ExportLName name) = format name

instance Format Decl where
  format (DcData d) = format d
  format DcValue {typ, def = def@DfValue {name}} =
    vcat
      [ hcat [format name, " :: ", format typ],
        format def
      ]
  format DcInstance {ctx, typ, defs} =
    hcat
      [ "instance ",
        format (TCtx ctx typ),
        whereBlock defs
      ]
  format (DcTypeInstance alias orig) =
    hcat
      [ "type instance ",
        format alias,
        " = ",
        format orig
      ]
  format (DcDeriveInstance ctx typ) =
    hcat
      [ "deriving instance ",
        format $ TCtx ctx typ
      ]
  format DcClass {ctx, name, params, funDeps, decls} =
    hcat
      [ "class ",
        format $
          TCtx ctx $
            TApp
              (TLName name)
              [TLName (Name.unQToLocal p) | p <- params],
        formatFunDeps funDeps,
        whereBlock decls
      ]

formatFunDeps :: [(T.Text, T.Text)] -> PP.Doc
formatFunDeps [] = ""
formatFunDeps deps =
  mconcat
    [ " | ",
      mconcat $
        intersperse
          ", "
          [ PP.textStrict a <> " -> " <> PP.textStrict b
            | (a, b) <- deps
          ]
    ]

whereBlock :: Format a => [a] -> PP.Doc
whereBlock [] = ""
whereBlock xs =
  vcat
    [ " where",
      indent $ vcat $ map format xs
    ]

instance Format ClassDecl where
  format (CdValueDecl name typ) = hcat [format name, " :: ", format typ]
  format (CdValueDef def) = format def
  format (CdMinimal names) =
    hcat
      [ "{-# MINIMAL ",
        hcat $ PP.punctuate PP.comma (map format names),
        " #-}"
      ]

instance Format InstanceDef where
  format (IdData d) = format d
  format (IdValue d) = format d
  format (IdType d) = format d

instance Format TypeAlias where
  format (TypeAlias name params value) =
    hcat
      [ "type ",
        format name,
        " ",
        mconcat $ intersperse " " (map format params),
        " = ",
        format value
      ]

instance Format DataDecl where
  format Data {dataName, typeArgs, dataVariants, dataInstance, dataNewtype, derives} =
    vcat
      [ hcat
          [ if dataNewtype
              then "newtype "
              else "data ",
            if dataInstance
              then "instance "
              else "",
            format dataName,
            " ",
            mconcat $ intersperse " " $ map format typeArgs
          ],
        indent $
          vcat
            [ case dataVariants of
                (d : ds) -> vcat $ ("= " <> format d) : map (("| " <>) . format) ds
                [] -> "",
              formatDerives derives
            ]
      ]

instance Format ValueDef where
  format DfValue {name, value, params} =
    hcat
      [ format name,
        " ",
        hcat $ intersperse " " $ map format params,
        " = ",
        format value
      ]

instance Format Exp where
  format (EApp e []) = format e
  format (EApp e es) =
    hcat
      [ "(",
        hcat $ intersperse " " $ map format (e : es),
        ")"
      ]
  format (EFApp e []) = format (EApp (eStd_ "pure") [e])
  format (EFApp e es) =
    hcat
      [ "(",
        format e,
        PP.encloseSep " <$> " ")" " <*> " $ map format es
      ]
  format (EGName e) = format e
  format (ELName e) = format e
  format (EVar name) = PP.textStrict name
  format (EInt n) = fromString (show n)
  format (EDo ds ex) =
    vcat
      [ "(do",
        indent $ vcat (map format ds ++ [format ex, ")"])
      ]
  format (EBind x f) = PP.parens (format x <> " >>= " <> format f)
  format (ETup es) = PP.tupled (map format es)
  format (EList es) = PP.list (map format es)
  format (ECase ex arms) =
    vcat
      [ hcat ["case ", format ex, " of"],
        indent $
          vcat
            [ vcat
                [ hcat [format p, " ->"],
                  indent (format e)
                ]
              | (p, e) <- arms
            ]
      ]
  format (ETypeAnno e ty) =
    PP.parens $
      hcat
        [ format e,
          " :: ",
          format ty
        ]
  format (EBytes bytes) = fromString (show bytes)
  format (ELambda params body) =
    PP.parens $
      hcat
        [ "\\",
          hcat (PP.punctuate " " (map format params)),
          " -> ",
          format body
        ]
  format (ERecord old updates) =
    format old
      <> PP.encloseSep
        "{"
        "}"
        ","
        [ hcat [format name, " = ", format value]
          | (name, value) <- updates
        ]
  format (ETypeApp e ts) =
    hcat ("(" : format e : [" @(" <> format t <> ")" | t <- ts] ++ [")"])
  format (ELabel name) = "#" <> format name

instance Format Do where
  format (DoBind var ex) = format var <> " <- " <> format ex
  format (DoE ex) = format ex

instance Format Pattern where
  format (PVar v) = PP.textStrict v
  format (PLCtor c ps) =
    hcat
      [ "(",
        mconcat $ intersperse " " (format c : map format ps),
        ")"
      ]
  format (PGCtor c ps) =
    hcat
      [ "(",
        mconcat $ intersperse " " (format c : map format ps),
        ")"
      ]
  format (PInt n) = fromString (show n)
  format (PLRecordWildCard name) = format name <> "{..}"

formatDerives :: [Name.UnQ] -> PP.Doc
formatDerives [] = ""
formatDerives ds = "deriving" <> PP.tupled (map format ds)

instance Format DataVariant where
  format DataVariant {dvCtorName, dvArgs} =
    hcat [format dvCtorName, " ", format dvArgs]

instance Format DataArgs where
  format (APos types) =
    mconcat $ intersperse " " (map format types)
  format (ARec fields) =
    PP.line
      <> indent
        ( PP.encloseSep
            "{"
            "}"
            ","
            [format name <> " :: " <> format ty | (name, ty) <- fields]
        )

instance Format Type where
  format (TGName ty) = format ty
  format (TLName ty) = format ty
  format (TVar txt) = PP.textStrict txt
  format (TApp t []) = format t
  format (TApp f xs) =
    "(" <> mconcat (intersperse " " $ map format (f : xs)) <> ")"
  format (TFn types) =
    mconcat $ intersperse " -> " $ map format types
  format (TCtx [] ty) = format ty
  format (TCtx constraints ty) =
    PP.tupled (map format constraints) <> " => " <> format ty
  format (TPrim ty) = format ty
  format TUnit = "()"
  format (TKindAnnotated ty kind) =
    "(" <> format ty <> " :: " <> format kind <> ")"
  format (TString str) = fromString $ show str

instance Format Name.GlobalQ where
  format Name.GlobalQ {local, globalNS = Name.NS parts} =
    mconcat
      [ mconcat $ intersperse "." (map PP.textStrict parts),
        ".",
        format local
      ]

instance Format Name.LocalQ where
  format = PP.textStrict . Name.renderLocalQ

instance Format Name.UnQ where
  format = PP.textStrict . Name.renderUnQ

instance Format Import where
  format ImportAs {importAs, parts} =
    hcat
      [ "import qualified ",
        implodeNS parts,
        " as ",
        format importAs
      ]
  format ImportQual {parts} = hcat ["import qualified ", implodeNS parts]
  format ImportAll {parts} = hcat ["import ", implodeNS parts]

implodeNS :: Format a => [a] -> PP.Doc
implodeNS parts = mconcat $ intersperse "." (map format parts)

instance Format C.PrimWord where
  format C.PrimBool = "Std_.Bool"
  format (C.PrimInt (C.IntType sign sz)) =
    let szDoc = fromString $ show $ C.sizeBits sz
        typePrefix = case sign of
          C.Signed -> "Int"
          C.Unsigned -> "Word"
     in "Std_." <> typePrefix <> szDoc
  format C.PrimFloat32 = "Std_.Float"
  format C.PrimFloat64 = "Std_.Double"

-- | Indent the argument by four spaces.
indent :: PP.Doc -> PP.Doc
indent = PP.indent 4