package com.example.aiapp

import android.util.Log
import java.io.BufferedWriter
import java.io.File
import java.io.FileWriter
import java.io.IOException
import java.io.RandomAccessFile

/**
 * This phone's copy of the transcripts it has already been sent, so reopening a session does not
 * download it again.
 *
 * What is stored is the server's own JSON for one event per line, in transcript order -- the
 * elements of a `/transcript` page and the payload of each SSE frame. Reading the cache means
 * running the same [parseSeqEvent] the network path runs, so a cached transcript and a fetched one
 * cannot draw differently, and an event type this build does not know ([SessionEvent.Unknown])
 * keeps every field it arrived with, on disk, for the build that will. Rows are deliberately *not*
 * what is stored: a row is a rendering of events, its shape changes whenever the fold does, and a
 * cache of rows would need throwing away on every app update that touched `foldEvent`.
 *
 * See TRANSCRIPT_CACHE.md for the design. Four rules run through all of it:
 * 1. what is on screen is what the server's transcript says, in order, with nothing missing -- the
 *    cache is a copy and is never inferred, folded or edited here;
 * 2. a cached line is never ahead of the live cursor, and the cursor never ahead of the cache;
 * 3. the cache is never load-bearing -- missing, evicted, damaged or unwritable all degrade to a
 *    cold open, never to a blank or a wrong screen;
 * 4. a line already on the phone is not fetched again.
 *
 * A plain [File] root and no Compose, `Context` or network, so the whole of the file logic runs
 * under the JVM unit tests. It is also why there is no JSON parser in here: what it needs off a
 * line is the sequence number and whether the line is a streamed delta, and both are read with a
 * regex over text the server wrote. A line it cannot read that way is treated as damage, which
 * gives the same answer as having no cache at all.
 *
 * [warn] is where failures are said, for the same reason -- `android.util.Log` is a stub that
 * throws under the JVM tests, and this file has to be exercisable there.
 */
class TranscriptCache(
    private val root: File,
    private val warn: (String) -> Unit = { Log.w("ai-app", it) },
) {
    /** The cache for one session, whether or not anything has been stored for it yet. */
    fun session(id: String): SessionCache = SessionCache(File(root, id), warn)

    /**
     * Deletes every session directory not in [ids], called after a successful list fetch.
     *
     * The path out for a session deleted on another device or at the backend: nothing here would
     * otherwise ever hear about it, and unlike a draft's few bytes what it leaves behind is
     * megabytes.
     */
    fun retainOnly(ids: Set<String>) =
        guardIo(Unit, warn) {
            sessionDirs().forEach { if (it.name !in ids) it.deleteRecursively() }
        }

    /**
     * Deletes least-recently-touched session directories, never [keep], until the whole of this
     * server's cache is under [budget].
     *
     * Least-recently-touched rather than largest: what a reader is likely to open again is what
     * they opened last, and evicting the big ones first would empty the cache for exactly the
     * conversations it exists for.
     */
    fun evictToBudget(keep: String, budget: Long = CACHE_BUDGET_BYTES) =
        guardIo(Unit, warn) {
            val dirs = sessionDirs().sortedBy { it.lastModified() }
            var total = dirs.sumOf { sizeOf(it) }
            for (dir in dirs) {
                if (total <= budget) break
                if (dir.name == keep) continue
                val was = sizeOf(dir)
                if (dir.deleteRecursively()) total -= was
            }
        }

    fun purgeAll() = guardIo(Unit, warn) { root.deleteRecursively() }

    private fun sessionDirs(): List<File> = root.listFiles()?.filter { it.isDirectory }.orEmpty()
}

/**
 * How much of this phone's cache directory all of one server's transcripts may take.
 *
 * A dozen of the largest transcripts seen in the dev VM (21 MB for 24,000 events) and a small
 * fraction of a phone. A number to revisit against real use rather than a measurement of anything.
 */
const val CACHE_BUDGET_BYTES: Long = 256L * 1000 * 1000

/**
 * What the newest cached line says, which is what the probe checks against the server.
 *
 * Both halves are wanted together and by the same caller: the seq is what the request asks about,
 * and the line is what its answer is compared with.
 */
