{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import qualified Data.ByteString as B
import Data.Int (Int32)
import Data.List (sortBy)
import Data.Maybe (isJust)
import Data.ProtoLens (decodeMessage, defMessage, encodeMessage)
import Data.ProtoLens.Compiler.ModuleName (protoModuleName)
import Data.ProtoLens.Labels ()
import qualified Data.Set as Set
import Data.String (fromString)
import qualified Data.Text as T
import Data.Text (Text, intercalate, pack, unpack)
import DynFlags (DynFlags, getDynFlags)
import GHC (runGhc)
import GHC.Paths (libdir)
import GHC.SourceGen
import GHC.SourceGen.Pretty (showPpr)
import GhcMonad (liftIO)
import Lens.Family2
import Proto.Google.Protobuf.Compiler.Plugin
  ( CodeGeneratorRequest,
    CodeGeneratorResponse,
  )
import Proto.Google.Protobuf.Descriptor
  ( DescriptorProto,
    EnumDescriptorProto,
    FieldDescriptorProto,
    FieldDescriptorProto'Label (FieldDescriptorProto'LABEL_REPEATED),
    FieldDescriptorProto'Type (FieldDescriptorProto'TYPE_MESSAGE),
    FileDescriptorProto,
  )
import System.Environment (getProgName)
import System.Exit (ExitCode (..), exitWith)
import qualified System.IO as IO
import Text.Casing (camel)

data ProtoMod
  = ProtoMod
      { modName :: String,
        modTypes :: [ProtoType]
      }
  deriving (Show)

data ProtoType
  = ProtoMsg String DescriptorProto
  | ProtoEnum String
  deriving (Show)

main :: IO ()
main = do
  contents <- B.getContents
  progName <- getProgName
  case decodeMessage contents of
    Left e -> IO.hPutStrLn IO.stderr e >> exitWith (ExitFailure 1)
    Right x -> runGhc (Just libdir) $ do
      dflags <- getDynFlags
      liftIO $ B.putStr $ encodeMessage $
        makeResponse dflags progName x

