ai-app/app/androidApp/src/main/kotlin/com/example/aiapp/FrameStats.kt

package com.example.aiapp

import android.app.Activity
import android.content.Context
import android.content.ContextWrapper
import android.os.Build
import android.os.Handler
import android.os.HandlerThread
import android.view.FrameMetrics
import android.view.Window
import androidx.compose.runtime.Composable
import androidx.compose.runtime.DisposableEffect
import androidx.compose.ui.platform.LocalContext

/**
 * How long each frame took, and which phase of it, taken from the platform rather than from a frame
 * counter of our own.
 *
 * The point of splitting it up is that "the scroll is laggy" has two completely different causes
 * and one appearance. If the layout-and-measure and draw figures are small and the total is large,
 * the time is going into rasterising and compositing, and no amount of doing less work per row will
 * move it. If they are large, the work per row is the problem and it is ours to fix.
 *
 * The phases are the platform's own: [FrameMetrics] reports each frame's cost in nanoseconds,
 * broken into the parts the UI thread is responsible for and the parts after it.
 *
 * One of these for the app, like [DebugStats], because the two are read as one report and
 * [drawAccounting] divides one by the other. Held per screen it was emptied by leaving a session
 * and the counters were not, so a report copied after visiting two sessions divided every session's
 * work by the newest one's frame count -- 36.8 seconds of placement inside a 13.5 second window.
 */
object FrameStats {
    private val total = ArrayList<Long>()
    private val waited = ArrayList<Long>()
    private val input = ArrayList<Long>()
    private val animation = ArrayList<Long>()
    private val layout = ArrayList<Long>()
    private val draw = ArrayList<Long>()
    private val sync = ArrayList<Long>()
    private val issue = ArrayList<Long>()
    private val swap = ArrayList<Long>()
    private val gpu = ArrayList<Long>()
    private var since = System.currentTimeMillis()

    /**
     * Where a named phase of a scripted run (bench v2's fling/stream/type/keyboard) started, as an
     * index into [total] and a wall-clock time -- not a second recorder, just a mark on this one,
     * so a phase's frames are the same [FrameMetrics] the whole-run report already has, sliced.
     */
    private data class PhaseMark(val name: String, val startIndex: Int, val startMs: Long)

    private val phaseMarks = ArrayList<PhaseMark>()

    /** Call at the start of each named phase of a scripted run; see [BenchRun]. */
    @Synchronized
    fun markPhase(name: String) {
        phaseMarks += PhaseMark(name, total.size, System.currentTimeMillis())
    }

