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.LowQualityExtractionEvent
import com.correx.core.events.events.SourceFetchedEvent
import com.correx.core.events.events.ToolExecutionCompletedEvent
import com.correx.core.events.events.ToolExecutionStartedEvent
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.FileAffectingTool
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.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

        // 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)

        // 3. validate args centrally. A malformed/missing-arg call is the model's mistake, not a
        // fatal stage failure: return recoverable so the orchestrator surfaces it as ERROR: + the
        // tool's arg schema and the model can correct + retry (bounded by MAX_TOOL_ROUNDS), instead
        // of stranding the stage with no artifact ("no transition condition matched").
        (tool.validateRequest(request) as? ValidationResult.Invalid)?.let { invalid ->
            return@withContext ToolResult.Failure(
                invocationId = invocationId,
                reason = invalid.reason,
                recoverable = true,
            )
        }

        // 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
        when (val result = delegate.execute(request)) {
            is ToolResult.Success -> {
                restoreOrClean(backupMap, success = true)
                cleanWorkingDir(workingDir)
                // ToolExecutionCompleted/Failed are recorded once, by SessionOrchestrator (the
                // authoritative execution recorder). Here we only emit what the orchestrator can't
                // reconstruct: pre/post-image hashes (reversibility) and research-source markers.
                emitResearchSourceEvents(sessionId, request.stageId, result)
                emitFileMutations(sessionId, invocationId, affectedPaths, preImages)
                reframeIfNoOpWrite(result, affectedPaths, preImages)
            }

            is ToolResult.Failure -> {
                restoreOrClean(backupMap, success = false)
                cleanWorkingDir(workingDir)
                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
            // Name the backup by the full source path, not just the basename — two affected paths
            // sharing a filename in different dirs would otherwise overwrite each other's backup and
            // restore the wrong content on failure. Hash keeps the name filesystem-safe and unique.
            val stem = original.fileName.toString()
            val hash = Integer.toHexString(original.toAbsolutePath().normalize().toString().hashCode())
            val backupPath = workingDir.resolve("$stem.$hash.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
        }
    }

    /**
     * A write whose result is byte-identical to what was already on disk is a no-op, but the delegate
     * still reports "written successfully" — a false progress signal that traps looping agents (they
     * re-write the same content, see success, and never learn nothing changed). Rewrite the receipt to
     * the truth. Only fires when EVERY affected path pre-existed with the same content, so a partial
     * change (one file touched, one identical) is still reported as a real write.
     * ponytail: hash-compares via the CAS, so only active when artifactStore is wired; else pass-through.
     */
    private suspend fun reframeIfNoOpWrite(
        result: ToolResult.Success,
        affectedPaths: Set<Path>,
        preImages: Map<Path, PreImage>,
    ): ToolResult.Success {
        if (artifactStore == null || affectedPaths.isEmpty()) return result
        val unchanged = affectedPaths.all { path ->
            val pre = preImages[path]
            pre?.existed == true && pre.hash != null && pre.hash == storeBytes(path)
        }
        if (!unchanged) return result
        val paths = affectedPaths.joinToString(", ") { it.toString() }
        return result.copy(
            output = "No change: $paths already contained this exact content; nothing was written. " +
                "Do not repeat this write — make a different change or complete the stage.",
            metadata = result.metadata + ("noop" to "true"),
        )
    }

    // --- event emission ---

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

    /**
     * Promotes the research fetch's quality + content-hash (BACKLOG §D) from the generic
     * [ToolExecutionCompletedEvent] metadata into first-class events, additively: the metadata write
     * above is left intact for existing consumers. Only fires for results that carry the fetch
     * markers ([url] + [content_sha256] + [quality]), so non-fetch tools are untouched. The
     * low-quality threshold reuses the extractor's own marker — [ExtractionQuality.LOW_QUALITY], whose
     * name is what [WebFetchTool] writes into `quality` — rather than inventing a new policy here.
     */
    private suspend fun emitResearchSourceEvents(
        sessionId: SessionId,
        stageId: StageId,
        result: ToolResult.Success,
    ) {
        val md = result.metadata
        val url = md["url"]
        val contentSha256 = md["content_sha256"]
        val quality = md["quality"]
        if (url == null || contentSha256 == null || quality == null) return
        val byteCount = md["fetched_bytes"]?.toIntOrNull() ?: 0

        emit(sessionId, SourceFetchedEvent(sessionId, stageId, url, contentSha256, quality, byteCount))

        // Source of truth for the bar is the extractor (ExtractionQuality.LOW_QUALITY); its enum name
        // is what WebFetchTool records in `quality`. Keep the marker as a literal to avoid a core ->
        // infrastructure dependency inversion.
        if (quality == LOW_QUALITY_MARKER) {
            val reason = "extraction quality below the minimum-content bar ($byteCount bytes fetched)"
            emit(sessionId, LowQualityExtractionEvent(sessionId, stageId, url, contentSha256, 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(),
                ),
            )
        }
    }

    private companion object {
        /** Mirrors `ExtractionQuality.LOW_QUALITY.name` — the marker WebFetchTool writes into `quality`. */
        const val LOW_QUALITY_MARKER = "LOW_QUALITY"
    }
}
