package com.defenseunicorns.peat_mesh.test

import android.os.Debug
import androidx.test.ext.junit.runners.AndroidJUnit4
import androidx.test.platform.app.InstrumentationRegistry
import com.defenseunicorns.peat.PeatJni
import java.io.File
import org.json.JSONObject
import org.junit.After
import org.junit.Assert.assertEquals
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Assert.fail
import org.junit.Test
import org.junit.runner.RunWith

/**
 * M4b (peat-mesh#138): the instrumented tests now drive peat-ffi's
 * production JNI surface (`PeatJni.*`) instead of a test-only shim.
 * The .so + the `external fun` stubs both ship from
 * `com.defenseunicorns:peat-ffi` on Maven Central; this file is the
 * Kotlin consumer that exercises them.
 *
 * What changed vs. M3:
 * - `TestJni.*` → `PeatJni.*` (different package, different method
 *   names — see android-tests/build.gradle.kts comments).
 * - `initBackend(formationKey)` → `createNodeJni(appId, sharedKey,
 *   storagePath)`: storage management moves to the Kotlin layer.
 * - `addPeerAddress` + `connectTo` → `connectPeerJni(handle, peer,
 *   address)`: peat-ffi populates iroh's lookup atomically.
 * - `forceSync` → `requestSyncJni` (same semantics — iterates
 *   connected peers and calls `sync_all_documents_with_peer`).
 * - Lost diagnostic helpers: `lastError`, `localDocCount`. peat-ffi
 *   doesn't expose either. Where they were used in M3's failure
 *   messages, those messages are now leaner. Failures still surface
 *   the test's own context (peer counts, JSON bodies).
 * - Hardcoded `SHARED_KEY` replaces `generateFormationKey()`.
 *   Both backends use the same key so their derived formation keys
 *   match → they trust each other and can sync.
 */
@RunWith(AndroidJUnit4::class)
class SyncProtocolTest {

    companion object {
        // Base64 of "test-key-1234567890123456789012345678901234".
        // Matches peat-ffi's own blob_tests test fixture; the exact
        // value doesn't matter as long as both backends use the
        // same value (so they derive the same formation key).
        private const val SHARED_KEY = "dGVzdC1rZXktMTIzNDU2Nzg5MDEyMzQ1Njc4OTAxMjM0"
        private const val APP_ID = "peat-mesh-m4b-test"
    }

    private data class BackendCtx(val handle: Long, val storageDir: File)

    private val backends = mutableListOf<BackendCtx>()

    @After
    fun teardown() {
        backends.forEach { ctx ->
            try {
                PeatJni.freeNodeJni(ctx.handle)
            } catch (t: Throwable) {
                // Best-effort cleanup; don't mask test failures with
                // teardown errors.
            }
            ctx.storageDir.deleteRecursively()
        }
        backends.clear()
    }

    /**
     * Create a fresh PeatNode with an isolated storage directory under
     * the instrumented test's cache dir. Returns the JNI handle so
     * subsequent calls can use it.
     */
    private fun createBackend(label: String): BackendCtx {
        val cacheDir = InstrumentationRegistry.getInstrumentation()
            .targetContext
            .cacheDir
        val storageDir = File(cacheDir, "peat-mesh-m4b-$label-${System.nanoTime()}").apply {
            mkdirs()
        }
        val handle = PeatJni.createNodeJni(APP_ID, SHARED_KEY, storageDir.absolutePath)
        assertTrue(
            "createNodeJni returned 0 for $label (Rust-side error; check logcat for PeatRust tag)",
            handle != 0L,
        )
        // peat-ffi's `create_node` doesn't auto-start sync; the
        // accept loop + observer forwarder come up only on
        // `start_sync()`. Without this, the receive side silently
        // never processes inbound sync messages — `deltaExchangeRoundTrip`
        // hangs at "B never received the doc", and bulk-publish tests
        // back up because the auto-sync push task never starts pushing.
        // (Discovered on peat-mesh#145 cb82dfb → 80d405c → this fix.)
        assertTrue(
            "startSyncJni failed for $label",
            PeatJni.startSyncJni(handle),
        )
        val ctx = BackendCtx(handle, storageDir)
        backends.add(ctx)
        return ctx
    }

