package {{PACKAGE_NAME}}

import android.app.Notification
import android.app.NotificationChannel
import android.app.NotificationManager
import android.app.Service
import android.content.Intent
import android.os.Build
import android.os.Bundle
import android.os.Debug
import android.os.Handler
import android.os.IBinder
import android.os.Looper
import android.os.ResultReceiver
import android.widget.TextView
import androidx.appcompat.app.AppCompatActivity
import com.sun.jna.Library
import com.sun.jna.Native
import com.sun.jna.Pointer
import com.sun.jna.Structure
import org.json.JSONArray
import org.json.JSONObject

private const val DEFAULT_FUNCTION = "{{DEFAULT_FUNCTION}}"
private const val DEFAULT_ITERATIONS = 20u
private const val DEFAULT_WARMUP = 3u
private const val FUNCTION_EXTRA = "bench_function"
private const val ITERATIONS_EXTRA = "bench_iterations"
private const val WARMUP_EXTRA = "bench_warmup"
private const val SPEC_ASSET = "bench_spec.json"
private const val RUN_BENCHMARK_ACTION = "{{PACKAGE_NAME}}.RUN_BENCHMARK"
private const val RESULT_RECEIVER_EXTRA = "bench_result_receiver"
private const val RESULT_DISPLAY_EXTRA = "bench_display"
private const val RESULT_ERROR_EXTRA = "bench_error"
private const val RESULT_OK = 1
private const val RESULT_ERROR = 2
private const val FOREGROUND_NOTIFICATION_ID = 1001
private const val FOREGROUND_CHANNEL_ID = "mobench_benchmark"
private const val FOREGROUND_CHANNEL_NAME = "Mobench benchmark"

class MainActivity : AppCompatActivity() {
    @Volatile private var benchmarkComplete = false

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContentView(R.layout.activity_main)

        val resultText = findViewById<TextView>(R.id.result_text)
        resultText?.text = "Running benchmark..."

        val params = resolveBenchParams()
        val resultReceiver = object : ResultReceiver(Handler(Looper.getMainLooper())) {
            override fun onReceiveResult(resultCode: Int, resultData: Bundle?) {
                val display = resultData?.getString(RESULT_DISPLAY_EXTRA)
                    ?: resultData?.getString(RESULT_ERROR_EXTRA)
                    ?: "Benchmark worker returned no result"
                resultText?.text = display
                benchmarkComplete = true

                if (resultCode == RESULT_OK) {
                    android.util.Log.i("BenchRunner", "Benchmark worker completed")
                } else {
                    android.util.Log.e("BenchRunner", display)
                }
            }
        }

        try {
            val intent = Intent(this, BenchmarkWorkerService::class.java).apply {
                action = RUN_BENCHMARK_ACTION
                putExtra(FUNCTION_EXTRA, params.function)
                putExtra(ITERATIONS_EXTRA, params.iterations.toInt())
                putExtra(WARMUP_EXTRA, params.warmup.toInt())
                putExtra(RESULT_RECEIVER_EXTRA, resultReceiver)
            }
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
                startForegroundService(intent)
            } else {
                startService(intent)
            }
        } catch (e: Exception) {
            android.util.Log.e("BenchRunner", "Failed to run benchmark worker", e)
            resultText?.text = "Failed to run benchmark worker: ${e.message}"
            benchmarkComplete = true
        }
    }

    fun isBenchmarkComplete(): Boolean = benchmarkComplete

    private fun resolveBenchParams(): BenchParams {
        val assetParams = loadBenchParamsFromAssets()
        val defaults = assetParams ?: BenchParams(DEFAULT_FUNCTION, DEFAULT_ITERATIONS, DEFAULT_WARMUP)

        val intentFunction = intent?.getStringExtra(FUNCTION_EXTRA)?.takeUnless { it.isBlank() }
        val intentIterations = intent?.let {
            val value = it.getIntExtra(ITERATIONS_EXTRA, -1)
            if (value >= 0) value.toUInt() else null
        }
        val intentWarmup = intent?.let {
            val value = it.getIntExtra(WARMUP_EXTRA, -1)
            if (value >= 0) value.toUInt() else null
        }

        val fn = intentFunction ?: defaults.function
        val iterations = intentIterations ?: defaults.iterations
        val warmup = intentWarmup ?: defaults.warmup
        android.util.Log.i("BenchRunner", "Resolved params: function=$fn, iterations=$iterations, warmup=$warmup")
        return BenchParams(fn, iterations, warmup)
    }

    private fun loadBenchParamsFromAssets(): BenchParams? {
        return try {
            val raw = assets.open(SPEC_ASSET).bufferedReader().use { it.readText() }
            if (raw.isBlank()) {
                android.util.Log.w("BenchRunner", "bench_spec.json exists but is empty, using defaults")
                null
            } else {
                val json = JSONObject(raw)
                val function = json.optString("function", DEFAULT_FUNCTION)
                val iterations = json.optInt("iterations", DEFAULT_ITERATIONS.toInt()).toUInt()
                val warmup = json.optInt("warmup", DEFAULT_WARMUP.toInt()).toUInt()
                android.util.Log.i("BenchRunner", "Loaded config from bench_spec.json: function=$function, iterations=$iterations, warmup=$warmup")
                BenchParams(function, iterations, warmup)
            }
        } catch (e: java.io.FileNotFoundException) {
            android.util.Log.d("BenchRunner", "No bench_spec.json in assets, will use intent extras or defaults")
            null
        } catch (e: Exception) {
            android.util.Log.e("BenchRunner", "Failed to parse bench_spec.json from assets", e)
            null
        }
    }
}