data class CachedTail(val seq: Long, val line: String)

/**
 * One session's cached lines, as a directory of chunks.
 *
 * A chunk is a set of lines *and a claim about what they cover*, and the two are not the same
 * thing: a coalesced page joins each run of streamed deltas into one event carrying the seq of the
 * run's oldest delta, so a page whose newest event is seq 1,200 may in fact cover everything up to
 * the 1,650 it was fetched with, and nothing in the lines says so. So coverage is the half-open
 * range in the file's name:
 * ```
 *   <first>-<end>.rows.jsonl     a coalesced page; end is the `before` it was fetched with
 *   <first>-<end>.raw.jsonl      an uncoalesced page, or a closed live run
 *   <first>-open.raw.jsonl       the live run; end is its last line's seq + 1
 * ```
 *
 * Two chunks are adjacent when one's `end` is the other's `first`. Only the contiguous run of
 * adjacent chunks ending at the newest chunk -- the **suffix** -- is ever served: chunks behind a
 * gap are kept, because the gap is usually closed by paging back through it, but nothing is served
 * across one.
 *
 * **The newest chunk is always raw**, which is what makes the stream cursor and the probe well
 * defined -- a raw chunk's last line is a real event at a real seq, and the server never coalesces
 * the newest window. It holds by construction (the opening window and every stream frame are raw)
 * and is checked on read: a `.rows` chunk at the newest end can only mean this app died between
 * closing one live run and opening the next, and it discards the session.
 *
 * Nothing here is load-bearing. Every operation that touches the disk answers as though the cache
 * were empty when it cannot, and a write failure disables writing for the rest of this instance's
 * life so that a full disk costs one log line rather than one per delta.
 *
 * Every operation is synchronized, because two of them really do run at once: the stream appends
 * live events from its own IO thread while a reader scrolling back reads pages from another. The
 * lock is uncontended in the ordinary case and what it buys is that the open chunk's name, its end
 * and its writer are never read half-rotated -- which would show up as a page silently fetched
 * again, or as a stored chunk overlapping the run it was written beside.
 */