    /**
     * Connect A → B and wait for the connection to register on both
     * sides. peat-ffi's `connectPeerJni` is atomic (populates iroh's
     * address lookup + dials in one call), so this is simpler than
     * M3's separate addPeerAddress+connectTo+waitForPeer.
     */
    private fun connectAndWaitSymmetric(a: BackendCtx, b: BackendCtx, timeoutMs: Long = 10_000) {
        val idB = PeatJni.nodeIdJni(b.handle)
        assertNotNull("B node id null", idB)
        val addrB = PeatJni.endpointSocketAddrJni(b.handle)
        assertNotNull("B endpoint socket addr null", addrB)

        assertTrue(
            "connectPeerJni(A→B) failed (idB=$idB addrB=$addrB)",
            PeatJni.connectPeerJni(a.handle, idB!!, addrB!!),
        )

        // Poll for symmetric connection. peat-ffi's `connectPeerJni`
        // returns true once iroh's dial succeeded on A's side, but B's
        // sync-ALPN accept happens asynchronously. Sync transfer needs
        // both sides to see each other → poll until peerCountJni ≥ 1
        // on both.
        val deadline = System.currentTimeMillis() + timeoutMs
        while (System.currentTimeMillis() < deadline) {
            if (PeatJni.peerCountJni(a.handle) >= 1 &&
                PeatJni.peerCountJni(b.handle) >= 1
            ) return
            Thread.sleep(100)
        }
        fail(
            "A and B did not reach symmetric connection in ${timeoutMs}ms; " +
                "A.peerCount=${PeatJni.peerCountJni(a.handle)} " +
                "B.peerCount=${PeatJni.peerCountJni(b.handle)}",
        )
    }

    // ------------------------------------------------------------------
    // M1 walking skeleton (carried over): proves the .so loads on the
    // device and a Java_*** symbol resolves. peat-ffi exposes
    // `peatVersion()` instead of M3's test-only `jniVersion()`.
    // ------------------------------------------------------------------
    @Test
    fun peatVersionIsReachable() {
        val version = PeatJni.peatVersion()
        assertTrue(
            "expected non-empty version string from peatVersion, got: '$version'",
            version.isNotEmpty(),
        )
    }

    // ------------------------------------------------------------------
    // M3 (#138) delta-exchange round-trip: A publishes, B receives.
    // LOCKS IN: peat-mesh#136 (post-convergence handling — exercised
    // transitively via requestSyncJni → sync_all_documents_with_peer
    // → initiate_sync → initiate_delta_sync).
    // ------------------------------------------------------------------
    @Test(timeout = 30_000)
    fun deltaExchangeRoundTrip() {
        val a = createBackend("delta-A")
        val b = createBackend("delta-B")
        connectAndWaitSymmetric(a, b)

        // `platforms` collection → SyncMode::LatestOnly per
        // peat-mesh's SyncModeRegistry::with_defaults — fire-and-forget
        // state snapshot, no bidirectional ALPN dance required.
        val collection = "platforms"
        val docId = "ANDROID-roundtrip"
        // The publishDocumentJni path expects the document id in the
        // JSON body (it's hoisted to Document::id during ingestion).
        val body = JSONObject()
            .put("id", docId)
            .put("name", "alpha")
            .put("severity", 3)
            .toString()

        val publishedId = PeatJni.publishDocumentJni(a.handle, collection, body)
        assertTrue(
            "publishDocumentJni on A returned empty id (signals an Err)",
            publishedId.isNotEmpty(),
        )
        assertEquals(
            "publishDocumentJni must return the id from the JSON body",
            docId,
            publishedId,
        )

        // Sanity: A reads back its own write.
        val onA = PeatJni.getDocumentJni(a.handle, collection, docId)
        assertNotNull("A doesn't see its own publish (null)", onA)
        assertEquals(
            "A's publish missing 'name' field. JSON: $onA",
            "alpha",
            JSONObject(onA!!).optString("name"),
        )

        assertTrue(
            "requestSyncJni on A failed",
            PeatJni.requestSyncJni(a.handle),
        )

        // Sync is push-based. Poll for actual field content on B,
        // re-triggering requestSync on each miss so a stalled sync
        // state advances.
        var fetched: String? = null
        val deadline = System.currentTimeMillis() + 15_000
        while (System.currentTimeMillis() < deadline) {
            fetched = PeatJni.getDocumentJni(b.handle, collection, docId)
            if (fetched != null) {
                val parsed = JSONObject(fetched)
                if (parsed.optString("name") == "alpha") break
            }
            PeatJni.requestSyncJni(a.handle)
            Thread.sleep(200)
        }
        assertNotNull(
            "B never received the doc within 15s. " +
                "A.peerCount=${PeatJni.peerCountJni(a.handle)} " +
                "B.peerCount=${PeatJni.peerCountJni(b.handle)}",
            fetched,
        )
        val parsed = JSONObject(fetched!!)
        assertEquals(
            "B's doc has wrong 'name' field. Raw: $fetched",
            "alpha",
            parsed.optString("name"),
        )
        assertEquals(
            "B's doc has wrong 'severity' field. Raw: $fetched",
            3,
            parsed.optInt("severity"),
        )
    }

