param(
    [string]$PgHost = "localhost",
    [int]$PgPort = 5432,
    [string]$PgUser = "postgres",
    [string]$BenchDb = "pgbench",
    [int]$Scale = 50,
    [int]$Duration = 60,
    [string]$Clients = "1 4 8 16 32",
    [int]$Jobs = 4,
    [int]$WarmupSeconds = 15,
    [bool]$RunReadOnly = $true,
    [bool]$RunCustom = $false,
    [string]$CustomSqlFile = "pgbench_custom_workload.sql",
    [string]$CustomModeName = "custom_sql",
    [ValidateRange(1, 99)][int]$CustomReadPct = 80,
    [ValidateRange(1, 100)][int]$CustomHotPct = 90,
    [ValidateRange(1, 4)][int]$CustomTxnSize = 2,
    [ValidateRange(100, 100000000)][int]$CustomHotAccounts = 10000,
    [bool]$Reinit = $false
)

$ErrorActionPreference = "Stop"

function Require-Command {
    param([Parameter(Mandatory = $true)][string]$Name)

    if (-not (Get-Command $Name -ErrorAction SilentlyContinue)) {
        throw "Missing required command: $Name"
    }
}

function Invoke-PsqlScalar {
    param(
        [Parameter(Mandatory = $true)][string]$Database,
        [Parameter(Mandatory = $true)][string]$Sql
    )

    $output = & psql -X -v ON_ERROR_STOP=1 -h $PgHost -p $PgPort -U $PgUser -d $Database -Atc $Sql
    if ($LASTEXITCODE -ne 0) {
        throw "psql failed running SQL: $Sql"
    }

    return ($output | Select-Object -First 1)
}

function Invoke-PsqlNonQuery {
    param(
        [Parameter(Mandatory = $true)][string]$Database,
        [Parameter(Mandatory = $true)][string]$Sql
    )

    & psql -X -v ON_ERROR_STOP=1 -h $PgHost -p $PgPort -U $PgUser -d $Database -c $Sql | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "psql failed running SQL: $Sql"
    }
}

function Ensure-BenchmarkDatabase {
    $exists = Invoke-PsqlScalar -Database "postgres" -Sql "SELECT 1 FROM pg_database WHERE datname = '$BenchDb';"
    if ($exists -ne "1") {
        Write-Host "Creating benchmark database: $BenchDb"
        Invoke-PsqlNonQuery -Database "postgres" -Sql "CREATE DATABASE $BenchDb;"
    }
}

function Needs-Init {
    $hasTable = Invoke-PsqlScalar -Database $BenchDb -Sql "SELECT to_regclass('public.pgbench_accounts') IS NOT NULL;"
    return $hasTable -ne "t"
}

function Init-OrReinitData {
    if ($Reinit -or (Needs-Init)) {
        Write-Host "Initializing pgbench schema/data (scale=$Scale)"
        & pgbench -h $PgHost -p $PgPort -U $PgUser -i -s $Scale $BenchDb | Out-Host
        if ($LASTEXITCODE -ne 0) {
            throw "pgbench initialization failed"
        }
    }
    else {
        Write-Host "Using existing pgbench data in $BenchDb"
    }
}

function Parse-Tps {
    param([Parameter(Mandatory = $true)][string]$Text)

    $matches = [regex]::Matches($Text, "(?m)^tps\s*=\s*([0-9]+(?:\.[0-9]+)?)")
    if ($matches.Count -gt 0) {
        return $matches[$matches.Count - 1].Groups[1].Value
    }

    return ""
}

function Parse-AvgLatencyMs {
    param([Parameter(Mandatory = $true)][string]$Text)

    $matches = [regex]::Matches($Text, "(?m)^latency average\s*=\s*([0-9]+(?:\.[0-9]+)?)\s*ms")
    if ($matches.Count -gt 0) {
        return $matches[$matches.Count - 1].Groups[1].Value
    }

    return ""
}

function Run-Case {
    param(
        [Parameter(Mandatory = $true)][string]$Mode,
        [Parameter(Mandatory = $true)][int]$ClientCount,
        [AllowEmptyCollection()][string[]]$ExtraArgs = @(),
        [Parameter(Mandatory = $true)][string]$ResultDir,
        [Parameter(Mandatory = $true)][string]$SummaryCsv
    )

    $logFile = Join-Path $ResultDir ("{0}_c{1}.txt" -f $Mode, $ClientCount)
    $txLogPrefix = Join-Path $ResultDir ("{0}_c{1}_txlog" -f $Mode, $ClientCount)
    Write-Host "Running mode=$Mode, clients=$ClientCount, duration=${Duration}s"

    # Warm cache and avoid measuring first-run effects.
    & pgbench -h $PgHost -p $PgPort -U $PgUser -n -M prepared -j $Jobs -c $ClientCount -T $WarmupSeconds @ExtraArgs $BenchDb | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "Warm-up run failed for mode=$Mode clients=$ClientCount"
    }

    $runOutput = & pgbench -h $PgHost -p $PgPort -U $PgUser -n -M prepared -j $Jobs -c $ClientCount -T $Duration -r -l --log-prefix $txLogPrefix @ExtraArgs $BenchDb 2>&1
    if ($LASTEXITCODE -ne 0) {
        throw "Measured run failed for mode=$Mode clients=$ClientCount"
    }

    $runOutput | Set-Content -Path $logFile -Encoding UTF8

    $text = ($runOutput -join [Environment]::NewLine)
    $tps = Parse-Tps -Text $text
    $latencyMs = Parse-AvgLatencyMs -Text $text
    $percentiles = Get-PercentileLatency -LogPrefix $txLogPrefix

    "$Mode,$ClientCount,$Duration,$Jobs,$tps,$latencyMs,$($percentiles.P95Ms),$($percentiles.P99Ms),$logFile" | Add-Content -Path $SummaryCsv
}