class SessionCache(
    private val dir: File,
    private val warn: (String) -> Unit = { Log.w("ai-app", it) },
) {
    /** Set by the first write that fails: a second would fail the same way, once per delta. */
    private var disabled = false
    /**
     * The open chunk's writer, its file, and the seq that chunk now ends at.
     *
     * Buffered, and flushed on [flush], because a delta is a hundred bytes and arrives dozens of
     * times a second while a reply streams -- a syscall each is the thing to avoid. What that costs
     * is the unflushed tail on a crash, which is safe: a shorter cache is a longer catch-up, never
     * a wrong one.
     */
    private var writer: BufferedWriter? = null
    private var openFile: File? = null
    private var openEnd: Long = 0

    /**
     * The newest line of the suffix, or null when there is none or the newest chunk is not raw.
     *
     * This is the cursor the live stream would resume from, so it is also what has to be shown to
     * still be the server's own line before anything is resumed from it -- see
     * `TranscriptSource.probe`.
     */
    @Synchronized
    fun tail(): CachedTail? =
        guard(null) {
            val newest = suffix().lastOrNull() ?: return@guard null
            var found: CachedTail? = null
            eachLine(newest) { line ->
                found = CachedTail(seqOf(line)!!, line)
                false
            }
            found
        }

    /** The newest [limit] lines of the suffix, oldest first -- the opening window. */
    @Synchronized
    fun newest(limit: Int): List<String> =
        guard(emptyList()) {
            val taken = ArrayDeque<String>()
            for (chunk in suffix().asReversed()) {
                if (taken.size >= limit) break
                eachLine(chunk) { line ->
                    taken.addFirst(line)
                    taken.size < limit
                }
            }
            taken.toList()
        }

    /**
     * The page of lines before [before], oldest first, or null when the cache cannot answer.
     *
     * Null is a miss -- the suffix does not cover the ground immediately below [before] -- and
     * means the server has to be asked. It is deliberately not an empty list: an empty page is how
     * the screen is told it has reached the start of the conversation, and a cache saying that of
     * history it merely does not hold would stop the transcript scrolling back for good.
     *
     * [before] is anywhere inside the suffix, not only at a chunk boundary. The cursor a warm open
     * leaves behind is in the middle of the live run -- the screen draws the newest eighty lines of
     * it -- so a cache that could only answer at a boundary would send the very first backwards
     * page to the server and, since that page would overlap the run, keep none of it.
     *
     * A short page is fine, and is what a walk that reaches the oldest chunk of the suffix returns:
     * the caller already treats a short page as a page.
     *
     * With [rows] the count is rows rather than lines, mirroring the server's `parse_coalesced`:
     * every event that is not a streamed delta is a row, and each maximal run of deltas is one row.
     * The deltas are not joined here -- `foldEvent` does that, and the joined row keeps the seq of
     * its first delta either way, so anchors and the next `before` land where they do today.
     */
    @Synchronized
    fun page(before: Long, limit: Int, rows: Boolean): List<String>? =
        guard(null) {
            val suffix = suffix()
            val newest = suffix.lastOrNull() ?: return@guard null
            // Above what is held, or at or below where it starts: either way the run the caller
            // is scrolling into is not continuous with this one, and only the server has it.
            if (before > newest.end || before <= suffix.first().first) return@guard null
            val taken = ArrayDeque<String>()
            var counted = 0
            var inRun = false
            var wanting = true
            for (chunk in suffix.asReversed()) {
                if (!wanting) break
                if (chunk.first >= before) continue
                eachLine(chunk) { line ->
                    // The page is what is *before* the cursor; the rows at or above it are the
                    // ones already on screen.
                    if (seqOf(line)!! >= before) return@eachLine true
                    if (rows) {
                        val delta = isDelta(line)
                        // Stop only between rows: a delta continuing the run being gathered is
                        // part of a row already counted, and breaking on it would drop the half
                        // of that row already taken.
                        if (counted >= limit && !(delta && inRun)) wanting = false
                        else {
                            if (!delta || !inRun) counted++
                            inRun = delta
                        }
                    } else if (taken.size >= limit) {
                        wanting = false
                    }
                    if (wanting) taken.addFirst(line)
                    wanting
                }
            }
            taken.toList()
        }

    /**
     * The `end` of the nearest chunk at or below [before], which is the floor a fetched page is
     * asked with so that it stops where this phone's copy starts. Null when there is no such chunk.
     *
     * Any chunk, not only the suffix's: the whole point is to reach the run behind a gap, so that
     * the gap is closed with exactly the bytes it is wide and the history behind it is served
     * locally from then on.
     */
    @Synchronized
    fun coveredUpTo(before: Long): Long? =
        guard(null) { chunks().map { it.end }.filter { it <= before }.maxOrNull() }

    /**
     * Stores a fetched page covering `[first, end)`; false when it was not stored.
     *
     * Refused when it overlaps a chunk already here, because there is no clean cut: a coalesced
     * event cannot be split at a seq inside its own delta run. `TranscriptSource` keeps that from
     * arising by bounding what it fetches, and this is the guard for a page that arrives anyway --
     * from a server without the `after` parameter, say. Such a page is still drawn; it is only not
     * kept.
     *
     * The newest chunk is never stored through here: the opening window and every live frame go
     * through [append], which is what keeps the newest chunk raw and open.
     */
    @Synchronized
    fun storePage(lines: List<String>, first: Long, end: Long, rows: Boolean): Boolean =
        guard(false) {
            if (disabled || lines.isEmpty() || end <= first) return@guard false
            if (chunks().any { first < it.end && it.first < end }) return@guard false
            dir.mkdirs()
            val kind = if (rows) "rows" else "raw"
            File(dir, "$first-$end.$kind.jsonl").writeText(lines.joinToString("\n", postfix = "\n"))
            true
        }

    /**
     * Appends one live event, which is also how a freshly fetched opening window is stored.
     *
     * A seq equal to the open chunk's end extends it. A larger one is a gap -- which is what a
     * `reset` looks like from here -- and closes the open chunk under the end it turned out to have
     * before starting a new one at [seq]. A smaller one is already covered and is ignored; the SSE
     * contract is `seq > after`, so that is a guard rather than a path.
     */
    @Synchronized
    fun append(line: String, seq: Long) =
        guard(Unit) {
            if (disabled) return@guard
            val writer = writerFor(seq) ?: return@guard
            // Written as it arrived. A newline inside it would split one event into two
            // unreadable halves, but neither source can produce one: SSE framing forbids it, and
            // a page's elements are re-serialized compactly, which escapes it.
            writer.write(line)
            writer.write("\n")
            openEnd = seq + 1
        }

    /**
     * Flushes what [append] has buffered. Called on each `Status` event -- the boundaries of a
     * turn, which is the granularity a crash may as well lose -- and when the stream closes.
     */
    @Synchronized fun flush() = guard(Unit) { writer?.flush() }

    /** What [purge] would discard, for the reload row in session settings. */
    @Synchronized fun bytes(): Long = guard(0L) { sizeOf(dir) }

    /** Marks this session as visited, which is what eviction ranks by. */
    @Synchronized
    fun touch() =
        guard(Unit) { if (dir.isDirectory) dir.setLastModified(System.currentTimeMillis()) }

    @Synchronized
    fun purge() =
        guard(Unit) {
            closeWriter()
            dir.deleteRecursively()
        }

    // -- chunks ------------------------------------------------------------------------------

    private data class Chunk(val file: File, val first: Long, val end: Long, val open: Boolean) {
        val rows: Boolean
            get() = file.name.endsWith(".rows.jsonl")
    }

    /**
     * Every chunk on disk, oldest first. A name this does not recognise is not ours and is ignored.
     *
     * Recomputed per operation rather than kept: another operation may have changed the directory,
     * and a hundred names is a directory listing.
     */
    private fun chunks(): List<Chunk> {
        writer?.flush()
        return dir.listFiles()
            .orEmpty()
            .mapNotNull { file ->
                val match = CHUNK_NAME.matchEntire(file.name) ?: return@mapNotNull null
                val first = match.groupValues[1].toLongOrNull() ?: return@mapNotNull null
                val open = match.groupValues[2] == "open"
                val end = if (open) openEndOf(file, first) else match.groupValues[2].toLongOrNull()
                // A chunk covering nothing is one that was created and never written to -- an
                // append whose very first write failed. It says nothing, so it is not a chunk.
                if (end == null || end <= first) null else Chunk(file, first, end, open)
            }
            .sortedBy { it.first }
    }

    /**
     * The open chunk's end: its last line's seq plus one, or the in-memory end while this instance
     * is the one writing it.
     *
     * An open chunk whose last line cannot be read is this app having died mid-write. That line is
     * dropped and the file truncated to the last good one before anything is served from it, which
     * is the one place damage is repaired rather than discarded: the tail of an append-only file is
     * the only place a partial line can be.
     */
    private fun openEndOf(file: File, first: Long): Long {
        if (openFile == file && openEnd > 0) return openEnd
        repairTail(file)
        var end = first
        eachLineBackwards(file) { _, line ->
            seqOf(line)?.let { end = it + 1 }
            false
        }
        return end
    }

    /**
     * The contiguous run of adjacent chunks ending at the newest one, oldest first.
     *
     * A newest chunk that is not raw cannot happen while this code is the only writer, and means
     * the directory is not to be trusted -- so the session is discarded rather than served across
     * whatever else is wrong with it.
     */
    private fun suffix(): List<Chunk> {
        val all = chunks()
        var index = all.size - 1
        val newest = all.lastOrNull() ?: return emptyList()
        if (newest.rows) throw Damaged(newest.file)
        val run = ArrayDeque<Chunk>()
        run.addFirst(newest)
        while (index > 0 && all[index - 1].end == run.first().first) {
            index--
            run.addFirst(all[index])
        }
        return run.toList()
    }

    /**
     * Each line of [chunk], newest first, until [take] says stop.
     *
     * Backwards and lazily, because every question this cache is asked is about the newest end --
     * the tail, the opening window, the page before a cursor -- and a live run grows to the size of
     * the conversation. Reading the file whole to answer with eighty lines of it is the cost the
     * server's own reader was rewritten to stop paying.
     *
     * Damage anywhere but at the tail of the open chunk was not written by this code, and there is
     * no honest way to say what a chunk covers with a line of it unreadable -- so it discards the
     * session rather than serving what it can read.
     */
    private fun eachLine(chunk: Chunk, take: (String) -> Boolean) {
        eachLineBackwards(chunk.file) { _, line ->
            if (seqOf(line) == null) throw Damaged(chunk.file)
            take(line)
        }
    }

    // -- writing -----------------------------------------------------------------------------

    /** The writer for the chunk [seq] belongs in, opening or rotating one as it has to. */
    private fun writerFor(seq: Long): BufferedWriter? {
        writer?.let { held ->
            if (seq == openEnd) return held
            if (seq < openEnd) return null
            // A gap: what this instance has written covers up to `openEnd`, and that is the name
            // the chunk gets before a new one starts at the arriving seq.
            closeOpenChunk(openEnd)
        }
        dir.mkdirs()
        // An open chunk left by an earlier instance, or by an earlier screen.
        chunks()
            .lastOrNull { it.open }
            ?.let { existing ->
                if (seq < existing.end) return null
                if (seq == existing.end) {
                    openFile = existing.file
                    openEnd = existing.end
                    return FileWriter(existing.file, true).buffered().also { writer = it }
                }
                rename(existing.file, existing.first, existing.end)
            }
        // A chunk that was created and never written to would otherwise be left behind under a
        // name a second one is about to want; it covers nothing, so nothing is lost with it.
        dir.listFiles().orEmpty().forEach {
            if (CHUNK_NAME.matchEntire(it.name)?.groupValues?.get(2) == "open" && it.length() == 0L)
                it.delete()
        }
        val file = File(dir, "$seq-open.raw.jsonl")
        openFile = file
        openEnd = seq
        return FileWriter(file, false).buffered().also { writer = it }
    }

    /** Renames the open chunk to the range it turned out to cover, so it stops being open. */
    private fun closeOpenChunk(end: Long) {
        val file = openFile
        closeWriter()
        if (file == null) return
        val first = CHUNK_NAME.matchEntire(file.name)?.groupValues?.get(1)?.toLongOrNull()
        if (first != null) rename(file, first, end)
    }

    private fun rename(file: File, first: Long, end: Long) {
        file.renameTo(File(dir, "$first-$end.raw.jsonl"))
    }

    private fun closeWriter() {
        try {
            writer?.close()
        } catch (_: IOException) {
            // Nothing left to do about it: the file is what it is, and the read path repairs a
            // half-written tail.
        }
        writer = null
        openFile = null
        openEnd = 0
    }

    // -- failure -----------------------------------------------------------------------------

    /** A chunk that cannot be read as what its name claims. */
    private class Damaged(val file: File) : RuntimeException()

    /**
     * Runs [body], answering [ifBroken] when the directory cannot give a real answer.
     *
     * None of this is reported on screen: none of it changes what the screen shows -- every read
     * here has a network path beside it producing the same result -- and the reader has nothing to
     * do about it. It is logged, and damage discards this session's cache, which is what makes the
     * next open an ordinary cold one.
     */
    private fun <T> guard(ifBroken: T, body: () -> T): T =
        // A disk that refused once will refuse again, once per delta, so the first refusal is
        // also the last: this instance stops writing rather than logging a line a token.
        guardIo(
            ifBroken,
            warn,
            onFailure = {
                disabled = true
                closeWriter()
            },
        ) {
            try {
                body()
            } catch (e: Damaged) {
                warn("transcript cache damaged at ${e.file}; discarding ${dir.name}")
                closeWriter()
                dir.deleteRecursively()
                ifBroken
            }
        }
}

