package com.correx.infrastructure.tools.shell

import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.tools.compression.CompressionRule.CollapseConsecutiveDuplicates
import com.correx.core.tools.compression.CompressionRule.HeadTail
import com.correx.core.tools.compression.CompressionRule.StripBlankLines
import com.correx.core.tools.compression.DeclarativeCompressor
import com.correx.core.tools.compression.OutputCompressionSpec
import com.correx.core.tools.compression.ToolOutputCompressor
import com.correx.core.tools.contract.ParamRole
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.async
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.buildJsonArray
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put
import kotlinx.serialization.json.putJsonObject
import java.nio.file.Path

class ShellTool(
    private val allowedExecutables: Set<String> = emptySet(),
    private val timeoutMs: Long = 30_000L,
    private val workingDir: Path? = null,
) : Tool, ToolExecutor {
    override val name: String = "shell"
    override val description: String = "Execute a shell command"
    override val parametersSchema: JsonObject = buildJsonObject {
        put("type", "object")
        putJsonObject("properties") {
            putJsonObject("argv") {
                put("type", "array")
                putJsonObject("items") { put("type", "string") }
                put("description", "Command and arguments as a list. Example: [\"bash\", \"script.sh\"]")
            }
        }
        put("required", buildJsonArray { add(JsonPrimitive("argv")) })
    }
    override val tier: Tier = Tier.T2
    override val requiredCapabilities: Set<ToolCapability> =
        setOf(ToolCapability.SHELL_EXEC, ToolCapability.PROCESS_SPAWN)
    override val paramRoles: Map<String, ParamRole> = mapOf("argv" to ParamRole.EXEC_COMMAND)
    override val outputCompressor: ToolOutputCompressor = DeclarativeCompressor(
        OutputCompressionSpec(
            // Collapse repeated identical lines (retries, warning spam) before truncating, so the
            // head/tail window holds distinct content rather than duplicates. Dedup is lossless
            // (count is preserved) and order-preserving, so it is safe as a default — unlike a
            // progress-dropping regex, which can silently eat real output.
            listOf(StripBlankLines, CollapseConsecutiveDuplicates, HeadTail(HEAD_LINES, TAIL_LINES)),
            bypassOnError = true,
        ),
    )

    override fun validateRequest(request: ToolRequest): ValidationResult {
        val argv = when (val raw = request.parameters["argv"]) {
            is List<*> -> raw.filterIsInstance<String>()
            is String -> runCatching {
                Json.decodeFromString<List<String>>(raw)
            }.getOrElse { emptyList() }
            else -> emptyList()
        }
        return when {
            argv.isEmpty() ->
                ValidationResult.Invalid("Missing or empty 'argv' parameter. Expected List<String>.")

            allowedExecutables.isNotEmpty() && argv[0] !in allowedExecutables ->
                ValidationResult.Invalid("Executable '${argv[0]}' is not in the allowed list.")

            else -> ValidationResult.Valid
        }
    }

    override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
        validateRequest(request).run {
            takeIf { this is ValidationResult.Valid }?.let {
                val argv = when (val raw = request.parameters["argv"]) {
                    is List<*> -> raw.filterIsInstance<String>()
                    is String -> runCatching {
                        Json.decodeFromString<List<String>>(raw)
                    }.getOrElse { emptyList() }
                    else -> emptyList()
                }
                val process = ProcessBuilder(argv).apply {
                    workingDir?.let { directory(it.toFile()) }
                }.start()
                runCatching {
                    runCmd(request, process)
                }.getOrElse {
                    process.destroyForcibly()
                    if (it is TimeoutCancellationException) {
                        ToolResult.Failure(
                            invocationId = request.invocationId,
                            reason = "Process timed out after ${timeoutMs}ms",
                            recoverable = false,
                        )
                    } else {
                        ToolResult.Failure(
                            invocationId = request.invocationId,
                            reason = it.message ?: "Unknown error occurred during execution",
                            recoverable = false,
                        )
                    }
                }
            } ?: ToolResult.Failure(
                invocationId = request.invocationId,
                reason = (this as ValidationResult.Invalid).reason,
                recoverable = false,
            )
        }
    }

    private companion object {
        const val HEAD_LINES = 40
        const val TAIL_LINES = 40
    }

    private suspend fun runCmd(request: ToolRequest, process: Process): ToolResult = withTimeout(timeoutMs) {
        val stdoutDeferred = async { process.inputStream.bufferedReader().use { it.readText() } }
        val stderrDeferred = async { process.errorStream.bufferedReader().use { it.readText() } }

        val exitCode = process.waitFor()
        val stdout = stdoutDeferred.await()
        val stderr = stderrDeferred.await()

        val metadata = if (stderr.isNotBlank()) mapOf("stderr" to stderr) else emptyMap()
        if (exitCode != 0) {
            ToolResult.Failure(
                invocationId = request.invocationId,
                reason = "Process exited with code $exitCode${if (stderr.isNotBlank()) ": $stderr" else ""}",
                recoverable = false,
            )
        } else {
            ToolResult.Success(
                invocationId = request.invocationId,
                output = stdout,
                exitCode = exitCode,
                metadata = metadata,
            )
        }
    }
}
