package net.kentonc.si.agent

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.add
import kotlinx.serialization.json.buildJsonArray
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put
import java.io.File
import java.nio.ByteBuffer
import java.nio.charset.CodingErrorAction
import java.util.concurrent.TimeUnit

/**
 * Tools the model can call: filesystem read/search/write and a shell, a port of
 * src/tools.rs. Every tool resolves its path relative to a root directory (the
 * working directory), so the model can't wander off into the rest of the device.
 * Writes and shell execution additionally require approval, supplied by the host
 * (AgentViewModel) as a suspend callback that shows a dialog.
 *
 * Tool results are Strings on purpose: error text is formatted like the Rust
 * original ("Error: ...") because the model reads it, not the UI.
 */
class ToolHost(private val rootDir: File) {

    /** JSON schema list sent with the request, declaring the available tools. */
    val specs: JsonElement = buildJsonArray {
        add(toolSpec("read_file", "Read a UTF-8 text file. Returns the file contents, truncated if long.", listOf(
            prop("path", "File path, relative to the working directory."),
        ), required = listOf("path")))
        add(toolSpec("list_dir", "List the entries of a directory, with sizes and a / suffix on directories.", listOf(
            prop("path", "Directory path, relative to the working directory. Defaults to '.'"),
        )))
        add(toolSpec("search", "Search file contents for a regular expression across the working directory. Returns path:line:match entries.", listOf(
            prop("pattern", "Regular expression to look for."),
            prop("path", "Subdirectory to search. Defaults to '.'"),
            prop("include", "Only search files whose name contains this, e.g. '.rs' or 'main'."),
        ), required = listOf("pattern")))
        add(toolSpec("write_file", "Write text to a file, creating parent directories as needed. Overwrites existing files. Requires user approval.", listOf(
            prop("path", "File path, relative to the working directory."),
            prop("content", "Full contents to write."),
        ), required = listOf("path", "content")))
        add(toolSpec("run_shell", "Run a command with sh (`sh -c`) in the working directory. Requires user approval.", listOf(
            prop("command", "The command line to execute."),
        ), required = listOf("command")))
    }

    /** Run one tool call, returning the text to hand back to the model. */
    suspend fun execute(
        call: ToolCall,
        autoApprove: Boolean,
        confirm: suspend (ToolCall) -> Boolean,
    ): String {
        val args: JsonObject = try {
            RootArgsJson.parseToJsonElement(call.function.arguments.ifBlank { "{}" }) as? JsonObject
                ?: JsonObject(emptyMap())
        } catch (e: Exception) {
            return "Error: arguments were not valid JSON: ${e.message}"
        }

        return try {
            when (call.function.name) {
                "read_file" -> readFile(arg(args, "path"))
                "list_dir" -> listDir(arg(args, "path"))
                "search" -> {
                    val pattern = arg(args, "pattern")
                    if (pattern.isEmpty()) "Error: `pattern` is required."
                    else search(pattern, arg(args, "path"), arg(args, "include"))
                }
                "write_file" -> {
                    if (approved(call, autoApprove, confirm)) writeFile(arg(args, "path"), arg(args, "content"))
                    else "Skipped: the user declined this write. Do not retry it; ask them what they'd prefer."
                }
                "run_shell" -> {
                    val command = arg(args, "command")
                    if (command.trim().isEmpty()) "Error: `command` is required."
                    else if (approved(call, autoApprove, confirm)) runShell(command)
                    else "Skipped: the user declined this command. Do not retry it; describe what you would have run instead."
                }
                else -> "Error: unknown tool `${call.function.name}`."
            }
        } catch (ce: kotlinx.coroutines.CancellationException) {
            throw ce
        } catch (e: Exception) {
            "Error: ${e.message ?: e.javaClass.simpleName}"
        }
    }

    // ─── filesystem ────────────────────────────────────────────────────────────

