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 com.correx.core.tools.process.ChildProcess
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.withContext
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
import java.util.concurrent.TimeUnit

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, salience = SALIENCE_PATTERN),
            ),
            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 {
                // Baseline denylist runs FIRST, ahead of the allowlist — a hard floor that holds even
                // when allowedExecutables is empty (allow-all). Blocks the class of command that
                // fetches+executes remote code or prompts interactively: an unattended run has no TTY,
                // so a create-* scaffolder / npx blocks forever (or balloons the machine). The T2 gate
                // is not enough — an auto-approve loop waves these straight through.
                deniedReason(parsed.argv) != null ->
                    ValidationResult.Invalid(deniedReason(parsed.argv)!!)
                // Check EVERY command position, not just argv[0]. An operator form runs via `sh -c`
                // (below), so `["ls", "&&", "rm", "-rf", "/"]` would slip past an argv[0]-only check
                // with allowlist {ls} and then execute `rm`. Shell builtins (cd, export…) are allowed
                // implicitly — they navigate, they don't exec an external program.
                allowedExecutables.isNotEmpty() ->
                    commandPositions(parsed.argv)
                        .firstOrNull { it !in SHELL_BUILTINS && it !in allowedExecutables }
                        ?.let { ValidationResult.Invalid("Executable '$it' is not in the allowed list.") }
                        ?: ValidationResult.Valid
                else -> ValidationResult.Valid
            }
        }

    // Return a rejection reason if argv fetches and runs an arbitrary *remote* package in one shot
    // (npx/bunx/pnpx), else null. That's remote code execution — a distinct threat we still deny.
    // Package-manager scaffolders (`npm create`, `npm init`, `npm install`) are NO LONGER blocked:
    // the hang they used to guard against (stdin prompt with no TTY) is now closed at the process
    // factory — every child gets stdin from /dev/null and a non-interactive env overlay
    // ([ChildProcess]). Matched by shape, not just argv[0], because a runner can hide inside a shell
    // command line run via `sh -c` (`cd web && npx create-foo`), where argv[0] is `cd`.
    private fun deniedReason(argv: List<String>): String? {
        // Scan every token, not just argv[0]: a runner can hide behind a shell builtin at argv[0].
        val progs = argv.map { it.substringAfterLast('/') }
        return progs.firstOrNull { it in REMOTE_EXEC_RUNNERS }?.let { runner ->
            "'$runner' downloads and runs an arbitrary remote package in one shot (remote code " +
                "execution). Use the package manager's own scaffold command instead — e.g. " +
                "`npm create vite@latest <dir> -- --template react-ts` — or write the files with file_write."
        }
    }

    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 }

    // The tokens that land in a command position once the argv is joined and run via `sh -c`:
    // argv[0], plus the token right after each command-separator (&&, ||, |, ;, &). Redirect
    // operators (>, >>, <, 2>…) are followed by a filename, not a command, so they open no new
    // command position. For a plain (non-operator) argv this is just [argv[0]].
    private fun commandPositions(argv: List<String>): List<String> = buildList {
        var expectCommand = true
        for (tok in argv) when {
            tok in COMMAND_SEPARATORS -> expectCommand = true
            tok in SHELL_OPERATORS -> expectCommand = false // redirect target follows, not a command
            expectCommand -> { add(tok); expectCommand = false }
        }
    }

    // argv[0] is the program name. A single element carrying whitespace is the weak-model "whole
    // command line in one string" collapse (["npm create vite@latest frontend"]) — split it into
    // tokens so the denylist, allowlist and executor all see a real argv, instead of rejecting it and
    // looping the stage on a call the model reliably re-emits. Weak models also leak the JSON array
    // separators into the individual tokens — trailing commas (["npm,","-v"]) or wrapping escaped-
    // quotes (["npm\",","\"-v\""]) — which the model re-emits identically until the stage loop breaks.
    // Strip stray leading/trailing quote/comma from each token so those recover too. (Internal commas,
    // --foo=a,b, are kept; a single token that is itself the whole collapsed array — "npm\",\"-v\"," —
    // still has interior quotes/commas and stays rejected: there is no runnable program to recover.
    // Splitting on whitespace loses quoting of args with spaces — same ceiling as the `sh -c` join
    // below; fine for agent shell use, not a general shell API.)
    private fun checkExecutable(raw: List<String>): ArgvParse {
        val stripped = raw.map { it.trim().trim { c -> c == '"' || c == ',' } }.filter { it.isNotEmpty() }
        val argv = if (stripped.size == 1 && stripped[0].any { it.isWhitespace() } &&
            stripped[0].none { it == '"' || it == ',' }
        ) {
            stripped[0].trim().split(WHITESPACE_RE)
        } else {
            stripped
        }
        return if (argv.isEmpty() || 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 {
                    ChildProcess.builder(command, workingDir?.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 {
                    killTree(process)
                    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
        // Case-insensitive: a decisive error line buried in the middle of a long build/test log
        // (line 250 of 500) survives the head/tail window instead of being silently dropped.
        const val SALIENCE_PATTERN = "error|fail|exception|panic|traceback|✗"
        const val EMPTY_MSG = "Missing or empty 'argv' parameter. Expected a JSON array of strings."
        val WHITESPACE_RE = Regex("\\s+")
        val SHELL_BUILTINS = setOf("cd", "export", "source", ".", "pushd", "popd", "umask", "ulimit")
        val SHELL_OPERATORS = setOf("&&", "||", "|", ">", ">>", "<", ";", "&", "2>", "2>&1")
        // Operators that chain a new command (its following token is a command position); the rest
        // of SHELL_OPERATORS are redirects whose following token is a filename.
        val COMMAND_SEPARATORS = setOf("&&", "||", "|", ";", "&")
        // Runners that download and execute an arbitrary remote package in one shot.
        val REMOTE_EXEC_RUNNERS = setOf("npx", "bunx", "pnpx")
    }

    // Combines a failed process's stdout and stderr into one labelled, head/tail-bounded block for
    // the model's failure feedback. Labels each stream (only when non-blank) so the model can tell
    // compiler diagnostics on stdout from runner errors on stderr; the same HEAD/TAIL window as the
    // success-path compressor keeps a large build log from flooding the next context window.
    private fun boundFailureOutput(stdout: String, stderr: String): String {
        val sections = buildList {
            if (stdout.isNotBlank()) add("stdout:\n${stdout.trimEnd()}")
            if (stderr.isNotBlank()) add("stderr:\n${stderr.trimEnd()}")
        }
        return boundLines(sections.joinToString("\n\n"))
    }

    // Keeps the first HEAD_LINES and last TAIL_LINES of a multi-line block, dropping the middle with
    // an elision marker — build errors cluster at both ends (first failures, final summary), so a
    // head+tail window preserves the actionable lines while bounding size.
    private fun boundLines(text: String): String {
        val lines = text.lines()
        if (lines.size <= HEAD_LINES + TAIL_LINES) return text
        val omitted = lines.size - HEAD_LINES - TAIL_LINES
        return (lines.take(HEAD_LINES) + "… ($omitted lines omitted) …" + lines.takeLast(TAIL_LINES))
            .joinToString("\n")
    }

    // Kill the whole process tree. destroyForcibly() signals only the direct child, so an `sh -c`
    // command's grandchildren (the actual npm/tsc/…) would survive as orphans; descendants() reaches
    // them. Snapshot descendants BEFORE destroying the parent — reparenting can hide them afterwards.
    private fun killTree(process: Process) {
        process.toHandle().descendants().forEach { it.destroyForcibly() }
        process.destroyForcibly()
    }

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

        // waitFor(timeout) enforces the deadline on the clock — the plain blocking waitFor() is not
        // coroutine-cancellable, so a withTimeout around it can't actually interrupt a hung process.
        val finished = withContext(Dispatchers.IO) { process.waitFor(timeoutMs, TimeUnit.MILLISECONDS) }
        if (!finished) {
            killTree(process)
            stdoutDeferred.cancel()
            stderrDeferred.cancel()
            // A timeout is NOT fatal — it's usually a long-running/watch/server command (npm run dev,
            // vite, a watcher) run in the foreground, where it never exits. Return recoverable so the
            // stage keeps going and coach the model to detach it, instead of killing the workflow with
            // a dead-end failure (no transition matched). ponytail: no background-process registry —
            // the model backgrounds it itself via `&`/nohup, which already routes through `sh -c`.
            return@coroutineScope ToolResult.Failure(
                invocationId = request.invocationId,
                reason = "Process did not exit within ${timeoutMs}ms and was killed. If this is a " +
                    "long-running command (a dev server, `npm run dev`/`vite`/`serve`, a `--watch` " +
                    "task), start it detached so it returns immediately — e.g. " +
                    "`nohup <cmd> > /tmp/dev.log 2>&1 &` — then verify separately (curl the port, or " +
                    "read the log). If you only needed its output, run a one-shot command that exits " +
                    "(e.g. `npm run build`, not `npm run dev`).",
                recoverable = true,
            )
        }
        val exitCode = process.exitValue()
        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 + output
            // and adapts, rather than aborting the stage on the first failed probe.
            //
            // BOTH streams, not just stderr: compilers and test runners (tsc, vite, jest, pytest…)
            // print their diagnostics to STDOUT, so a stderr-only failure message returns a bare
            // "exited with code 2" and the model fixes blind — it can see the build failed but not
            // which file/line. Combine stdout+stderr, head/tail-bounded so a large build log can't
            // flood the next context window.
            val detail = boundFailureOutput(stdout, stderr)
            ToolResult.Failure(
                invocationId = request.invocationId,
                reason = "Process exited with code $exitCode${if (detail.isNotBlank()) ":\n$detail" else ""}",
                recoverable = true,
            )
        } else {
            ToolResult.Success(
                invocationId = request.invocationId,
                output = stdout,
                exitCode = exitCode,
                metadata = metadata,
            )
        }
    }
}