function Get-PercentileLatency {
    param([Parameter(Mandatory = $true)][string]$LogPrefix)

    $logFiles = Get-ChildItem -Path ($LogPrefix + "*") -File -ErrorAction SilentlyContinue
    if (-not $logFiles -or $logFiles.Count -eq 0) {
        return @{ P95Ms = ""; P99Ms = "" }
    }

    # Keep memory bounded even for very large transaction logs.
    $maxSampleSize = 200000
    $reservoir = New-Object System.Collections.Generic.List[double]
    $reservoir.Capacity = $maxSampleSize
    $random = [System.Random]::new()
    $seen = 0

    foreach ($file in $logFiles) {
        $reader = [System.IO.File]::OpenText($file.FullName)
        try {
            while (($line = $reader.ReadLine()) -ne $null) {
                if ([string]::IsNullOrWhiteSpace($line) -or $line.StartsWith("#")) {
                    continue
                }

                $tokens = $line -split "[\s,]+" | Where-Object { $_ -ne "" }
                if ($tokens.Count -ge 3 -and $tokens[2] -match "^[0-9]+(?:\.[0-9]+)?$") {
                    $latUs = [double]$tokens[2]
                    $seen++

                    if ($reservoir.Count -lt $maxSampleSize) {
                        $reservoir.Add($latUs)
                    }
                    else {
                        $idx = $random.Next(0, $seen)
                        if ($idx -lt $maxSampleSize) {
                            $reservoir[$idx] = $latUs
                        }
                    }
                }
            }
        }
        finally {
            $reader.Dispose()
        }
    }

    if ($reservoir.Count -eq 0) {
        return @{ P95Ms = ""; P99Ms = "" }
    }

    $sorted = $reservoir | Sort-Object
    $p95Us = Select-PercentileValue -SortedValues $sorted -Percentile 95
    $p99Us = Select-PercentileValue -SortedValues $sorted -Percentile 99

    return @{
        P95Ms = [math]::Round(($p95Us / 1000.0), 3)
        P99Ms = [math]::Round(($p99Us / 1000.0), 3)
    }
}

function Select-PercentileValue {
    param(
        [Parameter(Mandatory = $true)][double[]]$SortedValues,
        [Parameter(Mandatory = $true)][ValidateRange(1, 99)][int]$Percentile
    )

    $n = $SortedValues.Count
    $rank = [math]::Ceiling(($Percentile / 100.0) * $n)
    if ($rank -lt 1) { $rank = 1 }
    if ($rank -gt $n) { $rank = $n }
    return $SortedValues[$rank - 1]
}

function Resolve-CustomSqlPath {
    if ([System.IO.Path]::IsPathRooted($CustomSqlFile)) {
        return $CustomSqlFile
    }

    return (Join-Path (Get-Location) $CustomSqlFile)
}

Require-Command -Name "psql"
Require-Command -Name "pgbench"

$timestamp = Get-Date -Format "yyyyMMdd_HHmmss"
$resultDir = Join-Path (Get-Location) ("pgbench_results_{0}" -f $timestamp)
$summaryCsv = Join-Path $resultDir "summary.csv"

New-Item -Path $resultDir -ItemType Directory -Force | Out-Null
"mode,clients,duration_seconds,jobs,tps,latency_ms,p95_ms,p99_ms,details_file" | Set-Content -Path $summaryCsv -Encoding UTF8

Ensure-BenchmarkDatabase
Init-OrReinitData

$customSqlPath = ""
if ($RunCustom) {
    $customSqlPath = Resolve-CustomSqlPath
    if (-not (Test-Path -Path $customSqlPath -PathType Leaf)) {
        throw "Custom SQL file not found: $customSqlPath"
    }
}

$clientValues = $Clients -split "\s+" | Where-Object { $_ -match "^[0-9]+$" }
if ($clientValues.Count -eq 0) {
    throw "No valid client values found. Example: -Clients '1 4 8 16 32'"
}

foreach ($c in $clientValues) {
    $clientInt = [int]$c
    Run-Case -Mode "read_write" -ClientCount $clientInt -ExtraArgs @() -ResultDir $resultDir -SummaryCsv $summaryCsv

    if ($RunReadOnly) {
        Run-Case -Mode "read_only" -ClientCount $clientInt -ExtraArgs @("-S") -ResultDir $resultDir -SummaryCsv $summaryCsv
    }

    if ($RunCustom) {
        Run-Case -Mode $CustomModeName -ClientCount $clientInt -ExtraArgs @(
            "-D", "read_pct=$CustomReadPct",
            "-D", "hot_pct=$CustomHotPct",
            "-D", "txn_size=$CustomTxnSize",
            "-D", "hot_accounts=$CustomHotAccounts",
            "-f", $customSqlPath
        ) -ResultDir $resultDir -SummaryCsv $summaryCsv
    }
}

Write-Host ""
Write-Host "Benchmark completed."
Write-Host "Summary: $summaryCsv"
Write-Host "Raw outputs: $resultDir"