module Unison.MCP (runOnStdIO, initServer) where

import Network.MCP.Server qualified as MCP
import Network.MCP.Server.StdIO qualified as MCP
import Network.MCP.Types
import Text.RawString.QQ (r)
import Unison.Auth.CredentialManager qualified as AuthN
import Unison.Auth.HTTPClient qualified as AuthN
import Unison.Auth.Tokens qualified as AuthN
import Unison.Codebase (Codebase)
import Unison.Codebase.Runtime (Runtime)
import Unison.MCP.Prompts (prompts)
import Unison.MCP.StaticResources (staticResources)
import Unison.MCP.Tools (tools)
import Unison.MCP.Types
import Unison.MCP.Wrapper qualified as MCPWrapper
import Unison.Parser.Ann (Ann)
import Unison.Prelude
import Unison.Symbol (Symbol)

serverDescription :: Text
serverDescription :: Text
serverDescription =
  Text
[r|
        This server provides tools for interacting with Unison Code locally, such as typechecking or reading documentation, as well as tools for searching Unison Share, which is a platform for sharing Unison projects and libraries.

        It also provides some mechanisms for editing and updating local Unison projects, such as installing libraries from Unison Share.

        Before doing any work in unison please read the file://unison-guide resource for information on how to write unison.
    |]

initServer :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> Maybe FilePath -> Text -> IO MCP.Server
initServer :: Codebase IO Symbol Ann
-> Runtime Symbol
-> Runtime Symbol
-> Maybe FilePath
-> Text
-> IO Server
initServer Codebase IO Symbol Ann
codebase Runtime Symbol
runtime Runtime Symbol
sbRuntime Maybe FilePath
workDir Text
ucmVersion = do
  CredentialManager
credMan <- IO CredentialManager
forall (m :: * -> *). MonadIO m => m CredentialManager
AuthN.newCredentialManager
  let tokenProvider :: AuthN.TokenProvider
      tokenProvider :: TokenProvider
tokenProvider = CredentialManager -> TokenProvider
AuthN.newTokenProvider CredentialManager
credMan
  AuthenticatedHttpClient
authenticatedHTTPClient <- TokenProvider -> Text -> IO AuthenticatedHttpClient
forall (m :: * -> *).
MonadIO m =>
TokenProvider -> Text -> m AuthenticatedHttpClient
AuthN.newAuthenticatedHTTPClient TokenProvider
tokenProvider Text
ucmVersion
  let env :: Env
env =
        Env
          { Codebase IO Symbol Ann
codebase :: Codebase IO Symbol Ann
$sel:codebase:Env :: Codebase IO Symbol Ann
codebase,
            Runtime Symbol
runtime :: Runtime Symbol
$sel:runtime:Env :: Runtime Symbol
runtime,
            Runtime Symbol
sbRuntime :: Runtime Symbol
$sel:sbRuntime:Env :: Runtime Symbol
sbRuntime,
            Text
ucmVersion :: Text
$sel:ucmVersion:Env :: Text
ucmVersion,
            Maybe FilePath
workDir :: Maybe FilePath
$sel:workDir:Env :: Maybe FilePath
workDir,
            AuthenticatedHttpClient
authenticatedHTTPClient :: AuthenticatedHttpClient
$sel:authenticatedHTTPClient:Env :: AuthenticatedHttpClient
authenticatedHTTPClient
          }
  -- Create server
  let serverInfo :: Implementation
serverInfo = Text -> Text -> Implementation
Implementation Text
"unison-mcp" Text
"0.0.1"

  Env -> MCP Server -> IO Server
forall a. Env -> MCP a -> IO a
runMCP Env
env (MCP Server -> IO Server) -> MCP Server -> IO Server
forall a b. (a -> b) -> a -> b
$ Implementation
-> Text
-> StaticResources
-> [Tool MCP]
-> [Prompt MCP]
-> MCP Server
forall (m :: * -> *).
MonadUnliftIO m =>
Implementation
-> Text -> StaticResources -> [Tool m] -> [Prompt m] -> m Server
MCPWrapper.mkServer Implementation
serverInfo Text
serverDescription StaticResources
staticResources [Tool MCP]
tools [Prompt MCP]
prompts

-- | Run the MCP server until we hit EOF.
runOnStdIO :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> FilePath -> Text -> IO ()
runOnStdIO :: Codebase IO Symbol Ann
-> Runtime Symbol -> Runtime Symbol -> FilePath -> Text -> IO ()
runOnStdIO Codebase IO Symbol Ann
codebase Runtime Symbol
runtime Runtime Symbol
sbRuntime FilePath
workDir Text
ucmVersion = do
  Server
server <- Codebase IO Symbol Ann
-> Runtime Symbol
-> Runtime Symbol
-> Maybe FilePath
-> Text
-> IO Server
initServer Codebase IO Symbol Ann
codebase Runtime Symbol
runtime Runtime Symbol
sbRuntime (FilePath -> Maybe FilePath
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FilePath
workDir) Text
ucmVersion
  -- Start the server with StdIO transport
  Server -> IO ()
MCP.runServerWithSTDIO Server
server