package link.socket.ampere.domain.capability

import kotlinx.coroutines.CoroutineScope
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.jsonPrimitive
import kotlinx.serialization.json.put
import link.socket.ampere.domain.agent.KoreAgentFactory
import link.socket.ampere.domain.agent.PromptSubAgentParams
import link.socket.ampere.domain.agent.bundled.agentArgsList
import link.socket.ampere.domain.ai.configuration.AIConfiguration
import link.socket.ampere.domain.tool.FunctionProvider
import link.socket.ampere.domain.tool.ParameterDefinition
import link.socket.ampere.util.logWith

/**
 * Represents a capability that an agent can possess.
 */
sealed class AgentCapability(open val agentTag: String) : Capability {

    override val tag: String
        get() = "$agentTag-Agent${super.tag}"

    /**
     * Capability to get a list of available LLM agents.
     */
    data class GetAgents(override val agentTag: String) : AgentCapability(agentTag) {
        override val tag: String = "${super.tag}-GetAgents"

        override val impl: Pair<String, FunctionProvider> =
            FunctionProvider.provide(
                name = "getAgents",
                description = "Returns a list of available LLM Agents, along with their respective arguments.",
                function = ::getAgents,
            )

        /**
         * Returns a newline-separated list of agent arguments.
         */
        private fun getAgents(): String =
            agentArgsList.joinToString("\n\n").also { agentList ->
                logWith(tag).i("\nResponse: $agentList")
            }
    }

    /**
     * Capability to prompt another agent for completion based on given inputs.
     *
     * @property conversationRepository Repository to manage conversation data.
     * @property scope Coroutine scope for managing asynchronous operations.
     */
    data class PromptAgent(
        override val agentTag: String,
        val scope: CoroutineScope,
        val config: AIConfiguration,
        val agentFactory: KoreAgentFactory,
    ) : AgentCapability(agentTag) {

        override val tag: String = "${super.tag}-PromptAgent"

        override val impl: Pair<String, FunctionProvider> =
            FunctionProvider.provideSuspend(
                name = "promptAgent",
                description = """
                    Requests a completion from another Agent with the given prompt and an example user's response to the LLM's initial Chat.
                    Returns the Chat completion generated by the sub-Agent.
                    You must provide either an Agent name or a prompt, but **never** both.
                """.trimIndent(),
                function = { args: JsonObject ->
                    val parentConversationId: String =
                        args.getOrElse("parentConversationId") {
                            JsonPrimitive("")
                        }.jsonPrimitive.content

                    val agent: String =
                        args.getOrElse("agent") {
                            JsonPrimitive("")
                        }.jsonPrimitive.content

                    val prompt: String =
                        args.getOrElse("prompt") {
                            JsonPrimitive("")
                        }.jsonPrimitive.content

                    val initialUserChat: String =
                        args.getOrElse("initialUserChat") {
                            JsonPrimitive("")
                        }.jsonPrimitive.content

                    val params = PromptSubAgentParams(
                        config = config,
                        parentConversationId = parentConversationId,
                        scope = scope,
                        agentName = agent,
                        prompt = prompt,
                        initialUserChat = initialUserChat,
                    )

                    agentFactory.promptSubAgent(params)
                },
                parameterList = listOf(
                    ParameterDefinition(
                        name = "parentConversationId",
                        isRequired = true,
                        definition = buildJsonObject {
                            put("type", "string")
                            put(
                                "description",
                                """
                                    The ID of the parent Conversation that is calling this function.
                                """.trimIndent(),
                            )
                        },
                    ),
                    ParameterDefinition(
                        name = "agent",
                        isRequired = false,
                        definition = buildJsonObject {
                            put("type", "string")
                            put(
                                "description",
                                """
                                    The name of the LLM Agent that should be completing the prompt. *Cannot* be used if prompt is being sent.
                                    This name *must* match the name returned by the `getAgents` function - i.e. do *not* include the word 'Agent' in this parameter.
                                """.trimIndent(),
                            )
                        },
                    ),
                    ParameterDefinition(
                        name = "prompt",
                        isRequired = false,
                        definition = buildJsonObject {
                            put("type", "string")
                            put(
                                "description",
                                """
                                    The system instructions to use in case an Agent is not specified.
                                    *Cannot* be used if an agent is being sent.
                                """.trimIndent(),
                            )
                        },
                    ),
                    ParameterDefinition(
                        name = "initialUserChat",
                        isRequired = false,
                        definition = buildJsonObject {
                            put("type", "string")
                            put(
                                "description",
                                """
                                    Send an initial User Chat message to the Agent, which will be submitted to 
                                    the Agent after the Agent's first response is generated.
                                        
                                    This property should be used when you know what question you want to ask the Agent
                                    and you don't want to wait for the Agent's initial response to ask the question.
                                """.trimIndent(),
                            )
                        },
                    ),
                ),
            )
    }
}
