package com.correx.infrastructure.tools.filesystem

import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.FileAffectingTool
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.CancellationException
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
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.io.IOException
import java.nio.file.Files
import java.nio.file.InvalidPathException
import java.nio.file.Path
import java.nio.file.Paths

@Suppress("TooManyFunctions")
class FileEditTool(
    allowedPaths: Set<Path> = emptySet(),
    private val workingDir: Path? = null,
) : Tool, FileAffectingTool, ToolExecutor {
    private val normalizedAllowedPaths: Set<Path> = allowedPaths.map { it.normalize().toAbsolutePath() }.toSet()

    private fun resolvePath(pathString: String): Path {
        val raw = Paths.get(pathString)
        return when {
            raw.isAbsolute -> raw.normalize()
            workingDir != null -> workingDir.resolve(raw).normalize()
            else -> raw.toAbsolutePath().normalize()
        }
    }

    override val name: String = "file_edit"
    override val description: String = "Edit a file. Prefer operation 'replace' (exact string " +
        "swap) for targeted edits; use file_write to replace the whole file. Use 'patch' only " +
        "when you can produce a valid unified diff that applies cleanly with 'patch -p1'."
    override val parametersSchema: JsonObject = buildJsonObject {
        put("type", "object")
        putJsonObject("properties") {
            putJsonObject("path") {
                put("type", "string")
                put("description", "Relative path to the file to edit")
            }
            putJsonObject("operation") {
                put("type", "string")
                put("description", "'replace' (recommended: exact string swap via target/" +
                    "replacement), 'append', or 'patch' (a unified diff; only if it applies cleanly)")
            }
            putJsonObject("content") {
                put("type", "string")
                put("description", "Content to append. Required for 'append'.")
            }
            putJsonObject("target") {
                put("type", "string")
                put("description", "Exact, unique string to find. Required for 'replace'.")
            }
            putJsonObject("replacement") {
                put("type", "string")
                put("description", "String to substitute for 'target'. Required for 'replace'.")
            }
            putJsonObject("patch") {
                put("type", "string")
                put("description", "A unified diff (patch -p1 format) to apply. Required for 'patch'.")
            }
        }
        put("required", buildJsonArray {
            add(JsonPrimitive("path"))
            add(JsonPrimitive("operation"))
        })
        // Per-operation requirements. `required` above can't express these (they depend on the
        // chosen operation), so a plain schema lets a constrained decoder finish a 'replace' call
        // with no target/replacement — a schema-valid but useless edit the runtime then has to
        // reject. The if/then conditionals make the operand params required once the operation is
        // fixed, so a schema-aware decoder emits them.
        put("allOf", buildJsonArray {
            add(conditionalRequirement("replace", "target", "replacement"))
            add(conditionalRequirement("append", "content"))
            add(conditionalRequirement("patch", "patch"))
        })
    }
    private fun conditionalRequirement(operation: String, vararg required: String): JsonObject =
        buildJsonObject {
            putJsonObject("if") {
                putJsonObject("properties") {
                    putJsonObject("operation") { put("const", operation) }
                }
                put("required", buildJsonArray { add(JsonPrimitive("operation")) })
            }
            putJsonObject("then") {
                put("required", buildJsonArray { required.forEach { add(JsonPrimitive(it)) } })
            }
        }

    override val tier: Tier = Tier.T3
    override val requiredCapabilities: Set<ToolCapability> = setOf(ToolCapability.FILE_WRITE)
    override val paramRoles: Map<String, ParamRole> = mapOf("path" to ParamRole.PATH)

    override fun affectedPaths(request: ToolRequest): Set<Path> {
        val pathString = request.parameters["path"] as? String ?: return emptySet()
        return setOf(resolvePath(pathString))
    }

    override fun validateRequest(request: ToolRequest): ValidationResult {
        val operation = request.parameters["operation"] as? String
        val pathString = request.parameters["path"] as? String

        return when {
            operation == null ->
                ValidationResult.Invalid("Missing 'operation' parameter. Expected 'append', 'replace', or 'patch'.")

            pathString == null ->
                ValidationResult.Invalid("Missing 'path' parameter. Expected String.")

            else -> runCatching {
                val path = resolvePath(pathString)
                checkPathAllowed(path, pathString, operation, request)
            }.getOrElse { e ->
                mapExceptionToValidationResult(e)
            }
        }
    }

    private fun checkPathAllowed(
        path: Path,
        pathString: String,
        operation: String,
        request: ToolRequest,
    ): ValidationResult {
        return when {
            normalizedAllowedPaths.isEmpty() -> ValidationResult.Invalid("No paths are allowed.")
            !isPathAllowed(path) ->
                ValidationResult.Invalid("Path '$pathString' is not in the allowed list.")

            !Files.exists(path) ->
                ValidationResult.Invalid("File not found: $pathString")

            operation == "append" && !request.parameters.containsKey("content") ->
                ValidationResult.Invalid("Missing 'content' parameter for 'append' operation.")

            operation == "replace" && missingReplaceParams(request) ->
                ValidationResult.Invalid(
                    "Missing 'target' or 'replacement' parameter for 'replace' operation. " +
                        "Provide target=\"exact string to find\" and replacement=\"new string\".",
                )

            // Anchor validity is checkable now, so reject a bad target BEFORE the approval gate
            // fires — mirroring read/write's file-exists / read-before-write pre-checks. Without
            // this the operator is prompted to approve an edit that then dies "Target not found"
            // at execute time. Recoverable: the model corrects the anchor and retries.
            operation == "replace" -> validateReplaceAnchor(path, pathString, request)

            operation == "patch" && !request.parameters.containsKey("patch") ->
                ValidationResult.Invalid("Missing 'patch' parameter for 'patch' operation.")

            // Don't point the model at 'file_read'/'list_dir' here — this rejection is commonly hit
            // in exactly the round those tools have been deliberately withheld (the read-loop
            // breaker forces a write-only round after repeated reads with no write), so telling it
            // to go use tools that aren't in its tool list this round would just compound the stall.
            operation == "read" ->
                ValidationResult.Invalid(
                    "'edit_file' has no 'read' operation. It only supports 'append', 'replace', or 'patch' — " +
                        "proceed with one of those using what you already know about the file.",
                )

            operation !in listOf("append", "replace", "patch") ->
                ValidationResult.Invalid("Unknown operation: $operation. Expected 'append', 'replace', or 'patch'.")

            else -> ValidationResult.Valid
        }
    }

    /** A replace needs a target and a replacement; 'content' is accepted as a replacement alias. */
    private fun missingReplaceParams(request: ToolRequest): Boolean =
        !request.parameters.containsKey("target") ||
            (!request.parameters.containsKey("replacement") && !request.parameters.containsKey("content"))

    private fun isPathAllowed(path: Path): Boolean =
        PathJail.isContained(path, normalizedAllowedPaths)

    /** Valid iff the replace target occurs exactly once (exact, else whitespace-flexible). */
    private fun validateReplaceAnchor(path: Path, pathString: String, request: ToolRequest): ValidationResult {
        val target = request.parameters["target"] as? String ?: return ValidationResult.Valid
        val content = Files.readString(path)
        return when (val occurrences = content.split(target).size - 1) {
            1 -> ValidationResult.Valid
            0 -> if (flexibleMatch(content, target) != null) ValidationResult.Valid
                else ValidationResult.Invalid(targetNotFoundMessage(pathString, content, target))
            else -> ValidationResult.Invalid(targetAmbiguousMessage(pathString, occurrences))
        }
    }

    /**
     * Locate [target] in [content] ignoring each line's leading/trailing whitespace — small models
     * routinely drop or misjudge indentation, so an exact-string miss is almost always an indent
     * mismatch, not a wrong edit. Returns the matched file-line range iff exactly one block matches.
     * ponytail: line-trim match only; mixed tab/space or a target spanning blank-line drift may miss.
     */
    private fun flexibleMatch(content: String, target: String): IntRange? {
        val fileLines = content.split("\n")
        val targetLines = target.split("\n").dropLastWhile { it.isBlank() }
        if (targetLines.isEmpty() || targetLines.size > fileLines.size) return null
        val normTarget = targetLines.map { it.trim() }
        val starts = (0..fileLines.size - targetLines.size).filter { start ->
            normTarget.indices.all { fileLines[start + it].trim() == normTarget[it] }
        }
        return if (starts.size == 1) starts[0] until (starts[0] + targetLines.size) else null
    }

    /** Rebase [replacement]'s indentation onto [baseIndent], preserving its own relative structure. */
    private fun reindent(replacement: String, baseIndent: String): String {
        val lines = replacement.split("\n")
        val replBase = lines.firstOrNull { it.isNotBlank() }?.takeWhile { it == ' ' || it == '\t' }.orEmpty()
        return lines.joinToString("\n") { line ->
            if (line.isBlank()) line
            else baseIndent + line.removePrefix(replBase).let { if (it == line) line.trimStart() else it }
        }
    }

    private fun targetNotFoundMessage(pathString: String, content: String, target: String): String =
        buildString {
            append("Target not found in $pathString")
            nearestLine(content, target)?.let {
                append(". Did you mean (whitespace/content may differ): $it ?")
            }
        }

    private fun targetAmbiguousMessage(pathString: String, occurrences: Int): String =
        "Target found $occurrences times in $pathString, expected exactly once — " +
            "extend it with surrounding lines to make it unique."

    private fun mapExceptionToValidationResult(e: Throwable): ValidationResult =
        when (e) {
            is InvalidPathException -> ValidationResult.Invalid("Invalid path format: ${e.message}")
            is IOException -> ValidationResult.Invalid("IO error: ${e.message}")
            is SecurityException -> ValidationResult.Invalid("Security error: ${e.message}")
            else -> ValidationResult.Invalid(e.message ?: "Unknown error occurred")
        }

    override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
        val validation = validateRequest(request)
        if (validation is ValidationResult.Invalid) {
            // Recoverable: malformed/missing params are model-correctable — feed the error back so
            // the model can fix the call (e.g. switch operation or supply the right argument)
            // instead of aborting the whole stage.
            return@withContext ToolResult.Failure(
                invocationId = request.invocationId,
                reason = validation.reason,
                recoverable = true,
            )
        }

        val operation = request.parameters["operation"] as String
        val pathString = request.parameters["path"] as String
        val path = resolvePath(pathString)
        val originalContent = runCatching { Files.readString(path) }.getOrDefault("")

        val result = runCatching {
            performOperation(operation, path, pathString, request)
        }.getOrElse { e ->
            handleExecutionException(e, request.invocationId, pathString)
        }

        // Attach a unified diff of the actual change so the side effect is auditable downstream
        // (the file_written artifact and its reviewer/operator surface carry real evidence).
        if (result is ToolResult.Success) {
            val updatedContent = runCatching { Files.readString(path) }.getOrDefault(originalContent)
            val diff = unifiedDiff(pathString, originalContent, updatedContent)
            result.copy(
                output = if (diff.isBlank()) result.output else "${result.output}\n\n$diff",
                metadata = result.metadata + ("diff" to diff),
            )
        } else {
            result
        }
    }

    private fun performOperation(
        operation: String,
        path: Path,
        pathString: String,
        request: ToolRequest,
    ): ToolResult = when (operation) {
        "append" -> append(path, request)
        "replace" -> replace(path, pathString, request)
        "patch" -> patch(path, pathString, request)
        else -> ToolResult.Failure(
            invocationId = request.invocationId,
            reason = "Unknown operation: $operation",
            recoverable = true,
        )
    }

    private fun append(path: Path, request: ToolRequest): ToolResult {
        val content = request.parameters["content"] as String
        // Read-modify-atomic-swap instead of an in-place APPEND: a crash mid-append leaves the file
        // whole (old content), never truncated. Same crash-safety FileWriteTool gets.
        AtomicFileWriter.write(path, (Files.readString(path) + content).toByteArray(Charsets.UTF_8))
        return ToolResult.Success(
            invocationId = request.invocationId,
            output = "Content appended to ${path.toAbsolutePath()}",
        )
    }

    private fun replace(path: Path, pathString: String, request: ToolRequest): ToolResult {
        val target = request.parameters["target"] as String
        // Small models routinely send 'content' (the append param) on a replace; the intent is
        // unambiguous, so accept it as replacement rather than rejecting a clearly-correct edit.
        val replacement = (request.parameters["replacement"] ?: request.parameters["content"]) as String
        val currentContent = Files.readString(path)

        val occurrences = currentContent.split(target).size - 1
        return when (occurrences) {
            1 -> {
                val newContent = currentContent.replace(target, replacement)
                writeReplaced(path, pathString, newContent, request)
            }
            // Exact miss is almost always indent drift — retry ignoring per-line whitespace, and
            // rebase the replacement onto the file's real indentation so the result stays well-formed.
            0 -> flexibleMatch(currentContent, target)?.let { range ->
                val fileLines = currentContent.split("\n")
                val baseIndent = fileLines[range.first].takeWhile { it == ' ' || it == '\t' }
                val newLines = fileLines.toMutableList()
                repeat(range.count()) { newLines.removeAt(range.first) }
                newLines.addAll(range.first, reindent(replacement, baseIndent).split("\n"))
                writeReplaced(path, pathString, newLines.joinToString("\n"), request)
            } ?: ToolResult.Failure(
                // Don't echo the (often whole-file) target back into context — it bloats the turn.
                // Point at the closest actual line instead.
                invocationId = request.invocationId,
                reason = targetNotFoundMessage(pathString, currentContent, target),
                recoverable = true,
            )
            else -> ToolResult.Failure(
                invocationId = request.invocationId,
                reason = targetAmbiguousMessage(pathString, occurrences),
                recoverable = true,
            )
        }
    }

    private fun writeReplaced(path: Path, pathString: String, newContent: String, request: ToolRequest): ToolResult {
        AtomicFileWriter.write(path, newContent.toByteArray(Charsets.UTF_8))
        return ToolResult.Success(
            invocationId = request.invocationId,
            output = "Target replaced in $pathString",
        )
    }

    /**
     * Best-effort "did you mean" for a not-found replace target: anchor on the target's first
     * non-blank line and return the file line most similar to it (levenshtein ratio), or null when
     * nothing is close enough to be a useful suggestion. Mirrors ListDirTool's fuzzy directory hint.
     */
    private fun nearestLine(content: String, target: String): String? {
        val anchor = target.lineSequence().firstOrNull { it.isNotBlank() }?.trim()?.take(MAX_CMP)
        if (anchor == null || anchor.length < MIN_ANCHOR) return null
        return content.lineSequence()
            .filter { it.isNotBlank() }
            .map { it.trim() to similarity(anchor, it.trim().take(MAX_CMP)) }
            .maxByOrNull { it.second }
            ?.takeIf { it.second >= MIN_SIMILARITY }
            ?.first
    }

    private fun similarity(a: String, b: String): Double {
        val maxLen = maxOf(a.length, b.length)
        return if (maxLen == 0) 1.0 else 1.0 - levenshtein(a, b).toDouble() / maxLen
    }

    private fun levenshtein(a: String, b: String): Int {
        val dp = IntArray(b.length + 1) { it }
        for (i in 1..a.length) {
            var prev = dp[0]
            dp[0] = i
            for (j in 1..b.length) {
                val tmp = dp[j]
                dp[j] = minOf(
                    dp[j] + 1,
                    dp[j - 1] + 1,
                    prev + if (a[i - 1].lowercaseChar() == b[j - 1].lowercaseChar()) 0 else 1,
                )
                prev = tmp
            }
        }
        return dp[b.length]
    }

    private fun patch(path: Path, pathString: String, request: ToolRequest): ToolResult {
        val patchContent = request.parameters["patch"] as String
        val parentDir = path.parent ?: Paths.get(".")
        val fileName = path.fileName.toString()

        val process = ProcessBuilder("patch", "-p1", fileName)
            .directory(parentDir.toFile())
            .redirectErrorStream(true)
            .start()

        return try {
            process.outputStream.use { it.write(patchContent.toByteArray()) }
            val output = process.inputStream.bufferedReader().use { it.readText() }
            if (!process.waitFor(PATCH_TIMEOUT_SECONDS, java.util.concurrent.TimeUnit.SECONDS)) {
                // A hung `patch` (e.g. prompting for a reject-file name it never gets) would block the
                // stage forever. Kill it and let the model fall back to replace/file_write.
                process.destroyForcibly()
                return ToolResult.Failure(
                    invocationId = request.invocationId,
                    reason = "Patch timed out after ${PATCH_TIMEOUT_SECONDS}s. " +
                        "Consider using operation 'replace' (exact string swap) or file_write instead.",
                    recoverable = true,
                )
            }
            val exitCode = process.exitValue()

            if (exitCode == 0) {
                ToolResult.Success(
                    invocationId = request.invocationId,
                    output = "Patch applied successfully to $pathString: $output",
                )
            } else {
                // Recoverable: a non-applying patch is model-correctable — feed the error back so
                // the model can fix the diff or fall back to 'replace'/file_write, not abort.
                ToolResult.Failure(
                    invocationId = request.invocationId,
                    reason = "Patch failed with exit code $exitCode: $output. " +
                        "Consider using operation 'replace' (exact string swap) or file_write instead.",
                    recoverable = true,
                )
            }
        } finally {
            process.destroy()
        }
    }

    private fun handleExecutionException(
        e: Throwable,
        invocationId: ToolInvocationId,
        pathString: String,
    ): ToolResult = when (e) {
        is CancellationException -> throw e
        is IOException -> ToolResult.Failure(
            invocationId = invocationId,
            reason = "IO error: ${e.message}",
            recoverable = false,
        )

        is SecurityException -> ToolResult.Failure(
            invocationId = invocationId,
            reason = "Access denied: $pathString, ${e.message}",
            recoverable = false,
        )

        else -> ToolResult.Failure(
            invocationId = invocationId,
            reason = e.message ?: "Unknown error occurred",
            recoverable = false,
        )
    }

    private companion object {
        const val MIN_ANCHOR = 3
        const val MAX_CMP = 400
        const val MIN_SIMILARITY = 0.5
        const val PATCH_TIMEOUT_SECONDS = 30L
    }
}