    @Synchronized
    fun add(metrics: FrameMetrics) {
        // The first frame after a window opens includes inflating it and is nobody's scroll.
        if (metrics.getMetric(FrameMetrics.FIRST_DRAW_FRAME) == 1L) return
        if (total.size >= CAP) return
        total += metrics.getMetric(FrameMetrics.TOTAL_DURATION)
        // How long the frame waited for the UI thread to be free before it could start. Reported
        // because the phases otherwise do not add up to the total, and the gap is the interesting
        // part: the frame being held up by work that is not the frame's.
        waited += metrics.getMetric(FrameMetrics.UNKNOWN_DELAY_DURATION)
        input += metrics.getMetric(FrameMetrics.INPUT_HANDLING_DURATION)
        animation += metrics.getMetric(FrameMetrics.ANIMATION_DURATION)
        layout += metrics.getMetric(FrameMetrics.LAYOUT_MEASURE_DURATION)
        draw += metrics.getMetric(FrameMetrics.DRAW_DURATION)
        sync += metrics.getMetric(FrameMetrics.SYNC_DURATION)
        issue += metrics.getMetric(FrameMetrics.COMMAND_ISSUE_DURATION)
        swap += metrics.getMetric(FrameMetrics.SWAP_BUFFERS_DURATION)
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.S) {
            gpu += metrics.getMetric(FrameMetrics.GPU_DURATION)
        }
    }

    @Synchronized
    fun reset() {
        listOf(total, waited, input, animation, layout, draw, sync, issue, swap, gpu).forEach {
            it.clear()
        }
        phaseMarks.clear()
        since = System.currentTimeMillis()
    }

    @Synchronized
    fun lines(refreshHz: Float): List<String> {
        if (total.isEmpty()) return listOf("  no frames recorded -- scroll first, then press this")
        val seconds = (System.currentTimeMillis() - since) / 1000.0
        val budget = if (refreshHz > 0) 1000.0 / refreshHz else 16.7
        val late = total.count { it / 1_000_000.0 > budget }
        return listOf(
            "  ${total.size} frames over ${"%.1f".format(seconds)}s" +
                " at ${"%.0f".format(refreshHz)}Hz (${"%.1f".format(budget)}ms budget)",
            "  late: $late (${percent(late, total.size)})" +
                if (total.size >= CAP) "  [capped]" else "",
            phase("total ", total),
            phase("waited", waited),
            phase("input ", input),
            phase("anim  ", animation),
            phase("layout", layout),
            phase("draw  ", draw),
            phase("sync  ", sync),
            phase("issue ", issue),
            phase("swap  ", swap),
        ) + if (gpu.isEmpty()) emptyList() else listOf(phase("gpu   ", gpu))
    }

    /**
     * One block per [markPhase] call: how many frames landed between that mark and the next (or the
     * end of the run, for the last one), how many were late, the total/p50/p90/p99, the worst
     * single frame, and how long the phase actually ran. Marks with no frames between them (a phase
     * that finished before a frame was drawn) still get a line rather than being silently dropped
     * -- UI_RULES' "say what you don't know" applies to a phase as much as to a single number.
     */
    @Synchronized
    fun phaseLines(refreshHz: Float): List<String> {
        if (phaseMarks.isEmpty()) return emptyList()
        val budget = if (refreshHz > 0) 1000.0 / refreshHz else 16.7
        val lines = ArrayList<String>()
        phaseMarks.forEachIndexed { i, mark ->
            val endIndex = if (i + 1 < phaseMarks.size) phaseMarks[i + 1].startIndex else total.size
            val endMs =
                if (i + 1 < phaseMarks.size) phaseMarks[i + 1].startMs
                else System.currentTimeMillis()
            val samples = total.subList(mark.startIndex, endIndex)
            val seconds = (endMs - mark.startMs) / 1000.0
            lines += "  ${mark.name}: ${samples.size} frames over ${"%.1f".format(seconds)}s"
            if (samples.isEmpty()) {
                lines += "    no frames recorded in this phase"
            } else {
                val late = samples.count { it / 1_000_000.0 > budget }
                lines += "    late: $late (${percent(late, samples.size)})"
                lines += "    " + phase("total ", samples)
                lines += "    worst ${"%.1fms".format(samples.max() / 1_000_000.0)}"
            }
        }
        return lines
    }

    /** How long the frames recorded here spent in their draw phase, and how many there were. */
    @Synchronized fun drawPhase(): Pair<Long, Int> = draw.sum() to draw.size

    private fun phase(name: String, samples: List<Long>): String {
        val sorted = samples.sorted()
        return "  $name p50 ${at(sorted, 50)}  p90 ${at(sorted, 90)}  p99 ${at(sorted, 99)}"
    }

    private fun at(sorted: List<Long>, percentile: Int): String {
        if (sorted.isEmpty()) return "-"
        val index = (sorted.size - 1) * percentile / 100
        return "%.1fms".format(sorted[index] / 1_000_000.0)
    }

    private fun percent(part: Int, whole: Int) = "%.1f%%".format(100.0 * part / whole)
}

/** Enough for a couple of minutes of scrolling; this is a diagnostic, not a log. */
private const val CAP = 20_000

/**
 * Records into [FrameStats] for as long as this screen is on it.
 *
 * The listener is what comes and goes; what it writes into does not, so a report covers the same
 * stretch of time as the counters beside it.
 *
 * The listener is handed its own thread because the platform calls it for every frame and the
 * documentation is explicit that doing that on the main thread taxes the very thing being measured.
 */
@Composable
fun RecordFrames() {
    val window = LocalContext.current.activity()?.window
    DisposableEffect(window) {
        if (window == null) return@DisposableEffect onDispose {}
        val thread = HandlerThread("frame-stats").apply { start() }
        val listener = Window.OnFrameMetricsAvailableListener { _, metrics, _ ->
            FrameStats.add(metrics)
        }
        window.addOnFrameMetricsAvailableListener(listener, Handler(thread.looper))
        onDispose {
            window.removeOnFrameMetricsAvailableListener(listener)
            thread.quitSafely()
        }
    }
}

/** The activity behind a composable's context, which is what owns the window. */
fun Context.activity(): Activity? {
    var context: Context? = this
    while (context is ContextWrapper) {
        if (context is Activity) return context
        context = context.baseContext
    }
    return null
}

/** What the display is actually refreshing at, so "late" is measured against the real budget. */
fun Context.refreshHz(): Float =
    @Suppress("DEPRECATION") (activity()?.windowManager?.defaultDisplay?.refreshRate ?: 60f)