    // ------------------------------------------------------------------
    // M3 (#138) OOM-class regression gate.
    // LOCKS IN: peat-mesh#134 (sync_all_documents_with_peer must not
    // materialize `Vec<(String, Automerge)>` via scan_prefix). The
    // structural pin in peat-mesh catches a source-level revert; this
    // test catches the runtime allocation pattern on real Android.
    // ------------------------------------------------------------------
    @Test(timeout = 180_000)
    fun scanDoesNotMaterializeBodies() {
        val a = createBackend("oom-A")
        val b = createBackend("oom-B")
        connectAndWaitSymmetric(a, b)

        // 30 docs × ~1 KB bodies under a LatestOnly collection. A
        // scan_prefix("") regression would Automerge::load every doc
        // into a single Vec for every requestSync; multi-GB peaks
        // observed pre-fix in field testing.
        val collection = "platforms"
        val docCount = 30
        val payload = "x".repeat(1024)
        for (i in 0 until docCount) {
            val body = JSONObject()
                .put("id", "doc-$i")
                .put("name", "ANDROID-$i")
                .put("payload", payload)
                .toString()
            val id = PeatJni.publishDocumentJni(a.handle, collection, body)
            assertTrue("publishDocumentJni(#$i) returned empty id", id.isNotEmpty())
        }

        // Settle: let publish-driven auto-sync drain before measuring.
        Thread.sleep(500)
        Runtime.getRuntime().gc()
        Thread.sleep(200)

        val heapBefore = Debug.getNativeHeapAllocatedSize()

        val cycles = 30
        for (i in 0 until cycles) {
            assertTrue("requestSyncJni cycle #$i failed", PeatJni.requestSyncJni(a.handle))
        }

        Runtime.getRuntime().gc()
        Thread.sleep(200)
        val heapAfter = Debug.getNativeHeapAllocatedSize()
        val deltaBytes = heapAfter - heapBefore
        val deltaMb = deltaBytes / (1024 * 1024)

        // 200 MB ceiling matches M3's threshold — well below the pre-
        // fix multi-GB regression, well above legitimate iroh +
        // Automerge sync-state scratch. Tighten in follow-ups once CI
        // data accumulates.
        val ceilingMb = 200L
        assertTrue(
            "native heap delta after $cycles requestSync cycles must stay under ${ceilingMb}MB; " +
                "measured ${deltaMb}MB (${deltaBytes} bytes). " +
                "Larger deltas indicate the scan_prefix → keys_with_prefix swap regressed.",
            deltaMb < ceilingMb,
        )
    }

    // ------------------------------------------------------------------
    // M4b carryover (from peat#879 QA round 1's deferred WARNING):
    // verify `getDocumentJni`'s `Err(_)` → `RuntimeException` branch
    // propagates across the JNI boundary. The in-crate Rust tests
    // cover `Ok(Some)` and `Ok(None)`; the throw path requires a JVM
    // and lands here.
    //
    // The earlier FS-deletion + cache-churn approach was non-
    // deterministic: Automerge's per-process LRU served the cached
    // copy and the test couldn't fail-closed on the throw contract.
    // peat-ffi 0.1.2 ships `forceStoreErrorForTestingJni` (peat#885)
    // that arms a one-shot flag — the next `getDocumentJni` call
    // unconditionally returns `Err(_)` from Rust, which propagates
    // as `RuntimeException` to Kotlin. The flag self-clears so
    // subsequent reads behave normally.
    // ------------------------------------------------------------------
    @Test(timeout = 30_000)
    fun getDocumentThrowsOnStoreError() {
        val ctx = createBackend("err")
        val collection = "platforms"
        val docId = "doc-err-target"

        // Arm the one-shot fault.
        assertTrue(
            "forceStoreErrorForTestingJni must return true on a valid handle",
            PeatJni.forceStoreErrorForTestingJni(ctx.handle),
        )

        // Next getDocumentJni call propagates Err(_) → RuntimeException.
        try {
            val result = PeatJni.getDocumentJni(ctx.handle, collection, docId)
            fail(
                "expected RuntimeException after arming fault injection; " +
                    "got result=$result instead",
            )
        } catch (e: RuntimeException) {
            // peat-ffi's `env.throw_new("java/lang/RuntimeException", ...)`
            // uses a getDocumentJni-prefixed message. Pin loosely so
            // a future reword doesn't break the test, but catches the
            // case where some OTHER RuntimeException slips through.
            assertNotNull("exception message must not be null", e.message)
            assertTrue(
                "RuntimeException did not carry the expected getDocumentJni prefix. " +
                    "Got: ${e.message}",
                e.message?.contains("getDocumentJni") == true,
            )
        }

        // Flag self-cleared: second call returns null (doc never existed),
        // not another throw.
        val secondResult = PeatJni.getDocumentJni(ctx.handle, collection, docId)
        assertNull(
            "second getDocumentJni after one-shot consumption must return null " +
                "(doc never existed; flag self-cleared)",
            secondResult,
        )
    }
}
