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, e.g. [\"npm\", \"install\"]. A shell command " +
                        "line also works — cd, &&, pipes and redirects run via a shell, e.g. " +
                        "[\"cd\", \"apps/web-ui\", \"&&\", \"npm\", \"install\"].",
                )
            }
        }
        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 =
        when (val parsed = parseArgv(request.parameters["argv"])) {
            is ArgvParse.Bad -> ValidationResult.Invalid(parsed.reason)
            is ArgvParse.Ok -> when {
                allowedExecutables.isNotEmpty() && parsed.argv[0] !in allowedExecutables ->
                    ValidationResult.Invalid("Executable '${parsed.argv[0]}' is not in the allowed list.")
                else -> ValidationResult.Valid
            }
        }

    private sealed interface ArgvParse {
        data class Ok(val argv: List<String>) : ArgvParse
        data class Bad(val reason: String) : ArgvParse
    }

    // Weak local models emit malformed argv in two recurring shapes: (1) flags as JSON numbers,
    // e.g. ["mkdir", -1, ...] where "-p" should be — AnyMapSerializer decodes these to Long/Double,
    // and blindly stringifying them feeds the shell a literal "-1" ("invalid option"); (2) the whole
    // array collapsed into one escaped string, e.g. ["npm\",\"-v\","] → argv[0] = npm","-v"," which
    // is not a runnable program. Reject both here with an actionable message so the model re-emits a
    // clean array on retry, instead of burning a stage on a cryptic exec failure.
    private fun parseArgv(raw: Any?): ArgvParse = when (raw) {
        is List<*> -> {
            val nonString = raw.withIndex().firstOrNull { it.value !is String }
            when {
                raw.isEmpty() -> ArgvParse.Bad(EMPTY_MSG)
                nonString != null -> {
                    val type = nonString.value?.let { it::class.simpleName } ?: "null"
                    ArgvParse.Bad(
                        "argv element at index ${nonString.index} is a $type (`${nonString.value}`), not a string. " +
                            "Emit argv as a JSON array of strings — quote every token, " +
                            "e.g. [\"mkdir\", \"-p\", \"dir\"] (not [\"mkdir\", -p]).",
                    )
                }
                else -> checkExecutable(raw.map { it as String })
            }
        }
        is String -> runCatching { Json.decodeFromString<List<String>>(raw) }
            .fold({ if (it.isEmpty()) ArgvParse.Bad("Empty 'argv'.") else checkExecutable(it) }, {
                ArgvParse.Bad("'argv' must be a JSON array of strings, e.g. [\"ls\", \"-la\"].")
            })
        else -> ArgvParse.Bad(EMPTY_MSG)
    }

    // The model meant a shell command line (not a lone program) when argv[0] is a shell builtin or
    // any token is a shell operator — route those through `sh -c` instead of exec-ing directly.
    private fun looksLikeShellCommand(argv: List<String>): Boolean =
        argv[0] in SHELL_BUILTINS || argv.any { it in SHELL_OPERATORS }

    // argv[0] is the program name; embedded quotes/commas/whitespace mean the array collapsed into
    // one string element and there is no real executable to run.
    private fun checkExecutable(argv: List<String>): ArgvParse =
        if (argv[0].any { it == '"' || it == ',' || it.isWhitespace() }) {
            ArgvParse.Bad(
                "argv[0] `${argv[0]}` is not a valid executable name — it looks like a collapsed array. " +
                    "Emit each token as a separate string element, e.g. [\"npm\", \"-v\"].",
            )
        } else ArgvParse.Ok(argv)

    override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
        validateRequest(request).run {
            takeIf { this is ValidationResult.Valid }?.let {
                val argv = (parseArgv(request.parameters["argv"]) as ArgvParse.Ok).argv
                // Weak models reliably emit shell command lines (`cd x && npm i`, pipes, redirects)
                // rather than a lone program — that's the near-universal agent expectation. Rather
                // than fight it, honour it: when argv looks like a shell command (a builtin at [0] or
                // an operator token), run it through `sh -c` so cd/&&/|/> work. A plain program still
                // execs directly (preserves exact arg quoting). ponytail: joining tokens for sh -c
                // loses quoting of args containing spaces — fine for agent shell use, not a general
                // shell API. A still-unrunnable program returns RECOVERABLE (model adapts, not fatal).
                val command = if (looksLikeShellCommand(argv)) listOf("sh", "-c", argv.joinToString(" ")) else argv
                val process = runCatching {
                    ProcessBuilder(command).apply { workingDir?.let { directory(it.toFile()) } }.start()
                }.getOrElse {
                    return@let ToolResult.Failure(
                        invocationId = request.invocationId,
                        reason = "Could not execute '${argv[0]}': ${it.message}. Pass a runnable program, " +
                            "or a shell command line (cd/&&/pipes are run via 'sh -c').",
                        recoverable = true,
                    )
                }
                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
        const val EMPTY_MSG = "Missing or empty 'argv' parameter. Expected a JSON array of strings."
        val SHELL_BUILTINS = setOf("cd", "export", "source", ".", "pushd", "popd", "umask", "ulimit")
        val SHELL_OPERATORS = setOf("&&", "||", "|", ">", ">>", "<", ";", "&", "2>", "2>&1")
    }

    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) {
            // A non-zero exit is normal command behaviour (missing path, no grep match, git
            // --exit-code, test, diff). Surface it recoverably so the model sees the code + stderr
            // and adapts, rather than aborting the stage on the first failed probe.
            ToolResult.Failure(
                invocationId = request.invocationId,
                reason = "Process exited with code $exitCode${if (stderr.isNotBlank()) ": $stderr" else ""}",
                recoverable = true,
            )
        } else {
            ToolResult.Success(
                invocationId = request.invocationId,
                output = stdout,
                exitCode = exitCode,
                metadata = metadata,
            )
        }
    }
}
