/*
 * This file is part of LibEuFin.
 * Copyright (C) 2024, 2025, 2026 Taler Systems S.A.

 * LibEuFin is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation; either version 3, or
 * (at your option) any later version.

 * LibEuFin is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General
 * Public License for more details.

 * You should have received a copy of the GNU Affero General Public
 * License along with LibEuFin; see the file COPYING.  If not, see
 * <http://www.gnu.org/licenses/>
 */

package tech.libeufin.common.test

import org.postgresql.copy.*
import org.postgresql.jdbc.*
import tech.libeufin.common.*
import kotlin.math.pow
import kotlin.math.sqrt
import kotlin.time.DurationUnit
import kotlin.time.measureTime
import kotlin.time.toDuration

fun PgConnection.genData(amount: Int, generators: Sequence<Pair<String, (Int) -> String>>) {
    for ((table, generator) in generators) {
        println("Gen rows for $table")
        PGCopyOutputStream(this, "COPY $table FROM STDIN", 16 * 1024 * 1024).use { out -> 
            repeat(amount) { 
                val str = generator(it+1)
                val bytes = str.toByteArray()
                out.write(bytes)
            }
        }
    }

    // Update database statistics for better perf
    this.execSQLUpdate("VACUUM FULL ANALYZE")
}

class Benchmark(private val iter: Int) {
    private val WARN = 4.toDuration(DurationUnit.MILLISECONDS)
    private val ERR = 50.toDuration(DurationUnit.MILLISECONDS)
    internal val measures: MutableList<List<String>> = mutableListOf()

    internal fun fmtMeasures(times: LongArray): List<String> {
        val min: Long = times.min()
        val max: Long = times.max()
        val mean: Long = times.average().toLong()
        val variance = times.map { (it.toDouble() - mean).pow(2) }.average()
        val stdVar: Long = sqrt(variance.toDouble()).toLong()
        return sequenceOf(min, mean, max, stdVar).map {
            val duration = it.toDuration(DurationUnit.MICROSECONDS)
            val str = duration.toString()
            if (duration > ERR) {
                ANSI.red(str)
            } else if (duration > WARN) {
                ANSI.yellow(str)
            } else {
                ANSI.green(str)
            }
            
        }.toList()
    }

    suspend fun <R> measureAction(name: String, lambda: suspend (Int) -> R): List<R> {
        println("Measure action $name")
        val results = mutableListOf<R>()
        val times = LongArray(iter) { idx ->
            measureTime { 
                val result = lambda(idx)
                results.add(result)
            }.inWholeMicroseconds
        }
        measures.add(listOf(ANSI.magenta(name)) + fmtMeasures(times))
        return results
    }
}

fun bench(lambda: Benchmark.(Int) -> Unit) {
    val ITER = System.getenv("BENCH_ITER")?.toIntOrNull() ?: 10
    val AMOUNT = System.getenv("BENCH_AMOUNT")?.toIntOrNull() ?: 100

    println("Bench $ITER times with $AMOUNT rows")

    val bench = Benchmark(ITER)
    bench.lambda(AMOUNT)
    printTable(
        listOf("benchmark", "min", "mean", "max", "std").map { ANSI.bold(it) },
        bench.measures,
        ' ',
        listOf(ColumnStyle.DEFAULT) + List(5) { ColumnStyle(false) }
    )
}