package com.correx.apps.server.undo

import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.FileWrittenEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.SessionWorkspaceBoundEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.ToolInvocationId
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Clock
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Test
import java.nio.file.Files
import java.nio.file.Path

class SessionUndoServiceTest {

    private class FakeArtifactStore : ArtifactStore {
        val blobs = linkedMapOf<String, ByteArray>()
        override suspend fun put(bytes: ByteArray): ArtifactId {
            val id = ArtifactId(bytes.toString(Charsets.UTF_8).hashCode().toString())
            blobs[id.value] = bytes
            return id
        }
        override suspend fun get(id: ArtifactId): ByteArray? = blobs[id.value]
        override suspend fun flushBefore(commit: suspend () -> Unit) = commit()
    }

    private class FakeEventStore(private val events: List<StoredEvent>) : EventStore {
        override suspend fun append(event: NewEvent): StoredEvent = error("unused")
        override suspend fun appendAll(events: List<NewEvent>): List<StoredEvent> = error("unused")
        override fun read(sessionId: SessionId): List<StoredEvent> = events
        override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> =
            events.filter { it.sequence >= fromSequence }
        override fun lastSequence(sessionId: SessionId): Long? = events.lastOrNull()?.sequence
        override fun subscribe(sessionId: SessionId): Flow<StoredEvent> = emptyFlow()
        override fun subscribeAll(): Flow<StoredEvent> = emptyFlow()
        override suspend fun lastGlobalSequence(): Long = 0L
        override fun allEvents(): Sequence<StoredEvent> = emptySequence()
        override fun allSessionIds(): Set<SessionId> = emptySet()
    }

    private fun storedEvent(payload: EventPayload, sessionId: SessionId, seq: Long): StoredEvent =
        StoredEvent(
            metadata = EventMetadata(
                eventId = EventId("e$seq"),
                sessionId = sessionId,
                timestamp = Clock.System.now(),
                schemaVersion = 1,
                causationId = null,
                correlationId = null,
            ),
            sequence = seq,
            sessionSequence = seq,
            payload = payload,
        )

    @Test
    fun `undo restores a file that was overwritten`(): Unit = runBlocking {
        val dir = Files.createTempDirectory("undo-test").toRealPath()
        val target = dir.resolve("file.txt")
        Files.writeString(target, "NEW")

        val fakeStore = FakeArtifactStore()
        val preImageHash = fakeStore.put("OLD".toByteArray()).value

        val sessionId = SessionId("session-overwrite")
        val event = FileWrittenEvent(
            invocationId = ToolInvocationId("inv-1"),
            sessionId = sessionId,
            path = target.toString(),
            preImageHash = preImageHash,
            postImageHash = null,
            preExisted = true,
            timestampMs = 0L,
        )
        val fakeEventStore = FakeEventStore(listOf(storedEvent(event, sessionId, 1L)))
        val service = SessionUndoService(fakeEventStore, fakeStore, setOf(dir))

        val summary = service.undo(sessionId)

        assertEquals(1, summary.reverted)
        assertEquals(0, summary.deleted)
        assertEquals(0, summary.rejected)
        assertEquals(0, summary.failed)
        assertEquals("OLD", Files.readString(target))
    }

    @Test
    fun `undo deletes a file that was newly created`(): Unit = runBlocking {
        val dir = Files.createTempDirectory("undo-new").toRealPath()
        val target = dir.resolve("new.txt")
        Files.writeString(target, "CREATED")

        val fakeStore = FakeArtifactStore()
        val sessionId = SessionId("session-create")
        val event = FileWrittenEvent(
            invocationId = ToolInvocationId("inv-2"),
            sessionId = sessionId,
            path = target.toString(),
            preImageHash = null,
            postImageHash = null,
            preExisted = false,
            timestampMs = 0L,
        )
        val fakeEventStore = FakeEventStore(listOf(storedEvent(event, sessionId, 1L)))
        val service = SessionUndoService(fakeEventStore, fakeStore, setOf(dir))

        val summary = service.undo(sessionId)

        assertEquals(0, summary.reverted)
        assertEquals(1, summary.deleted)
        assertEquals(0, summary.rejected)
        assertEquals(0, summary.failed)
        assertFalse(Files.exists(target))
    }

