{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE RankNTypes      #-}
{-# LANGUAGE TemplateHaskell #-}
module Language.LSP.Protocol.Utils.Misc
  ( rdrop
  , makeSingletonFromJSON
  , makeRegHelper
  , lspOptions
  , lspOptionsUntagged
  , prettyJSON
  , ViaJSON (..)
  ) where

import           Control.Monad
import           Data.Aeson    
import           Data.Aeson.Text     as Aeson
import qualified Data.Foldable       as F
import qualified Data.Foldable.WithIndex as F
import qualified Data.Functor.WithIndex.Instances ()
import           Data.List           hiding (group)
import           Data.Maybe          (mapMaybe)
import           Language.Haskell.TH
import           Prettyprinter

-- ---------------------------------------------------------------------

rdrop :: Int -> [a] -> [a]
rdrop cnt = reverse . drop cnt . reverse

-- | Given a wrapper and a singleton GADT, construct FromJSON
-- instances for each constructor return type by invoking the
-- FromJSON instance for the wrapper and unwrapping
makeSingletonFromJSON :: Name -> Name -> [Name] -> Q [Dec]
makeSingletonFromJSON wrap gadt skip = do
  TyConI (DataD _ _ _ _ cons _) <- reify gadt
  concat <$> (sequence $ mapMaybe (makeInst wrap skip) cons)

{-
instance FromJSON (SMethod $method) where
  parseJSON = parseJSON >=> \case
      SomeMethod $singleton-method -> pure $singleton-method
      _ -> mempty
-}
makeInst :: Name -> [Name] -> Con -> Maybe (Q [Dec])
makeInst _ skip (GadtC [sConstructor] _ _) | sConstructor `elem` skip = Nothing
makeInst wrap _ (GadtC [sConstructor] args t) = Just $ do
  ns <- replicateM (length args) (newName "x")
  let wrappedPat = conP wrap [conP sConstructor (map varP ns)]
      unwrappedE = pure $ foldl' AppE (ConE sConstructor) (map VarE ns)
  [d| instance FromJSON $(pure t) where
        parseJSON = parseJSON >=> \case
          $wrappedPat -> pure $unwrappedE
          _           -> mempty
    |]
makeInst wrap skip (ForallC _ _ con) = makeInst wrap skip con -- Cancel and Custom requests
makeInst _ _ _ = Just $ fail "makeInst only defined for GADT constructors"

makeRegHelper :: Name -> DecsQ
makeRegHelper regOptTypeName = do
  Just sMethodTypeName <- lookupTypeName "SMethod"
  Just fromClientName <- lookupValueName "ClientToServer"
  TyConI (DataD _ _ _ _ allCons _) <- reify sMethodTypeName

  let isConsFromClient (GadtC _ _ (AppT _ method)) = isMethodFromClient method
      isConsFromClient _                           = return False
      isMethodFromClient :: Type -> Q Bool
      isMethodFromClient (PromotedT method) = do
        DataConI _ typ _ <- reify method
        case typ of
          AppT (AppT _ (PromotedT n)) _ -> return $ n == fromClientName
          _                             -> return False
      isMethodFromClient _ = fail "Didn't expect this type of Method!"

  cons <- filterM isConsFromClient allCons

  let conNames = mapMaybe (\case { (GadtC [name] _ _) -> Just name; _ -> Nothing; }) cons
      helperName = mkName "regHelper"
      mkClause name = do
        x <- newName "x"
        clause [ conP name [], varP x ]
               (normalB (varE x))
               []
      regOptTcon = conT regOptTypeName
  fun <- funD helperName (map mkClause conNames)

  typSig <- sigD helperName $
    [t| forall m x. $(conT sMethodTypeName) m
        -> (Show ($regOptTcon m) => ToJSON ($regOptTcon m) => FromJSON ($regOptTcon m) => x)
        -> x |]
  return [typSig, fun]

-- | Standard options for use when generating JSON instances
-- NOTE: This needs to be in a separate file because of the TH stage restriction
lspOptions :: Options
lspOptions = defaultOptions { omitNothingFields = True, fieldLabelModifier = modifier }
  where
  modifier :: String -> String
  -- For fields called data and type in the spec, we call them xdata and xtype
  -- in haskell-lsp-types to avoid it clashing with the Haskell keywords. This
  -- fixes up the json derivation
  modifier "_xdata" = "data"
  modifier "_xtype" = "type"
  modifier xs       = drop 1 xs

-- | Standard options for use when generating JSON instances for an untagged union
lspOptionsUntagged :: Options
lspOptionsUntagged = lspOptions { sumEncoding = UntaggedValue }

prettyJSON :: Value -> Doc ann
prettyJSON = \case
  Array vec ->
    let docs = fmap prettyJSON (F.toList vec)
        separator = ","
    in group $ nest 2 ("[" <> line <> vsep (punctuate separator docs)) <> line <> "]"
  Object km ->
    let docs = fmap (\(k, v) -> pretty (show k) <> ":" <+> prettyJSON v) (F.itoList km)
        separator = ","
    in group $ nest 2 ("{" <> line <> vsep (punctuate separator docs)) <> line <> "}"
  -- for atomic objects, piggyback off aeson's encoding
  v -> pretty $ Aeson.encodeToLazyText v

newtype ViaJSON a = ViaJSON a

instance ToJSON a => Pretty (ViaJSON a) where
  pretty (ViaJSON a) = prettyJSON $ toJSON a
