repositories {
    mavenCentral()
}

dependencies {
    implementation("com.agentclientprotocol:acp:0.3.0-SNAPSHOT")
    // Optional extras:
    // implementation("com.agentclientprotocol:acp-ktor-client:0.3.0-SNAPSHOT")
    // implementation("com.agentclientprotocol:acp-ktor-server:0.3.0-SNAPSHOT")
}

import com.agentclientprotocol.agent.Agent
import com.agentclientprotocol.agent.AgentInfo
import com.agentclientprotocol.agent.AgentSession
import com.agentclientprotocol.agent.AgentSupport
import com.agentclientprotocol.agent.clientInfo
import com.agentclientprotocol.client.ClientInfo
import com.agentclientprotocol.client.FileSystemOperations
import com.agentclientprotocol.common.Event
import com.agentclientprotocol.common.SessionParameters
import com.agentclientprotocol.common.remoteSessionOperations
import com.agentclientprotocol.model.AgentCapabilities
import com.agentclientprotocol.model.ContentBlock
import com.agentclientprotocol.model.LATEST_PROTOCOL_VERSION
import com.agentclientprotocol.model.PromptResponse
import com.agentclientprotocol.model.SessionId
import com.agentclientprotocol.model.SessionUpdate
import com.agentclientprotocol.model.StopReason
import com.agentclientprotocol.protocol.Protocol
import com.agentclientprotocol.transport.StdioTransport
import kotlinx.coroutines.currentCoroutineContext
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.JsonElement
import kotlinx.io.asSink
import kotlinx.io.asSource
import kotlinx.io.buffered

// 1. Build a dedicated AgentSession implementation for each connection.
private class TerminalAgentSession(
    override val sessionId: SessionId
) : AgentSession {
    override suspend fun prompt(
        content: List<ContentBlock>,
        _meta: JsonElement?
    ): Flow<Event> = flow {
        // Echo back what the user typed.
        val userText = content.filterIsInstance<ContentBlock.Text>().joinToString(" ") { it.text }
        emit(
            Event.SessionUpdateEvent(
                SessionUpdate.AgentMessageChunk(ContentBlock.Text("Agent heard: $userText"))
            )
        )

        // Optional extension call via the coroutine context.
        val context = currentCoroutineContext()
        val clientCapabilities = context.clientInfo.capabilities
        if (clientCapabilities.fs?.readTextFile == true) {
            val fs = context.remoteSessionOperations(FileSystemOperations)
            val readmeSnippet = fs.fsReadTextFile("README.md").content.take(120)
            emit(
                Event.SessionUpdateEvent(
                    SessionUpdate.AgentMessageChunk(ContentBlock.Text("README preview: $readmeSnippet…"))
                )
            )
        }

        // Finish the turn once updates are sent.
        emit(Event.PromptResponseEvent(PromptResponse(StopReason.END_TURN)))
    }

    override suspend fun cancel() {
        // No long-running work in this demo, so nothing to clean up yet.
    }
}

// 2. Implement AgentSupport: negotiate capabilities and build per-session handlers.
private class TerminalAgentSupport : AgentSupport {
    override suspend fun initialize(clientInfo: ClientInfo) = AgentInfo(
        protocolVersion = LATEST_PROTOCOL_VERSION,
        capabilities = AgentCapabilities() // advertise baseline agent features
    )

    override suspend fun createSession(sessionParameters: SessionParameters): AgentSession {
        // 3. Instantiate the session implementation defined above.
        val sessionId = SessionId("session-${System.currentTimeMillis()}")
        return TerminalAgentSession(sessionId)
    }

    override suspend fun loadSession(sessionId: SessionId, sessionParameters: SessionParameters): AgentSession =
        // Rehydrate existing sessions with the provided identifier.
        TerminalAgentSession(sessionId)
}

