{-# LANGUAGE DataKinds #-}

module Unison.LSP.EditDefinition
  ( editDefinitionHandler,
  )
where

import Control.Lens hiding (List)
import Control.Monad.Except
import Control.Monad.Reader
import Data.Aeson qualified as Aeson
import Data.Text qualified as Text
import Language.LSP.Protocol.Lens qualified as LSP
import Language.LSP.Protocol.Message qualified as Msg
import Language.LSP.Protocol.Types
import Language.LSP.Protocol.Types qualified as LSP
import Unison.LSP.FileAnalysis (ppedForFile)
import Unison.LSP.Queries qualified as LSPQ
import Unison.LSP.Types
import Unison.LSP.Util.Wrappers (editDefinitionByFQN)
import Unison.LabeledDependency qualified as LD
import Unison.Prelude
import Unison.PrettyPrintEnv qualified as PPE
import Unison.PrettyPrintEnvDecl qualified as PPED
import Unison.Syntax.HashQualified qualified as SyntaxHQ

data EditDefinitionParams = EditDefinitionParams
  { EditDefinitionParams -> TextDocumentIdentifier
textDocument :: LSP.TextDocumentIdentifier,
    EditDefinitionParams -> Maybe Position
position :: Maybe LSP.Position,
    EditDefinitionParams -> Maybe Text
fqn :: Maybe Text
  }
  deriving (Int -> EditDefinitionParams -> ShowS
[EditDefinitionParams] -> ShowS
EditDefinitionParams -> String
(Int -> EditDefinitionParams -> ShowS)
-> (EditDefinitionParams -> String)
-> ([EditDefinitionParams] -> ShowS)
-> Show EditDefinitionParams
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EditDefinitionParams -> ShowS
showsPrec :: Int -> EditDefinitionParams -> ShowS
$cshow :: EditDefinitionParams -> String
show :: EditDefinitionParams -> String
$cshowList :: [EditDefinitionParams] -> ShowS
showList :: [EditDefinitionParams] -> ShowS
Show, EditDefinitionParams -> EditDefinitionParams -> Bool
(EditDefinitionParams -> EditDefinitionParams -> Bool)
-> (EditDefinitionParams -> EditDefinitionParams -> Bool)
-> Eq EditDefinitionParams
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EditDefinitionParams -> EditDefinitionParams -> Bool
== :: EditDefinitionParams -> EditDefinitionParams -> Bool
$c/= :: EditDefinitionParams -> EditDefinitionParams -> Bool
/= :: EditDefinitionParams -> EditDefinitionParams -> Bool
Eq)

instance Aeson.FromJSON EditDefinitionParams where
  parseJSON :: Value -> Parser EditDefinitionParams
parseJSON = String
-> (Object -> Parser EditDefinitionParams)
-> Value
-> Parser EditDefinitionParams
forall a. String -> (Object -> Parser a) -> Value -> Parser a
Aeson.withObject String
"EditDefinitionParams" ((Object -> Parser EditDefinitionParams)
 -> Value -> Parser EditDefinitionParams)
-> (Object -> Parser EditDefinitionParams)
-> Value
-> Parser EditDefinitionParams
forall a b. (a -> b) -> a -> b
$ \Object
v -> do
    TextDocumentIdentifier
textDocument <- Object
v Object -> Key -> Parser TextDocumentIdentifier
forall a. FromJSON a => Object -> Key -> Parser a
Aeson..: Key
"textDocument"
    Maybe Position
position <- Object
v Object -> Key -> Parser (Maybe Position)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
Aeson..:? Key
"position"
    Maybe Text
fqn <- Object
v Object -> Key -> Parser (Maybe Text)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
Aeson..:? Key
"fqn"
    EditDefinitionParams -> Parser EditDefinitionParams
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure EditDefinitionParams {TextDocumentIdentifier
$sel:textDocument:EditDefinitionParams :: TextDocumentIdentifier
textDocument :: TextDocumentIdentifier
textDocument, Maybe Position
$sel:position:EditDefinitionParams :: Maybe Position
position :: Maybe Position
position, Maybe Text
$sel:fqn:EditDefinitionParams :: Maybe Text
fqn :: Maybe Text
fqn}

data EditDefinitionResponse
  = EditDefinitionError Text
  | EditDefinitionSuccess Bool {- Whether the definition was newly added -}
  deriving (Int -> EditDefinitionResponse -> ShowS
[EditDefinitionResponse] -> ShowS
EditDefinitionResponse -> String
(Int -> EditDefinitionResponse -> ShowS)
-> (EditDefinitionResponse -> String)
-> ([EditDefinitionResponse] -> ShowS)
-> Show EditDefinitionResponse
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EditDefinitionResponse -> ShowS
showsPrec :: Int -> EditDefinitionResponse -> ShowS
$cshow :: EditDefinitionResponse -> String
show :: EditDefinitionResponse -> String
$cshowList :: [EditDefinitionResponse] -> ShowS
showList :: [EditDefinitionResponse] -> ShowS
Show, EditDefinitionResponse -> EditDefinitionResponse -> Bool
(EditDefinitionResponse -> EditDefinitionResponse -> Bool)
-> (EditDefinitionResponse -> EditDefinitionResponse -> Bool)
-> Eq EditDefinitionResponse
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EditDefinitionResponse -> EditDefinitionResponse -> Bool
== :: EditDefinitionResponse -> EditDefinitionResponse -> Bool
$c/= :: EditDefinitionResponse -> EditDefinitionResponse -> Bool
/= :: EditDefinitionResponse -> EditDefinitionResponse -> Bool
Eq)

instance Aeson.ToJSON EditDefinitionResponse where
  toJSON :: EditDefinitionResponse -> Value
toJSON = \case
    EditDefinitionSuccess Bool
newlyAdded ->
      [Pair] -> Value
Aeson.object
        [ Key
"newlyAdded" Key -> Bool -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
Aeson..= Bool
newlyAdded
        ]
    EditDefinitionError Text
err ->
      [Pair] -> Value
Aeson.object
        [ Key
"error" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
Aeson..= Text
err
        ]

-- | Handler for the 'unison/editDefinition' custom LSP request.
-- This resolves the symbol at the given position to its FQN and adds it to the current file.
editDefinitionHandler ::
  Msg.TRequestMessage ('Msg.Method_CustomMethod "unison/editDefinition") ->
  (Either Msg.ResponseError Aeson.Value -> Lsp ()) ->
  Lsp ()
editDefinitionHandler :: forall {f :: MessageDirection}.
TRequestMessage ('Method_CustomMethod "unison/editDefinition")
-> (Either ResponseError Value -> Lsp ()) -> Lsp ()
editDefinitionHandler TRequestMessage ('Method_CustomMethod "unison/editDefinition")
m Either ResponseError Value -> Lsp ()
respond = do
  Either Text Bool
result <- ExceptT Text Lsp Bool -> Lsp (Either Text Bool)
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (ExceptT Text Lsp Bool -> Lsp (Either Text Bool))
-> ExceptT Text Lsp Bool -> Lsp (Either Text Bool)
forall a b. (a -> b) -> a -> b
$ do
    let paramsJSON :: Value
paramsJSON = TRequestMessage ('Method_CustomMethod "unison/editDefinition")
m TRequestMessage ('Method_CustomMethod "unison/editDefinition")
-> Getting
     Value
     (TRequestMessage ('Method_CustomMethod "unison/editDefinition"))
     Value
-> Value
forall s a. s -> Getting a s a -> a
^. Getting
  Value
  (TRequestMessage ('Method_CustomMethod "unison/editDefinition"))
  Value
forall s a. HasParams s a => Lens' s a
Lens'
  (TRequestMessage ('Method_CustomMethod "unison/editDefinition"))
  Value
LSP.params
    EditDefinitionParams {TextDocumentIdentifier
$sel:textDocument:EditDefinitionParams :: EditDefinitionParams -> TextDocumentIdentifier
textDocument :: TextDocumentIdentifier
textDocument, Maybe Position
$sel:position:EditDefinitionParams :: EditDefinitionParams -> Maybe Position
position :: Maybe Position
position, Maybe Text
$sel:fqn:EditDefinitionParams :: EditDefinitionParams -> Maybe Text
fqn :: Maybe Text
fqn} <- case Value -> Result EditDefinitionParams
forall a. FromJSON a => Value -> Result a
Aeson.fromJSON Value
paramsJSON of
      Aeson.Error String
err -> Text -> ExceptT Text Lsp EditDefinitionParams
forall a. Text -> ExceptT Text Lsp a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text Lsp EditDefinitionParams)
-> Text -> ExceptT Text Lsp EditDefinitionParams
forall a b. (a -> b) -> a -> b
$ Text
"Invalid parameters: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
Text.pack String
err
      Aeson.Success EditDefinitionParams
p -> EditDefinitionParams -> ExceptT Text Lsp EditDefinitionParams
forall a. a -> ExceptT Text Lsp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure EditDefinitionParams
p

    let fileURI :: Uri
fileURI = TextDocumentIdentifier
textDocument._uri
    -- Get the FQN either directly or by resolving the symbol at the position
    Text
fqnText <- case (Maybe Text
fqn, Maybe Position
position) of
      (Just Text
directFqn, Maybe Position
_) ->
        -- Use the provided FQN directly
        Text -> ExceptT Text Lsp Text
forall a. a -> ExceptT Text Lsp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
directFqn
      (Maybe Text
Nothing, Just Position
pos) -> do
        -- Get the symbol reference at the position
        LabeledDependency
ref <- Text
-> Lsp (Maybe LabeledDependency)
-> ExceptT Text Lsp LabeledDependency
forall a. Text -> Lsp (Maybe a) -> ExceptT Text Lsp a
orFail Text
"Error: Definition not found in Codebase" (Lsp (Maybe LabeledDependency)
 -> ExceptT Text Lsp LabeledDependency)
-> (MaybeT Lsp LabeledDependency -> Lsp (Maybe LabeledDependency))
-> MaybeT Lsp LabeledDependency
-> ExceptT Text Lsp LabeledDependency
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MaybeT Lsp LabeledDependency -> Lsp (Maybe LabeledDependency)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT Lsp LabeledDependency
 -> ExceptT Text Lsp LabeledDependency)
-> MaybeT Lsp LabeledDependency
-> ExceptT Text Lsp LabeledDependency
forall a b. (a -> b) -> a -> b
$ Uri -> Position -> MaybeT Lsp LabeledDependency
forall (m :: * -> *).
Lspish m =>
Uri -> Position -> MaybeT m LabeledDependency
LSPQ.refAtPosition Uri
fileURI Position
pos

        -- Get the FQN for the reference
        PrettyPrintEnvDecl
pped <- Lsp PrettyPrintEnvDecl -> ExceptT Text Lsp PrettyPrintEnvDecl
forall (m :: * -> *) a. Monad m => m a -> ExceptT Text m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Lsp PrettyPrintEnvDecl -> ExceptT Text Lsp PrettyPrintEnvDecl)
-> Lsp PrettyPrintEnvDecl -> ExceptT Text Lsp PrettyPrintEnvDecl
forall a b. (a -> b) -> a -> b
$ Uri -> Lsp PrettyPrintEnvDecl
forall (m :: * -> *). Lspish m => Uri -> m PrettyPrintEnvDecl
ppedForFile Uri
fileURI
        let unsuffixifiedPPE :: PrettyPrintEnv
unsuffixifiedPPE = PrettyPrintEnvDecl -> PrettyPrintEnv
PPED.unsuffixifiedPPE PrettyPrintEnvDecl
pped
        let fqnName :: HashQualified Name
fqnName = case LabeledDependency
ref of
              LD.TypeReference TypeReference
typeRef -> PrettyPrintEnv -> TypeReference -> HashQualified Name
PPE.typeName PrettyPrintEnv
unsuffixifiedPPE TypeReference
typeRef
              LD.TermReferent Referent
termRef -> PrettyPrintEnv -> Referent -> HashQualified Name
PPE.termName PrettyPrintEnv
unsuffixifiedPPE Referent
termRef
        Text -> ExceptT Text Lsp Text
forall a. a -> ExceptT Text Lsp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> ExceptT Text Lsp Text) -> Text -> ExceptT Text Lsp Text
forall a b. (a -> b) -> a -> b
$ HashQualified Name -> Text
SyntaxHQ.toText HashQualified Name
fqnName
      (Maybe Text
Nothing, Maybe Position
Nothing) ->
        Text -> ExceptT Text Lsp Text
forall a. Text -> ExceptT Text Lsp a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError Text
"Either 'position' or 'fqn' must be provided"

    -- Call the editDefinitionByFQN utility
    Uri -> Text -> ExceptT Text Lsp Bool
editDefinitionByFQN Uri
fileURI Text
fqnText

  -- Send the response
  case Either Text Bool
result of
    (Left Text
errMsg) -> Either ResponseError Value -> Lsp ()
respond (Value -> Either ResponseError Value
forall a b. b -> Either a b
Right (Value -> Either ResponseError Value)
-> Value -> Either ResponseError Value
forall a b. (a -> b) -> a -> b
$ EditDefinitionResponse -> Value
forall a. ToJSON a => a -> Value
Aeson.toJSON (EditDefinitionResponse -> Value)
-> EditDefinitionResponse -> Value
forall a b. (a -> b) -> a -> b
$ Text -> EditDefinitionResponse
EditDefinitionError Text
errMsg)
    Right Bool
isNewDefinition ->
      Either ResponseError Value -> Lsp ()
respond (Value -> Either ResponseError Value
forall a b. b -> Either a b
Right (Value -> Either ResponseError Value)
-> Value -> Either ResponseError Value
forall a b. (a -> b) -> a -> b
$ EditDefinitionResponse -> Value
forall a. ToJSON a => a -> Value
Aeson.toJSON (EditDefinitionResponse -> Value)
-> EditDefinitionResponse -> Value
forall a b. (a -> b) -> a -> b
$ Bool -> EditDefinitionResponse
EditDefinitionSuccess Bool
isNewDefinition)
  where
    orFail :: Text -> Lsp (Maybe a) -> ExceptT Text Lsp a
    orFail :: forall a. Text -> Lsp (Maybe a) -> ExceptT Text Lsp a
orFail Text
err Lsp (Maybe a)
action = do
      Maybe a
ma <- Lsp (Maybe a) -> ExceptT Text Lsp (Maybe a)
forall (m :: * -> *) a. Monad m => m a -> ExceptT Text m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift Lsp (Maybe a)
action
      case Maybe a
ma of
        Just a
a -> a -> ExceptT Text Lsp a
forall a. a -> ExceptT Text Lsp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
        Maybe a
Nothing -> Text -> ExceptT Text Lsp a
forall a. Text -> ExceptT Text Lsp a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError Text
err