makeResponse :: DynFlags -> String -> CodeGeneratorRequest -> CodeGeneratorResponse
makeResponse dflags prog req =
  defMessage
    & #file
      .~ [ defMessage
             & #name .~ "Proto/SignableOrphan.hs"
             & #content .~ header <> "\n\n" <> body
         ]
  where
    protoMods :: [ProtoMod]
    protoMods =
      (\x -> ProtoMod (parseModName x) $ parseModTypes x)
        <$> req ^. #protoFile
    imports :: [ImportDecl']
    imports =
      qualified'
        <$> ( [ import' "Universum",
                import' "Data.Signable",
                import' "GHC.List"
              ]
                <> ( protoMods
                       >>= ( \x ->
                               let n = modName x
                                in import' . fromString <$> [n, n <> "_Fields"]
                           )
                   )
            )
    body :: Text
    body =
      pack . showPpr dflags $
        module'
          (Just "Proto.SignableOrphan")
          (Just [])
          imports
          (protoMods >>= mkImpls)
    header :: Text
    header =
      Data.Text.intercalate "\n" $
        [ "{- This file was auto-generated by the "
            <> pack prog
            <> " program. -}",
          "{-# OPTIONS_GHC -fno-warn-orphans #-}",
          "{-# LANGUAGE NoImplicitPrelude #-}"
        ]

parseModName :: FileDescriptorProto -> String
parseModName fd = protoModuleName (T.unpack $ fd ^. #name)

parseModTypes :: FileDescriptorProto -> [ProtoType]
parseModTypes x =
  (parseEnum mempty <$> x ^. #enumType)
    <> ((x ^. #messageType) >>= parseMsg mempty)

parseEnum :: String -> EnumDescriptorProto -> ProtoType
parseEnum ns x = ProtoEnum $ ns <> unpack (x ^. #name)

parseMsg :: String -> DescriptorProto -> [ProtoType]
parseMsg ns0 x =
  ProtoMsg n x
    : (parseEnum ns <$> x ^. #enumType)
    <> ((x ^. #nestedType) >>= parseMsg ns)
  where
    n = ns0 <> unpack (x ^. #name)
    ns = n <> "'"

mkImpls :: ProtoMod -> [HsDecl']
mkImpls x =
  mk (modName x) <$> modTypes x
  where
    mk m = \case
      ProtoMsg t d -> mkMsgImpl m t d
      ProtoEnum t -> mkEnumImpl m t

mkMsgImpl :: String -> String -> DescriptorProto -> HsDecl'
mkMsgImpl m t d =
  instance'
    (var "Data.Signable.Signable" @@ var (fromString $ m <> "." <> t))
    [ funBind "toBinary" $
        match
          []
          ( op
              (var "Universum.mconcat")
              compose
              ( op
                  ( var "Universum.<&>"
                      @@ ( list
                             . (mkMsgChunk m <$>)
                             . sortBy
                               ( \x y ->
                                   compare
                                     (x ^. #number)
                                     (y ^. #number)
                               )
                             $ d ^. #field
                         )
                  )
                  compose
                  (var "Universum.&")
              )
          )
    ]

mkMsgChunk :: String -> FieldDescriptorProto -> HsExpr'
mkMsgChunk m d
  | d ^. #label == FieldDescriptorProto'LABEL_REPEATED =
    rExpr
  | (d ^. #type' == FieldDescriptorProto'TYPE_MESSAGE)
      || (isJust $ d ^. #maybe'oneofIndex) =
    mExpr
  | otherwise =
    expr
  where
    n0 = unReserve . camel . unpack $ d ^. #name
    tag = case safeFromIntegral $ d ^. #number :: Maybe Int32 of
      Just v ->
        var "Data.Signable.toBinary"
          @@ int (fromIntegral v :: Integer) @::@ var "Universum.Int32"
      Nothing -> error "TAG_OVERFLOW"
    rExpr =
      op
        ( var "Data.Signable.ifThenElse"
            @@ (var "GHC.List.null")
            @@ (var "Universum.const" @@ var "Universum.mempty")
            @@ ( op
                   (var "Universum.<>" @@ tag)
                   compose
                   (var "Data.Signable.toBinary")
               )
        )
        compose
        (var "Universum.view" @@ var (fromString $ m <> "_Fields." <> n0))
    mExpr =
      op
        ( var "Universum.maybe"
            @@ (var "Universum.mempty")
            @@ ( op
                   (var "Universum.<>" @@ tag)
                   compose
                   (var "Data.Signable.toBinary")
               )
        )
        compose
        (var "Universum.view" @@ var (fromString $ m <> "_Fields.maybe'" <> n0))
    expr =
      op
        (var "Universum.<>" @@ tag)
        compose
        ( op
            (var "Data.Signable.toBinary")
            compose
            (var "Universum.view" @@ var (fromString $ m <> "_Fields." <> n0))
        )

mkEnumImpl :: String -> String -> HsDecl'
mkEnumImpl m t =
  instance'
    (var "Data.Signable.Signable" @@ var (fromString $ m <> "." <> t))
    [ funBind "toBinary" $
        match
          []
          ( op
              ( (var "Universum.maybe")
                  @@ (var "Universum.error" @@ string "ENUM_OVERFLOW")
                  @@ (var "Data.Signable.toBinary")
              )
              compose
              ( op
                  ( (var "Data.Signable.safeFromIntegral")
                      @::@ ( var "Universum.Int"
                               --> var "Universum.Maybe"
                               @@ var "Universum.Int32"
                           )
                  )
                  compose
                  (var "Universum.fromEnum")
              )
          )
    ]

compose :: RdrNameStr
compose = "Universum.."

unReserve :: String -> String
unReserve x =
  if x `Set.member` reservedKeywords
    then x <> "'"
    else x

-- | A list of reserved keywords that aren't valid as variable names.
reservedKeywords :: Set.Set String
reservedKeywords =
  Set.fromList $
    -- Haskell2010 keywords:
    -- https://www.haskell.org/onlinereport/haskell2010/haskellch2.html#x7-180002.4
    -- We don't include keywords that are allowed to be variable names,
    -- in particular: "as", "forall", and "hiding".
    [ "case",
      "class",
      "data",
      "default",
      "deriving",
      "do",
      "else",
      "foreign",
      "if",
      "import",
      "in",
      "infix",
      "infixl",
      "infixr",
      "instance",
      "let",
      "module",
      "newtype",
      "of",
      "then",
      "type",
      "where"
    ]
      ++ [ "mdo", -- Nonstandard extensions
           -- RecursiveDo
           "rec", -- Arrows, RecursiveDo
           "pattern", -- PatternSynonyms
           "proc" -- Arrows
         ]

safeFromIntegral ::
  forall a b. (Integral a, Integral b, Bounded b) => a -> Maybe b
safeFromIntegral x =
  if (intX >= intMin) && (intX <= intMax)
    then Just $ fromIntegral x
    else Nothing
  where
    intX = fromIntegral x :: Integer
    intMin = fromIntegral (minBound :: b) :: Integer
    intMax = fromIntegral (maxBound :: b) :: Integer