fun main(): Unit = runBlocking {
    // 4. Bridge STDIO to the Protocol so the agent can speak ACP over stdin/stdout.
    val transport = StdioTransport(
        parentScope = this,
        input = System.`in`.asSource().buffered(),
        output = System.out.asSink().buffered()
    )
    val protocol = Protocol(this, transport)

    // 5. Register the agent and declare which remote extensions it will use.
    Agent(
        protocol = protocol,
        agentSupport = TerminalAgentSupport(),
        remoteSideExtensions = listOf(FileSystemOperations)
    )

    // 6. Start listening for messages from the client.
    protocol.start()
}

import com.agentclientprotocol.client.*
import com.agentclientprotocol.common.ClientSessionOperations
import com.agentclientprotocol.common.Event
import com.agentclientprotocol.common.SessionParameters
import com.agentclientprotocol.model.*
import com.agentclientprotocol.protocol.Protocol
import com.agentclientprotocol.transport.StdioTransport
import kotlinx.coroutines.flow.collect
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.JsonElement
import kotlinx.io.asSink
import kotlinx.io.asSource
import kotlinx.io.buffered
import java.nio.file.Path
import java.nio.file.Paths
import kotlin.io.path.readText
import kotlin.io.path.writeText

// 1. Describe how the client should create session handlers for each connection.
private class TerminalClientSupport(private val projectDir: Path) : ClientSupport {
    override suspend fun createClientSession(
        session: ClientSession,
        _sessionResponseMeta: JsonElement?
    ): ClientSessionOperations = TerminalSession(projectDir)
}

private class TerminalSession(
    private val projectDir: Path
) : ClientSessionOperations, FileSystemOperations {
    override suspend fun requestPermissions(
        toolCall: SessionUpdate.ToolCallUpdate,
        permissions: List<PermissionOption>,
        _meta: JsonElement?
    ): RequestPermissionResponse =
        // Grant whichever option was first in the list (swap for real UX).
        RequestPermissionResponse(RequestPermissionOutcome.Selected(permissions.first().optionId))

    override suspend fun notify(notification: SessionUpdate, _meta: JsonElement?) {
        // Surface streaming updates back to the host application.
        println("Agent update: $notification")
    }

    override suspend fun fsReadTextFile(
        path: String,
        line: UInt?,
        limit: UInt?,
        _meta: JsonElement?
    ): ReadTextFileResponse =
        // Resolve file paths relative to the workspace root the client chose.
        ReadTextFileResponse(projectDir.resolve(path).readText())

    override suspend fun fsWriteTextFile(
        path: String,
        content: String,
        _meta: JsonElement?
    ): WriteTextFileResponse {
        // Allow the agent to write files through the same extension API.
        projectDir.resolve(path).writeText(content)
        return WriteTextFileResponse()
    }
}

fun main(): Unit = runBlocking {
    val transport = StdioTransport(
        parentScope = this,
        input = System.`in`.asSource().buffered(),
        output = System.out.asSink().buffered()
    )
    val protocol = Protocol(this, transport)

    val projectRoot = Paths.get("").toAbsolutePath()

    val client = Client(
        // 2. Register the client support and advertise which extensions you implement.
        protocol = protocol,
        clientSupport = TerminalClientSupport(projectRoot),
        handlerSideExtensions = listOf(FileSystemOperations)
    )

    protocol.start()

    client.initialize(
        // 3. Negotiate capabilities so the agent knows extensions are available.
        ClientInfo(
            capabilities = ClientCapabilities(
                fs = FileSystemCapability(readTextFile = true, writeTextFile = true)
            )
        )
    )

    val session = client.newSession(
        // 4. Launch a session pointing at the project workspace.
        SessionParameters(
            cwd = projectRoot.toString(),
            mcpServers = emptyList()
        )
    )

    session.prompt(listOf(ContentBlock.Text("Hello agent!"))).collect { event ->
        when (event) {
            // 5. React to streaming updates and final responses.
            is Event.SessionUpdateEvent -> println("Agent update: ${event.update}")
            is Event.PromptResponseEvent -> println("Prompt finished: ${event.response.stopReason}")
        }
    }
}