class BenchmarkWorkerService : Service() {
    override fun onBind(intent: Intent?): IBinder? = null

    override fun onStartCommand(intent: Intent?, flags: Int, startId: Int): Int {
        if (intent?.action != RUN_BENCHMARK_ACTION) {
            stopSelf(startId)
            return START_NOT_STICKY
        }

        val resultReceiver = intent.resultReceiverExtra()
        val params = BenchParams(
            function = intent.getStringExtra(FUNCTION_EXTRA)?.takeUnless { it.isBlank() } ?: DEFAULT_FUNCTION,
            iterations = intent.getIntExtra(ITERATIONS_EXTRA, DEFAULT_ITERATIONS.toInt()).toUInt(),
            warmup = intent.getIntExtra(WARMUP_EXTRA, DEFAULT_WARMUP.toInt()).toUInt(),
        )

        startBenchmarkForeground()
        Thread {
            try {
                val result = runBenchmarkInWorker(params)
                val bundle = Bundle().apply {
                    putString(RESULT_DISPLAY_EXTRA, result.displayText)
                    result.errorMessage?.let { putString(RESULT_ERROR_EXTRA, it) }
                }
                resultReceiver?.send(if (result.errorMessage == null) RESULT_OK else RESULT_ERROR, bundle)
            } finally {
                stopBenchmarkForeground()
                stopSelf(startId)
            }
        }.apply {
            name = "mobench-benchmark-worker"
            start()
        }

        return START_NOT_STICKY
    }

    private fun startBenchmarkForeground() {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
            val manager = getSystemService(NotificationManager::class.java)
            manager?.createNotificationChannel(
                NotificationChannel(
                    FOREGROUND_CHANNEL_ID,
                    FOREGROUND_CHANNEL_NAME,
                    NotificationManager.IMPORTANCE_LOW
                )
            )
        }
        startForeground(FOREGROUND_NOTIFICATION_ID, buildBenchmarkNotification())
    }

    @Suppress("DEPRECATION")
    private fun stopBenchmarkForeground() {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
            stopForeground(Service.STOP_FOREGROUND_REMOVE)
        } else {
            stopForeground(true)
        }
    }

    @Suppress("DEPRECATION")
    private fun buildBenchmarkNotification(): Notification {
        val builder = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
            Notification.Builder(this, FOREGROUND_CHANNEL_ID)
        } else {
            Notification.Builder(this).setPriority(Notification.PRIORITY_LOW)
        }

        return builder
            .setSmallIcon(android.R.drawable.ic_dialog_info)
            .setContentTitle("Running benchmark")
            .setContentText("Mobench isolated worker is measuring this run")
            .setOngoing(true)
            .setOnlyAlertOnce(true)
            .setShowWhen(false)
            .setCategory(Notification.CATEGORY_SERVICE)
            .build()
    }

    private fun runBenchmarkInWorker(params: BenchParams): WorkerBenchmarkResult {
        val display = try {
            val processMemorySampler = ProcessMemorySampler()
            var runProcessPeakMemoryKb: Long?
            processMemorySampler.start()
            val rawReport = try {
                NativeBenchRunner.run(params)
            } finally {
                runProcessPeakMemoryKb = processMemorySampler.stop()
            }
            val json = buildNativeBenchReportJson(rawReport, runProcessPeakMemoryKb)
            android.util.Log.i("BenchRunner", "BENCH_JSON ${json}")
            formatNativeBenchReport(json)
        } catch (e: Exception) {
            android.util.Log.e("BenchRunner", "Benchmark error: ${e.message}", e)
            return WorkerBenchmarkResult(
                displayText = "Benchmark error: ${e.message}",
                errorMessage = "Benchmark error: ${e.message}"
            )
        }

        return WorkerBenchmarkResult(displayText = display, errorMessage = null)
    }
}

private data class BenchParams(
    val function: String,
    val iterations: UInt,
    val warmup: UInt,
)