/** `<first>-<end|open>.<rows|raw>.jsonl`; anything else in the directory is not ours. */
private val CHUNK_NAME = Regex("""^(\d+)-(\d+|open)\.(rows|raw)\.jsonl$""")

private val SEQ_IN_LINE = Regex(""""seq"\s*:\s*(\d+)""")
private val TYPE_IN_LINE = Regex(""""type"\s*:\s*"([^"]*)"""")

/**
 * One line's sequence number, or null when the line is not one of ours.
 *
 * A regex rather than a JSON parse, so that this file carries no parser and runs under the JVM
 * tests: the seq is the first field the server writes (`SeqEvent`'s declaration order, with the
 * event flattened after it), so the first match is the top-level one.
 */
private fun seqOf(line: String): Long? = SEQ_IN_LINE.find(line)?.groupValues?.get(1)?.toLongOrNull()

/** Whether a line is one streamed piece of a reply, which is what makes a run of them one row. */
private fun isDelta(line: String): Boolean =
    TYPE_IN_LINE.find(line)?.groupValues?.get(1) == "assistantText"

/**
 * How much of a file is read at a time when walking it backwards. One block covers a page of a
 * transcript comfortably, and the walk stops as soon as the caller has what it asked for.
 */
private const val READ_BLOCK = 64 * 1024

/**
 * Calls [onLine] with each non-blank line of [file], **newest first**, along with the byte offset
 * it starts at, until [onLine] answers false.
 *
 * Every question the cache is asked is about the newest end of a chunk, and a live run reaches the
 * size of the conversation, so reading forwards means reading a transcript to answer with the last
 * eighty lines of it. This reads blocks from the end and stops where the caller stops.
 *
 * Splitting on bytes is safe because the separator is `\n`, which cannot occur inside a multi-byte
 * UTF-8 sequence; each line is decoded whole, so nothing is cut through a character. A missing file
 * yields nothing, which is the same answer as an empty one.
 */
private fun eachLineBackwards(file: File, onLine: (offset: Long, line: String) -> Boolean) {
    if (!file.isFile) return
    RandomAccessFile(file, "r").use { handle ->
        // Bytes below `unread` have not been looked at; `pending` is the oldest line so far, which
        // is incomplete until a newline is found before it in an older block.
        var unread = handle.length()
        var pending = ByteArray(0)
        while (unread > 0) {
            val take = minOf(READ_BLOCK.toLong(), unread).toInt()
            val start = unread - take
            val block = ByteArray(take)
            handle.seek(start)
            handle.readFully(block)
            val buffer = if (pending.isEmpty()) block else block + pending
            var lineEnd = buffer.size
            var at = buffer.size - 1
            while (at >= 0) {
                if (buffer[at] == NEWLINE) {
                    val line = String(buffer, at + 1, lineEnd - at - 1, Charsets.UTF_8)
                    if (line.isNotBlank() && !onLine(start + at + 1, line)) return
                    lineEnd = at
                }
                at--
            }
            pending = buffer.copyOfRange(0, lineEnd)
            unread = start
        }
        // The first line of a file has no newline before it to be found.
        val first = String(pending, Charsets.UTF_8)
        if (first.isNotBlank()) onLine(0, first)
    }
}

private const val NEWLINE = '\n'.code.toByte()

/**
 * Drops a final line that is not one of ours, by truncating the file to where it starts.
 *
 * This app having died mid-write is the one kind of damage that is repaired rather than discarded:
 * the tail of an append-only file is the only place a partial line can be, and everything before it
 * is intact. A second bad line is not this, and is left for the read path to notice.
 */
private fun repairTail(file: File) {
    var truncateTo = -1L
    eachLineBackwards(file) { offset, line ->
        if (seqOf(line) == null) truncateTo = offset
        false
    }
    if (truncateTo >= 0) RandomAccessFile(file, "rw").use { it.setLength(truncateTo) }
}

private fun sizeOf(file: File): Long =
    if (file.isDirectory) file.listFiles().orEmpty().sumOf { sizeOf(it) } else file.length()

/**
 * The disk half of [SessionCache.guard], shared with [TranscriptCache]'s own maintenance.
 *
 * [onFailure] is what the caller does about it beyond answering [ifBroken] -- for a session's
 * cache, giving up on writing.
 */
private fun <T> guardIo(
    ifBroken: T,
    warn: (String) -> Unit,
    onFailure: () -> Unit = {},
    body: () -> T,
): T =
    try {
        body()
    } catch (e: IOException) {
        warn("transcript cache unusable: ${e.message}")
        onFailure()
        ifBroken
    } catch (e: SecurityException) {
        warn("transcript cache unreadable: ${e.message}")
        onFailure()
        ifBroken
    }