    private fun readFile(path: String): String {
        if (path.trim().isEmpty()) error("`path` is required")
        val file = resolve(path, mustExist = true)
        val bytes = file.readBytes()
        // Strict decoding so a binary file fails the same way the Rust one does.
        val text = try {
            Charsets.UTF_8.newDecoder()
                .onMalformedInput(CodingErrorAction.REPORT)
                .onUnmappableCharacter(CodingErrorAction.REPORT)
                .decode(ByteBuffer.wrap(bytes)).toString()
        } catch (e: Exception) {
            error("${file.path} is not a UTF-8 text file")
        }
        return truncate(text)
    }

    private fun listDir(path: String): String {
        val dir = resolve(path, mustExist = true)
        if (!dir.isDirectory) error("${dir.path} is not a directory")

        val entries = (dir.listFiles() ?: emptyArray())
            .map { Triple(it.name, it.isDirectory, it.length()) }
        if (entries.isEmpty()) return "(empty directory)"

        return entries
            .sortedWith(compareByDescending<Triple<String, Boolean, Long>> { it.second }.thenBy { it.first })
            .take(500)
            .joinToString("\n") { (name, isDir, size) ->
                if (isDir) "$name/" else "$name (${humanSize(size)})"
            }
    }

    private fun writeFile(path: String, content: String): String {
        if (path.trim().isEmpty()) error("`path` is required")
        val file = resolve(path, mustExist = false)
        file.parentFile?.mkdirs()
        file.writeText(content)
        return "Wrote ${humanSize(content.length.toLong())} to ${file.path}."
    }

    // ─── search ────────────────────────────────────────────────────────────────

    /** Directories we never descend into (SKIP_DIRS in tools.rs). */
    private val skipDirs = setOf("target", ".git", "node_modules", "__pycache__", ".venv", "dist", "build")

    private fun search(pattern: String, path: String, include: String): String {
        val regex = Regex(pattern)
        val start = resolve(path, mustExist = true)
        if (!start.isDirectory) error("${start.path} is not a directory")

        val results = ArrayList<String>()
        var filesScanned = 0
        var truncated = false
        val stack = ArrayDeque(listOf(start))

        outer@ while (stack.isNotEmpty()) {
            val dir = stack.removeLast()
            for (entry in dir.listFiles() ?: continue) {
                if (truncated) break@outer
                val name = entry.name
                if (entry.isDirectory) {
                    if (name !in skipDirs && !name.startsWith(".")) stack.addLast(entry)
                    continue
                }
                if (include.isNotEmpty() && !name.contains(include)) continue
                if (entry.length() > 1_000_000) continue  // skip binaries and huge files
                val text = runCatching { entry.readText() }.getOrNull() ?: continue
                filesScanned++
                for ((index, line) in text.lineSequence().withIndex()) {
                    if (regex.containsMatchIn(line)) {
                        if (results.size >= 200) { truncated = true; break@outer }
                        results.add("${entry.path}:${index + 1}:${snippet(line.trim(), 300)}")
                    }
                }
            }
        }

        if (results.isEmpty()) return "No matches for /$pattern/ in $filesScanned file(s)."
        var out = results.joinToString("\n")
        if (truncated) out += "\n… (truncated; narrow the pattern or path)"
        return truncate(out)
    }

    // ─── shell ─────────────────────────────────────────────────────────────────

    private suspend fun runShell(command: String): String = withContext(Dispatchers.IO) {
        val process = ProcessBuilder("sh", "-c", command)
            .directory(rootDir)
            // Android's Redirect has no NULL constant (Java 9's DISCARD
            // arrived much later); INHERIT is harmless: an app process has
            // nothing meaningful on stdin for a shell tool to read.
            .apply { redirectInput(ProcessBuilder.Redirect.INHERIT) }
            .start()

        val stdout = StringBuilder()
        val stderr = StringBuilder()
        // Both pipes drained on threads so a big output can't block the process.
        val outBuf = java.io.ByteArrayOutputStream()
        val errBuf = java.io.ByteArrayOutputStream()
        val t1 = Thread { process.inputStream.use { it.copyTo(outBuf) } }
        val t2 = Thread { process.errorStream.use { it.copyTo(errBuf) } }
        t1.start(); t2.start()

        val finished = process.waitFor(SHELL_TIMEOUT_SECS, TimeUnit.SECONDS)
        if (!finished) {
            process.destroyForcibly()
            t1.join(1000); t2.join(1000)
            error("timed out after ${SHELL_TIMEOUT_SECS}s")
        }
        t1.join(); t2.join()

        var out = ""
        if (outBuf.size() > 0) out += String(outBuf.toByteArray(), Charsets.UTF_8).trimEnd()
        if (errBuf.size() > 0) {
            if (out.isNotEmpty()) out += "\n\n[stderr]"
            out += String(errBuf.toByteArray(), Charsets.UTF_8).trimEnd()
        }
        if (out.isEmpty()) out = "(no output)"
        if (process.exitValue() != 0) out += "\n\n[exit status: ${process.exitValue()}]"
        truncate(out.trim())
    }