private data class WorkerBenchmarkResult(
    val displayText: String,
    val errorMessage: String?,
)

private object NativeBenchRunner {
    fun run(params: BenchParams): JSONObject {
        val spec = JSONObject()
            .put("name", params.function)
            .put("iterations", params.iterations.toInt())
            .put("warmup", params.warmup.toInt())
        val specBytes = spec.toString().toByteArray(Charsets.UTF_8)
        val out = MobenchBuf()
        val status = MobenchNativeLibrary.api.mobench_run_benchmark_json(
            specBytes,
            specBytes.size.toLong(),
            out
        )
        out.read()

        if (status != 0) {
            val error = MobenchNativeLibrary.lastError()
            MobenchNativeLibrary.api.mobench_free_buf(out)
            throw IllegalStateException(error)
        }

        return try {
            val ptr = out.ptr ?: throw IllegalStateException("native benchmark returned a null report buffer")
            val bytes = ptr.getByteArray(0, out.len.toInt())
            JSONObject(bytes.toString(Charsets.UTF_8))
        } finally {
            MobenchNativeLibrary.api.mobench_free_buf(out)
        }
    }
}

private interface MobenchNativeAbi : Library {
    fun mobench_run_benchmark_json(specPtr: ByteArray, specLen: Long, out: MobenchBuf): Int
    fun mobench_free_buf(buf: MobenchBuf)
    fun mobench_last_error_message(): Pointer?
}

@Suppress("unused")
private class MobenchBuf : Structure() {
    @JvmField var ptr: Pointer? = null
    @JvmField var len: Long = 0
    @JvmField var cap: Long = 0

    override fun getFieldOrder(): List<String> = listOf("ptr", "len", "cap")
}

private object MobenchNativeLibrary {
    val api: MobenchNativeAbi by lazy {
        Native.load("{{LIBRARY_NAME}}", MobenchNativeAbi::class.java)
    }

    fun lastError(): String {
        return api.mobench_last_error_message()?.getString(0) ?: "native benchmark failed"
    }
}

@Suppress("DEPRECATION")
private fun Intent.resultReceiverExtra(): ResultReceiver? {
    return if (Build.VERSION.SDK_INT >= 33) {
        getParcelableExtra(RESULT_RECEIVER_EXTRA, ResultReceiver::class.java)
    } else {
        getParcelableExtra(RESULT_RECEIVER_EXTRA) as? ResultReceiver
    }
}

private class ProcessMemorySampler(private val sampleIntervalMs: Long = 1000L) {
    @Volatile private var running = false
    @Volatile private var peakKb = 0L
    private var samplerThread: Thread? = null

    fun start() {
        if (running) {
            return
        }
        running = true
        recordCurrent()
        samplerThread = Thread {
            while (running) {
                recordCurrent()
                try {
                    Thread.sleep(sampleIntervalMs)
                } catch (e: InterruptedException) {
                    Thread.currentThread().interrupt()
                    break
                }
            }
            recordCurrent()
        }.apply {
            name = "mobench-process-memory-sampler"
            isDaemon = true
            start()
        }
    }

    fun stop(): Long? {
        running = false
        try {
            samplerThread?.join(sampleIntervalMs * 2)
        } catch (e: InterruptedException) {
            Thread.currentThread().interrupt()
        }
        recordCurrent()
        return peakKb.takeIf { it > 0L }
    }

    @Synchronized
    private fun recordCurrent() {
        currentProcessPssKb()?.let { observedKb ->
            if (observedKb > peakKb) {
                peakKb = observedKb
            }
        }
    }
}

private fun currentProcessPssKb(): Long? {
    readProcMemoryKb("/proc/self/smaps_rollup", "Pss:")?.let { return it }
    readProcMemoryKb("/proc/self/status", "VmHWM:")?.let { return it }
    readProcMemoryKb("/proc/self/status", "VmRSS:")?.let { return it }

    val memInfo = Debug.MemoryInfo()
    return try {
        Debug.getMemoryInfo(memInfo)
        memInfo.totalPss.toLong().takeIf { it > 0L }
    } catch (e: Exception) {
        android.util.Log.d("BenchRunner", "Unable to read process memory from Debug.getMemoryInfo", e)
        null
    }
}

private fun readProcMemoryKb(path: String, label: String): Long? {
    return try {
        java.io.File(path).bufferedReader().use { reader ->
            var line: String? = reader.readLine()
            while (line != null) {
                if (line.startsWith(label)) {
                    return@use parseMemoryKb(line)
                }
                line = reader.readLine()
            }
            null
        }
    } catch (e: Exception) {
        null
    }
}

private fun parseMemoryKb(line: String): Long? {
    val value = line
        .substringAfter(':', "")
        .trim()
        .split(' ')
        .firstOrNull { it.isNotBlank() }
    return value?.toLongOrNull()?.takeIf { it > 0L }
}

