package com.correx.infrastructure.tools

import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.EventDispatcher
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.FileWrittenEvent
import com.correx.core.events.events.ToolExecutionCompletedEvent
import com.correx.core.events.events.ToolExecutionFailedEvent
import com.correx.core.events.events.ToolExecutionStartedEvent
import com.correx.core.events.events.ToolReceipt
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.FileAffectingTool
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.registry.ToolRegistry
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import kotlinx.datetime.Clock
import java.io.IOException
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.StandardCopyOption

@SuppressWarnings("TooManyFunctions", "LongParameterList")
class SandboxedToolExecutor(
    private val delegate: ToolExecutor,
    private val registry: ToolRegistry,
    private val eventDispatcher: EventDispatcher,
    private val workDir: Path = Path.of("/tmp/correx-sandbox"),
    private val artifactStore: ArtifactStore? = null,
) : ToolExecutor {

    private data class PreImage(val existed: Boolean, val hash: String?)

    override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
        val sessionId = request.sessionId
        val invocationId = request.invocationId
        val toolName = request.toolName
        val startTime = System.currentTimeMillis()

        // 1. resolve tool
        val tool = registry.resolve(toolName)
            ?: return@withContext ToolResult.Failure(
                invocationId = invocationId,
                reason = "tool not found: $toolName",
                recoverable = false,
            )

        // 2. emit started
        emitStarted(sessionId, invocationId, toolName)

        // 4. create working dir
        val workingDir = workDir.resolve(sessionId.value).resolve(invocationId.value)
        Files.createDirectories(workingDir)

        // 5. compute affected paths once (A3: avoid double-invocation divergence)
        val affectedPaths: Set<Path> = if (tool is FileAffectingTool) tool.affectedPaths(request) else emptySet()

        // 5b. capture pre-image bytes into CAS for event-sourced reversibility (invariants #1/#5)
        val preImages: Map<Path, PreImage> = capturePreImages(affectedPaths)

        // 6. backup affected files
        val backupMap: Map<Path, Path> = try {
            backupAffectedFiles(affectedPaths, workingDir)
        } catch (e: IOException) {
            return@withContext ToolResult.Failure(
                invocationId = invocationId,
                reason = "backup failed: ${e.message}",
                recoverable = false,
            )
        }

        // 7. delegate
        val durationMs = { System.currentTimeMillis() - startTime }

        when (val result = delegate.execute(request)) {
            is ToolResult.Success -> {
                val diff = computeDiff(affectedPaths, backupMap)
                restoreOrClean(backupMap, success = true)
                cleanWorkingDir(workingDir)
                emitCompleted(sessionId, invocationId, toolName, result, tool, affectedPaths, durationMs(), diff)
                emitFileMutations(sessionId, invocationId, affectedPaths, preImages)
                result
            }

            is ToolResult.Failure -> {
                restoreOrClean(backupMap, success = false)
                cleanWorkingDir(workingDir)
                emitFailed(sessionId, invocationId, toolName, result.reason)
                result
            }
        }
    }

    // --- backup/restore ---

    /**
     * Returns a map of originalPath -> backupPath for all affected files.
     * Throws IOException if any backup fails — caller handles this.
     */
    private fun backupAffectedFiles(
        affectedPaths: Collection<Path>,
        workingDir: Path,
    ): Map<Path, Path> = buildMap {
        for (original in affectedPaths) {
            // A1: skip files that don't yet exist (new-file tools have nothing to back up)
            if (!Files.exists(original)) continue
            val backupPath = workingDir.resolve("${original.fileName}.bak")
            Files.copy(original, backupPath, StandardCopyOption.REPLACE_EXISTING)
            put(original, backupPath)
        }
    }

    /**
     * On success: delete backups.
     * On failure: restore originals from backups.
     */
    private fun restoreOrClean(backupMap: Map<Path, Path>, success: Boolean) {
        for ((original, backup) in backupMap) {
            try {
                if (success) {
                    Files.deleteIfExists(backup)
                } else {
                    Files.move(backup, original, StandardCopyOption.REPLACE_EXISTING)
                }
            } catch (_: IOException) {
                // best effort
            }
        }
    }

    private fun cleanWorkingDir(workingDir: Path) {
        try {
            Files.walk(workingDir)
                .sorted(Comparator.reverseOrder())
                .forEach { Files.deleteIfExists(it) }
        } catch (_: IOException) {
            // best effort
        }
    }

    // --- event emission ---

    private suspend fun emitStarted(
        sessionId: SessionId,
        invocationId: ToolInvocationId,
        toolName: String,
    ) = emit(sessionId, ToolExecutionStartedEvent(invocationId, sessionId, toolName))

    private suspend fun emitCompleted(
        sessionId: SessionId,
        invocationId: ToolInvocationId,
        toolName: String,
        result: ToolResult.Success,
        tool: Tool,
        affectedPaths: Set<Path>,
        durationMs: Long,
        diff: String? = null,
    ) {
        val affectedEntities = affectedPaths.map { it.toString() }
        emit(
            sessionId,
            ToolExecutionCompletedEvent(
                invocationId = invocationId,
                sessionId = sessionId,
                toolName = toolName,
                receipt = ToolReceipt(
                    invocationId = invocationId,
                    toolName = toolName,
                    exitCode = result.exitCode,
                    outputSummary = result.output,
                    structuredOutput = result.metadata,
                    affectedEntities = affectedEntities,
                    durationMs = durationMs,
                    tier = tool.tier,
                    timestamp = Clock.System.now(),
                    diff = diff,
                ),
            ),
        )
    }

    @Suppress("MaxLineLength")
    private fun computeDiff(affectedPaths: Set<Path>, backupMap: Map<Path, Path>): String? {
        val diffs = mutableListOf<String>()

        // Files that existed before the tool ran (backed up)
        for ((original, backup) in backupMap) {
            runCatching {
                val oldContent = Files.readString(backup)
                val newContent = if (Files.exists(original)) Files.readString(original) else ""
                if (oldContent != newContent) {
                    diffs.add(DiffUtil.unifiedDiff(oldContent, newContent, original.toString()))
                }
            }
        }

        // Newly created files (no backup, but affected and now exist)
        for (path in affectedPaths) {
            if (!backupMap.containsKey(path) && Files.exists(path)) {
                runCatching {
                    val newContent = Files.readString(path)
                    if (newContent.isNotEmpty()) {
                        diffs.add(DiffUtil.unifiedDiff("", newContent, path.toString()))
                    }
                }
            }
        }

        if (diffs.isEmpty()) return null
        return diffs.joinToString("\n")
    }

    private suspend fun emitFailed(
        sessionId: SessionId,
        invocationId: ToolInvocationId,
        toolName: String,
        reason: String,
    ) = emit(sessionId, ToolExecutionFailedEvent(invocationId, sessionId, toolName, reason))

    private suspend fun emit(sessionId: SessionId, payload: EventPayload) {
        eventDispatcher.emit(payload, sessionId)
    }

    // --- event-sourced reversibility (CAS pre/post images) ---

    private suspend fun capturePreImages(paths: Set<Path>): Map<Path, PreImage> {
        if (artifactStore == null) return emptyMap()
        return paths.associateWith { path ->
            val existed = Files.exists(path)
            PreImage(existed = existed, hash = if (existed) storeBytes(path) else null)
        }
    }

    private suspend fun storeBytes(path: Path): String? {
        val store = artifactStore ?: return null
        val bytes = runCatching { Files.readAllBytes(path) }.getOrNull() ?: return null
        return store.put(bytes).value
    }

    /** Emits one [FileWrittenEvent] per affected path, carrying the CAS hashes that make the
     *  mutation reversible. Post-image is read after the swap; null ⇒ the path was deleted. */
    private suspend fun emitFileMutations(
        sessionId: SessionId,
        invocationId: ToolInvocationId,
        affectedPaths: Set<Path>,
        preImages: Map<Path, PreImage>,
    ) {
        if (artifactStore == null) return
        for (path in affectedPaths) {
            val pre = preImages[path]
            emit(
                sessionId,
                FileWrittenEvent(
                    invocationId = invocationId,
                    sessionId = sessionId,
                    path = path.toString(),
                    preImageHash = pre?.hash,
                    postImageHash = storeBytes(path),
                    preExisted = pre?.existed ?: false,
                    timestampMs = Clock.System.now().toEpochMilliseconds(),
                ),
            )
        }
    }
}