    // ─── approval, root resolution, helpers ────────────────────────────────────

    private suspend fun approved(call: ToolCall, autoApprove: Boolean, confirm: suspend (ToolCall) -> Boolean) =
        autoApprove || confirm(call)

    /**
     * Resolve a path against the working directory, refusing to leave it
     * (resolve() in tools.rs).
     */
    private fun resolve(raw: String, mustExist: Boolean): File {
        val root = rootDir.canonicalFile
        val joined = if (raw.trim().isEmpty() || raw.trim() == ".") root
        else {
            val candidate = File(raw)
            if (candidate.isAbsolute) candidate else File(root, raw)
        }

        val resolved = if (mustExist) {
            joined.canonicalFile
        } else {
            // The file may not exist yet, so anchor on the nearest existing ancestor.
            joined.parentFile?.canonicalFile?.resolve(joined.name) ?: joined
        }

        if (resolved.canonicalPath != root.canonicalPath &&
            !resolved.canonicalPath.startsWith(root.canonicalPath + File.separator)
        ) {
            error("${resolved.path} is outside the working directory (${root.path})")
        }
        return resolved
    }

    private fun arg(args: JsonObject, key: String): String =
        (args[key] as? JsonPrimitive)?.content ?: ""

    /** Cut at a char boundary so the result fits, appending a marker. */
    private fun truncate(text: String, max: Int = TOOL_OUTPUT_MAX_CHARS) =
        if (text.length <= max) text else text.take(max) + "\n… (truncated)"

    private fun snippet(text: String, max: Int) =
        if (text.length <= max) text else text.take(max) + "…"

    companion object {
        private const val TOOL_OUTPUT_MAX_CHARS = 16_000
        private const val SHELL_TIMEOUT_SECS = 60L
        private val RootArgsJson = kotlinx.serialization.json.Json { ignoreUnknownKeys = true }
    }
}

/** 500 B -> "500 B", 2048 -> "2.0 KB", etc. (human_size in tools.rs) */
fun humanSize(bytes: Long): String {
    val units = listOf("B", "KB", "MB", "GB", "TB")
    var value = bytes.toDouble()
    var unit = 0
    while (value >= 1024.0 && unit < units.size - 1) { value /= 1024.0; unit++ }
    return if (unit == 0) "$bytes B" else "%.1f %s".format(value, units[unit])
}

/** MIME types accepted for image attachments, keyed by file extension. */
val IMAGE_MIMES = mapOf(
    "png" to "image/png",
    "jpg" to "image/jpeg",
    "jpeg" to "image/jpeg",
    "gif" to "image/gif",
    "webp" to "image/webp",
)

// ─── JSON schema helpers used by `specs` ───────────────────────────────────────

private fun prop(name: String, description: String) = name to buildJsonObject {
    put("type", "string")
    put("description", description)
}

private fun toolSpec(
    name: String,
    description: String,
    props: List<Pair<String, JsonObject>>,
    required: List<String> = emptyList(),
): JsonObject = buildJsonObject {
    put("type", "function")
    put("function", buildJsonObject {
        put("name", name)
        put("description", description)
        put("parameters", buildJsonObject {
            put("type", "object")
            put("properties", JsonObject(props.toMap()))
            if (required.isNotEmpty()) put("required", JsonArray(required.map { JsonPrimitive(it) }))
        })
    })
}