    /**
     * When the session log contains a [SessionWorkspaceBoundEvent] the workspace root
     * is unioned into the jail roots, so a file written under that workspace (which is
     * NOT in the boot roots) is successfully reverted rather than rejected.
     *
     * The sibling assertion that proves the workspace exclusion: a service built with
     * boot roots that do NOT include [workspaceDir] would return rejected=1 instead of
     * reverted=1. We verify this via a separate call using an unrelated boot root.
     */
    @Test
    fun `undo uses session workspace root from SessionWorkspaceBoundEvent`(): Unit = runBlocking {
        val bootDir = Files.createTempDirectory("undo-boot").toRealPath()
        val workspaceDir = Files.createTempDirectory("undo-workspace").toRealPath()
        val target = workspaceDir.resolve("ws-file.txt")
        Files.writeString(target, "CURRENT")

        val fakeStore = FakeArtifactStore()
        val preImageHash = fakeStore.put("ORIGINAL".toByteArray()).value

        val sessionId = SessionId("session-workspace")

        val workspaceBound = SessionWorkspaceBoundEvent(
            sessionId = sessionId,
            workspaceRoot = workspaceDir.toString(),
            allowedPaths = emptyList(),
        )
        val fileWritten = FileWrittenEvent(
            invocationId = ToolInvocationId("inv-ws"),
            sessionId = sessionId,
            path = target.toString(),
            preImageHash = preImageHash,
            postImageHash = null,
            preExisted = true,
            timestampMs = 0L,
        )

        val eventList = listOf(
            storedEvent(workspaceBound, sessionId, 1L),
            storedEvent(fileWritten, sessionId, 2L),
        )
        val fakeEventStore = FakeEventStore(eventList)

        // Boot roots do NOT include workspaceDir — the workspace event unlocks it.
        val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))
        val summary = service.undo(sessionId)

        assertEquals(1, summary.reverted)
        assertEquals(0, summary.rejected, "workspace file must not be jailed out when bound event is present")
        assertEquals("ORIGINAL", Files.readString(target))

        // Prove that without the workspace event the revert would have been jailed:
        // a service with boot-only roots (excluding workspaceDir) and no bound event rejects the file.
        val storeNoWorkspace = FakeArtifactStore().also { it.blobs[preImageHash] = "ORIGINAL".toByteArray() }
        val noWorkspaceEventStore = FakeEventStore(listOf(storedEvent(fileWritten, sessionId, 1L)))
        val serviceBootOnly = SessionUndoService(noWorkspaceEventStore, storeNoWorkspace, setOf(bootDir))
        val summaryBootOnly = serviceBootOnly.undo(sessionId)
        assertEquals(
            1,
            summaryBootOnly.rejected,
            "boot-only roots must jail the workspace-dir file when no bound event exists",
        )
    }

    /**
     * A session with NO [SessionWorkspaceBoundEvent] uses boot roots exactly as before.
     * Files inside boot roots are reverted; files outside are rejected.
     */
    @Test
    fun `undo with no SessionWorkspaceBoundEvent falls back to boot roots only`(): Unit = runBlocking {
        val bootDir = Files.createTempDirectory("undo-boot-compat").toRealPath()
        val outsideDir = Files.createTempDirectory("undo-outside").toRealPath()

        val fakeStore = FakeArtifactStore()
        val preImageHash = fakeStore.put("PREV".toByteArray()).value

        val sessionId = SessionId("session-no-workspace")

        // File inside boot root — should be reverted.
        val insideTarget = bootDir.resolve("inside.txt")
        Files.writeString(insideTarget, "NEW")

        // File outside boot root and no bound event — should be rejected.
        val outsideTarget = outsideDir.resolve("outside.txt")
        Files.writeString(outsideTarget, "NEW")

        val insideEvent = FileWrittenEvent(
            invocationId = ToolInvocationId("inv-in"),
            sessionId = sessionId,
            path = insideTarget.toString(),
            preImageHash = preImageHash,
            postImageHash = null,
            preExisted = true,
            timestampMs = 0L,
        )
        val outsideEvent = FileWrittenEvent(
            invocationId = ToolInvocationId("inv-out"),
            sessionId = sessionId,
            path = outsideTarget.toString(),
            preImageHash = preImageHash,
            postImageHash = null,
            preExisted = true,
            timestampMs = 0L,
        )

        val fakeEventStore = FakeEventStore(
            listOf(storedEvent(insideEvent, sessionId, 1L), storedEvent(outsideEvent, sessionId, 2L)),
        )
        val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))

        val summary = service.undo(sessionId)

        assertEquals(1, summary.reverted, "file inside boot root must be reverted")
        assertEquals(1, summary.rejected, "file outside boot root must be rejected when no workspace event exists")
    }

    /**
     * When [SessionWorkspaceBoundEvent.allowedPaths] carries additional paths beyond
     * [SessionWorkspaceBoundEvent.workspaceRoot], those paths are also included in the
     * effective jail roots.
     */
    @Test
    fun `undo includes allowedPaths from SessionWorkspaceBoundEvent`(): Unit = runBlocking {
        val bootDir = Files.createTempDirectory("undo-boot-allowed").toRealPath()
        val workspaceDir = Files.createTempDirectory("undo-ws-root").toRealPath()
        val extraAllowedDir = Files.createTempDirectory("undo-extra-allowed").toRealPath()
        val target = extraAllowedDir.resolve("extra.txt")
        Files.writeString(target, "STALE")

        val fakeStore = FakeArtifactStore()
        val preImageHash = fakeStore.put("FRESH".toByteArray()).value

        val sessionId = SessionId("session-allowed-paths")

        val workspaceBound = SessionWorkspaceBoundEvent(
            sessionId = sessionId,
            workspaceRoot = workspaceDir.toString(),
            allowedPaths = listOf(extraAllowedDir.toString()),
        )
        val fileWritten = FileWrittenEvent(
            invocationId = ToolInvocationId("inv-extra"),
            sessionId = sessionId,
            path = target.toString(),
            preImageHash = preImageHash,
            postImageHash = null,
            preExisted = true,
            timestampMs = 0L,
        )

        val fakeEventStore = FakeEventStore(
            listOf(storedEvent(workspaceBound, sessionId, 1L), storedEvent(fileWritten, sessionId, 2L)),
        )
        val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))

        val summary = service.undo(sessionId)

        assertEquals(1, summary.reverted)
        assertEquals(0, summary.rejected, "extra allowed path must not be jailed")
        assertEquals("FRESH", Files.readString(target))
    }

    /**
     * When multiple [SessionWorkspaceBoundEvent]s appear the LAST one wins.
     */
    @Test
    fun `undo uses last SessionWorkspaceBoundEvent when multiple are present`(): Unit = runBlocking {
        val bootDir = Files.createTempDirectory("undo-boot-multi").toRealPath()
        val firstWorkspace = Files.createTempDirectory("undo-ws-first").toRealPath()
        val lastWorkspace = Files.createTempDirectory("undo-ws-last").toRealPath()
        val target = lastWorkspace.resolve("last-ws-file.txt")
        Files.writeString(target, "MODIFIED")

        val fakeStore = FakeArtifactStore()
        val preImageHash = fakeStore.put("ORIGINAL".toByteArray()).value

        val sessionId = SessionId("session-multi-workspace")

        val firstBound = SessionWorkspaceBoundEvent(
            sessionId = sessionId,
            workspaceRoot = firstWorkspace.toString(),
            allowedPaths = emptyList(),
        )
        val lastBound = SessionWorkspaceBoundEvent(
            sessionId = sessionId,
            workspaceRoot = lastWorkspace.toString(),
            allowedPaths = emptyList(),
        )
        val fileWritten = FileWrittenEvent(
            invocationId = ToolInvocationId("inv-multi"),
            sessionId = sessionId,
            path = target.toString(),
            preImageHash = preImageHash,
            postImageHash = null,
            preExisted = true,
            timestampMs = 0L,
        )

        val fakeEventStore = FakeEventStore(
            listOf(
                storedEvent(firstBound, sessionId, 1L),
                storedEvent(lastBound, sessionId, 2L),
                storedEvent(fileWritten, sessionId, 3L),
            ),
        )
        val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))

        val summary = service.undo(sessionId)

        assertEquals(1, summary.reverted)
        assertEquals(0, summary.rejected, "file under last workspace must be reverted")
        assertEquals("ORIGINAL", Files.readString(target))
    }
}
