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
import java.nio.file.StandardOpenOption

@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"))
        })
    }
    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" && (
                    !request.parameters.containsKey("target") ||
                            !request.parameters.containsKey("replacement")
                    ) ->
                ValidationResult.Invalid("Missing 'target' or 'replacement' parameter for 'replace' operation.")

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

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

            else -> ValidationResult.Valid
        }
    }

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

    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
        Files.writeString(path, content, StandardOpenOption.APPEND)
        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
        val replacement = request.parameters["replacement"] as String
        val currentContent = Files.readString(path)

        val occurrences = currentContent.split(target).size - 1
        return if (occurrences == 1) {
            val newContent = currentContent.replace(target, replacement)
            Files.writeString(path, newContent)
            ToolResult.Success(
                invocationId = request.invocationId,
                output = "Target replaced in $pathString",
            )
        } else {
            ToolResult.Failure(
                invocationId = request.invocationId,
                reason = if (occurrences == 0) {
                    "Target '$target' not found in $pathString"
                } else {
                    "Target '$target' found $occurrences times, expected exactly once"
                },
                // Recoverable: the model can adjust the target string and retry.
                recoverable = true,
            )
        }
    }

    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() }
            val exitCode = process.waitFor()

            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,
        )
    }
}