private fun median(values: List<Long>): Long {
    val sorted = values.sorted()
    if (sorted.isEmpty()) {
        return 0L
    }
    val middle = sorted.size / 2
    return if (sorted.size % 2 == 0) {
        (sorted[middle - 1] + sorted[middle]) / 2
    } else {
        sorted[middle]
    }
}

private fun formatDuration(ns: Long): String {
    val ms = ns.toDouble() / 1_000_000.0
    return if (ms >= 1000.0) {
        String.format("%.3fs", ms / 1000.0)
    } else {
        String.format("%.3fms", ms)
    }
}

private fun buildNativeBenchReportJson(report: JSONObject, runProcessPeakMemoryKb: Long?): JSONObject {
    val json = JSONObject(report.toString())
    val spec = json.optJSONObject("spec") ?: JSONObject()
    val function = spec.optString("name", DEFAULT_FUNCTION)
    json.put("function", function)

    val samples = json.optJSONArray("samples") ?: JSONArray()
    val durations = mutableListOf<Long>()
    val cpuSamplesMs = mutableListOf<Long>()
    val peakSamplesKb = mutableListOf<Long>()
    val processPeakSamplesKb = mutableListOf<Long>()
    val samplesNs = JSONArray()
    for (index in 0 until samples.length()) {
        val sample = samples.optJSONObject(index) ?: continue
        val duration = sample.optLong("duration_ns", -1L)
        if (duration >= 0L) {
            durations.add(duration)
            samplesNs.put(duration)
        }
        optionalLong(sample, "cpu_time_ms")?.let { cpuSamplesMs.add(it) }
        optionalLong(sample, "peak_memory_kb")?.let { peakSamplesKb.add(it) }
        optionalLong(sample, "process_peak_memory_kb")?.let { processPeakSamplesKb.add(it) }
    }
    json.put("samples_ns", samplesNs)

    if (durations.isNotEmpty()) {
        val stats = JSONObject()
        stats.put("min_ns", durations.minOrNull() ?: 0L)
        stats.put("max_ns", durations.maxOrNull() ?: 0L)
        stats.put("avg_ns", durations.sum().toDouble() / durations.size.toDouble())
        stats.put("mean_ns", (durations.sum().toDouble() / durations.size.toDouble()).toLong())
        stats.put("median_ns", median(durations))
        json.put("stats", stats)
    }

    val resources = JSONObject()
    resources.put("platform", "android")
    resources.put("timestamp_ms", System.currentTimeMillis())
    resources.put("memory_process", "isolated_worker")
    if (cpuSamplesMs.isNotEmpty()) {
        val total = cpuSamplesMs.sum()
        resources.put("cpu_total_ms", total)
        resources.put("cpu_median_ms", median(cpuSamplesMs))
        resources.put("elapsed_cpu_ms", total)
    }
    peakSamplesKb.maxOrNull()?.let {
        resources.put("peak_memory_kb", it)
        resources.put("peak_memory_growth_kb", it)
    }
    (processPeakSamplesKb.maxOrNull() ?: runProcessPeakMemoryKb)?.let {
        resources.put("process_peak_memory_kb", it)
    }
    json.put("resources", resources)

    return json
}

private fun optionalLong(json: JSONObject, key: String): Long? {
    if (!json.has(key) || json.isNull(key)) {
        return null
    }
    return json.optLong(key).takeIf { it >= 0L }
}

private fun formatNativeBenchReport(report: JSONObject): String = buildString {
    val spec = report.optJSONObject("spec") ?: JSONObject()
    val function = spec.optString("name", DEFAULT_FUNCTION)
    val iterations = spec.optInt("iterations", DEFAULT_ITERATIONS.toInt())
    val warmup = spec.optInt("warmup", DEFAULT_WARMUP.toInt())
    val samples = report.optJSONArray("samples") ?: JSONArray()

    appendLine("=== Benchmark Results ===")
    appendLine()
    appendLine("Function: $function")
    appendLine("Iterations: $iterations")
    appendLine("Warmup: $warmup")
    appendLine()
    appendLine("Samples (${samples.length()}):")
    for (index in 0 until samples.length()) {
        val duration = samples.optJSONObject(index)?.optLong("duration_ns", 0L) ?: 0L
        appendLine("  ${index + 1}. ${formatDuration(duration)}")
    }

    val stats = report.optJSONObject("stats")
    if (stats != null) {
        appendLine()
        appendLine("Statistics:")
        appendLine("  Min: ${formatDuration(stats.optLong("min_ns", 0L))}")
        appendLine("  Max: ${formatDuration(stats.optLong("max_ns", 0L))}")
        appendLine("  Avg: ${formatDuration(stats.optLong("mean_ns", 0L))}")
    }
}
