[CmdletBinding()]
param(
    [ValidateSet('Release', 'ReleaseFast', 'ReleaseSize')]
    [string]$BuildProfile = 'ReleaseFast',
    [ValidateSet('Debug', 'Release', 'ReleaseFast', 'ReleaseSize')]
    [string]$FlamegraphBuildProfile = 'Release',
    [int]$Warmup = 2,
    [int]$Runs = 6,
    [switch]$SkipBenchmarks,
    [switch]$IncludeBloat,
    [switch]$IncludeLlvmLines,
    [switch]$IncludeBuildTimings,
    [switch]$IncludeFlamegraph,
    [string[]]$FlamegraphPackages = @('codeshape', 'refs')
)

$ErrorActionPreference = 'Stop'
Set-StrictMode -Version Latest

. (Join-Path $PSScriptRoot 'toolbox-commands.ps1')

function Invoke-StrictNative {
    param(
        [Parameter(Mandatory = $true)]
        [string]$FilePath,
        [Parameter(Mandatory = $true)]
        [string[]]$ArgumentList
    )

    & $FilePath @ArgumentList
    if ($LASTEXITCODE -ne 0) {
        throw "Command failed with exit code ${LASTEXITCODE}: $FilePath $($ArgumentList -join ' ')"
    }
}

function Get-RequiredCommandPath {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Name
    )

    $command = Get-Command -Name $Name -ErrorAction SilentlyContinue
    if ($null -eq $command) {
        throw "Required command not found on PATH: $Name"
    }

    return $command.Source
}

function Get-OptionalCommandPath {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Name
    )

    $command = Get-Command -Name $Name -ErrorAction SilentlyContinue
    if ($null -eq $command) {
        return $null
    }

    return $command.Source
}

function Resolve-XperfPath {
    $commandPath = Get-OptionalCommandPath -Name 'xperf'
    if (-not [string]::IsNullOrWhiteSpace($commandPath)) {
        return $commandPath
    }

    $candidates = @()
    foreach ($root in @($env:ProgramFiles, ${env:ProgramFiles(x86)})) {
        if (-not [string]::IsNullOrWhiteSpace($root)) {
            $candidates += (Join-Path $root 'Windows Kits\10\Windows Performance Toolkit\xperf.exe')
        }
    }

    foreach ($candidate in $candidates) {
        if (Test-Path -LiteralPath $candidate) {
            return $candidate
        }
    }

    return $null
}

function Test-IsElevated {
    if (-not $IsWindows) {
        return $false
    }

    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Get-FlamegraphCapability {
    $capability = [ordered]@{
        platform = if ($IsWindows) { 'windows' } else { 'other' }
        can_attempt = $true
        cargo_flamegraph_backend = if ($IsWindows) { 'unknown' } else { 'platform-default' }
        is_elevated = $false
        elevation_wrapper = '.\scripts\cargo-flamegraph-windows.ps1'
        dtrace_path = $null
        dtrace_source = $null
        xperf_path = $null
        samply_path = $null
        samply_available = $false
        reason = 'Non-Windows platform; defer to cargo flamegraph defaults.'
        note = $null
    }

    if (-not $IsWindows) {
        return [pscustomobject]$capability
    }

    $capability.is_elevated = Test-IsElevated
    $capability.samply_path = Get-OptionalCommandPath -Name 'samply'
    $capability.xperf_path = Resolve-XperfPath
    $capability.samply_available = (
        -not [string]::IsNullOrWhiteSpace([string]$capability.samply_path) -and
        -not [string]::IsNullOrWhiteSpace([string]$capability.xperf_path)
    )

    $dtraceOverride = $env:DTRACE
    if (-not [string]::IsNullOrWhiteSpace($dtraceOverride)) {
        $resolvedOverride = $null
        if (Test-Path -LiteralPath $dtraceOverride) {
            $resolvedOverride = (Resolve-Path -LiteralPath $dtraceOverride).Path
        }
        else {
            $resolvedCommand = Get-Command -Name $dtraceOverride -ErrorAction SilentlyContinue
            if ($null -ne $resolvedCommand) {
                $resolvedOverride = $resolvedCommand.Source
            }
        }

        if (-not [string]::IsNullOrWhiteSpace($resolvedOverride)) {
            $capability.dtrace_path = $resolvedOverride
            $capability.dtrace_source = 'env:DTRACE'
        }
        else {
            $capability.note = "DTRACE is set but does not currently resolve: $dtraceOverride"
        }
    }

    if ([string]::IsNullOrWhiteSpace([string]$capability.dtrace_path)) {
        $dtracePath = Get-OptionalCommandPath -Name 'dtrace'
        if (-not [string]::IsNullOrWhiteSpace($dtracePath)) {
            $capability.dtrace_path = $dtracePath
            $capability.dtrace_source = 'PATH'
        }
    }

    if (-not [string]::IsNullOrWhiteSpace([string]$capability.dtrace_path)) {
        $capability.cargo_flamegraph_backend = 'dtrace'
        $capability.reason = 'DTrace detected; cargo flamegraph will prefer DTrace on Windows.'
        if (-not $capability.is_elevated -and [string]::IsNullOrWhiteSpace([string]$capability.note)) {
            $capability.note = 'Microsoft documents DTrace setup and validation from an elevated prompt; if collection fails, rerun from an elevated shell.'
        }

        return [pscustomobject]$capability
    }

    $capability.cargo_flamegraph_backend = 'blondie'
    if ($capability.is_elevated) {
        $capability.reason = 'No DTrace detected; cargo flamegraph will fall back to blondie, and the shell is elevated.'
    }
    else {
        $capability.can_attempt = $false
        $capability.reason = 'No DTrace detected; cargo flamegraph will fall back to blondie, which requires an elevated shell on Windows because it uses the ETW Kernel Logger session.'
        if ($capability.samply_available) {
            $capability.note = 'Use .\scripts\cargo-flamegraph-windows.ps1 for an auto-elevated one-shot flamegraph. samply is also available separately because xperf is installed, but cargo flamegraph itself still needs elevation or DTrace.'
        }
        elseif ([string]::IsNullOrWhiteSpace([string]$capability.note)) {
            $capability.note = 'Use .\scripts\cargo-flamegraph-windows.ps1 for an auto-elevated one-shot flamegraph. samply is not a usable fallback on this machine because xperf from Windows Performance Toolkit is not installed.'
        }
    }

    return [pscustomobject]$capability
}

function Get-BinaryTable {
    param(
        [Parameter(Mandatory = $true)]
        [string]$WorkspaceRoot,
        [Parameter(Mandatory = $true)]
        [string]$TargetDir
    )

    return @(
        Get-ChildItem (Join-Path $WorkspaceRoot $TargetDir) -Filter '*.exe' |
            Sort-Object Length -Descending |
            Select-Object Name,
                @{ Name = 'SizeKB'; Expression = { [math]::Round($_.Length / 1KB, 1) } },
                @{ Name = 'SizeBytes'; Expression = { $_.Length } }
    )
}

function New-BenchmarkCommands {
    param(
        [Parameter(Mandatory = $true)]
        [string]$WorkspaceRoot,
        [Parameter(Mandatory = $true)]
        [string]$TargetDir
    )

    $binaryRoot = Join-Path $WorkspaceRoot $TargetDir
    $fixtureRoot = Join-Path $WorkspaceRoot 'fixtures'

    return @(
        ('"{0}" --sort-keys "{1}"' -f (Join-Path $binaryRoot 'cjson.exe'), (Join-Path $fixtureRoot 'cjson\sample.json')),
        ('"{0}" level=error "{1}" --pick ts,msg' -f (Join-Path $binaryRoot 'jsonlgrep.exe'), (Join-Path $fixtureRoot 'jsonl\events.jsonl')),
        ('"{0}" "{1}" --json' -f (Join-Path $binaryRoot 'jsonshape.exe'), (Join-Path $fixtureRoot 'toon\config.json')),
        ('"{0}" --root "{1}" --since 9999h --ext rs --limit 20' -f (Join-Path $binaryRoot 'recent.exe'), $fixtureRoot),
        ('"{0}" "{1}" --max-depth 2 --limit-per-file 12' -f (Join-Path $binaryRoot 'codeshape.exe'), (Join-Path $fixtureRoot 'polyglot\repo')),
        ('"{0}" helper "{1}" --limit 6' -f (Join-Path $binaryRoot 'defsnip.exe'), (Join-Path $fixtureRoot 'polyglot\repo')),
        ('"{0}" helper "{1}" --limit 12' -f (Join-Path $binaryRoot 'refs.exe'), (Join-Path $fixtureRoot 'polyglot\repo')),
        ('"{0}" --context 1 "{1}:18" "{1}:26"' -f (Join-Path $binaryRoot 'hitsnip.exe'), (Join-Path $fixtureRoot 'reading\sample.rs')),
        ('"{0}" "{1}"' -f (Join-Path $binaryRoot 'stringscan.exe'), (Join-Path $fixtureRoot 'binaries\stringscan-sample.bin')),
        ('"{0}" "{1}" --json' -f (Join-Path $binaryRoot 'toon.exe'), (Join-Path $fixtureRoot 'toon\config.toon')),
        ('"{0}" rg' -f (Join-Path $binaryRoot 'pathshadow.exe')),
        ('"{0}" --env none --group shell' -f (Join-Path $binaryRoot 'sysshape.exe'))
    )
}

function Get-FlamegraphTarget {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Package,
        [Parameter(Mandatory = $true)]
        [string]$WorkspaceRoot
    )

    $fixtureRoot = Join-Path $WorkspaceRoot 'fixtures'

    switch ($Package) {
        'codeshape' {
            return @{
                Package = 'codeshape'
                Binary = 'codeshape'
                Arguments = @((Join-Path $fixtureRoot 'polyglot\repo'), '--max-depth', '2', '--limit-per-file', '12')
            }
        }
        'refs' {
            return @{
                Package = 'refs'
                Binary = 'refs'
                Arguments = @('helper', (Join-Path $fixtureRoot 'polyglot\repo'), '--limit', '12')
            }
        }
        'jsonshape' {
            return @{
                Package = 'jsonshape'
                Binary = 'jsonshape'
                Arguments = @((Join-Path $fixtureRoot 'toon\config.json'), '--json')
            }
        }
        default {
            return @{
                Package = $Package
                Binary = $Package
                Arguments = @('--help')
            }
        }
    }
}

function Assert-FlamegraphPackageName {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Package
    )

    if ([string]::IsNullOrWhiteSpace($Package)) {
        throw 'Flamegraph package name must not be empty.'
    }

    if ([System.IO.Path]::IsPathRooted($Package) -or $Package.Contains([System.IO.Path]::DirectorySeparatorChar) -or $Package.Contains([System.IO.Path]::AltDirectorySeparatorChar)) {
        throw "Flamegraph package name must be a single package token: $Package"
    }

    if ($Package.IndexOfAny([System.IO.Path]::GetInvalidFileNameChars()) -ge 0) {
        throw "Flamegraph package name contains invalid filename characters: $Package"
    }
}

function Write-JsonFile {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path,
        [Parameter(Mandatory = $true)]
        [object]$Value
    )

    $directory = Split-Path -Parent $Path
    if (-not [string]::IsNullOrWhiteSpace($directory)) {
        New-Item -ItemType Directory -Force -Path $directory | Out-Null
    }

    $Value | ConvertTo-Json -Depth 6 | Set-Content -LiteralPath $Path -Encoding utf8NoBOM
}

$workspaceRoot = Split-Path -Parent $PSScriptRoot
$cargoPath = Get-RequiredCommandPath -Name 'cargo'
$batPath = Get-OptionalCommandPath -Name 'bat'
$targetProfile = Resolve-ToolboxProfileName -Configuration $BuildProfile
$flamegraphTargetProfile = Resolve-ToolboxProfileName -Configuration $FlamegraphBuildProfile
$targetDir = Join-Path 'target' $targetProfile
$outputRoot = Join-Path (Join-Path $workspaceRoot 'dist') 'probe'
$profileOutputRoot = Join-Path $outputRoot $targetProfile
$benchmarkJsonPath = Join-Path $profileOutputRoot 'hyperfine.json'
$sizeJsonPath = Join-Path $profileOutputRoot 'binary-sizes.json'
$bloatRoot = Join-Path $profileOutputRoot 'bloat'
$llvmLinesRoot = Join-Path $profileOutputRoot 'llvm-lines'
$flamegraphRoot = Join-Path $profileOutputRoot 'flamegraphs'

Write-Host "Mercury Toolbox probe"
Write-Host "workspace: $workspaceRoot"
Write-Host "profile: $BuildProfile"
Write-Host "target dir: $targetDir"
Write-Host "output root: $profileOutputRoot"

New-Item -ItemType Directory -Force -Path $profileOutputRoot | Out-Null
if ($IncludeBuildTimings) {
    Invoke-StrictToolboxNative -FilePath $cargoPath -ArgumentList ((Get-ToolboxBuildArguments -Configuration $BuildProfile) + '--timings')
    if ($BuildProfile -ne 'Debug') {
        Invoke-StrictToolboxNative -FilePath $cargoPath -ArgumentList (
            Get-ToolboxBuildArguments -Configuration $BuildProfile -Packages (Get-ToolboxSlimBuildCommandNames)
        )
    }
}
else {
    Invoke-ToolboxBuild -CargoPath $cargoPath -Configuration $BuildProfile
}

$binaryTable = Get-BinaryTable -WorkspaceRoot $workspaceRoot -TargetDir $targetDir

Write-Host ''
Write-Host 'Binary sizes:'
$binaryTable | Format-Table -AutoSize
Write-JsonFile -Path $sizeJsonPath -Value $binaryTable
Write-Host "Size report: $sizeJsonPath"

if ($IncludeBuildTimings) {
    $timingsDir = Join-Path (Join-Path $workspaceRoot 'target') 'cargo-timings'
    $latestTiming = Get-ChildItem -LiteralPath $timingsDir -Filter 'cargo-timing*.html' -File -ErrorAction SilentlyContinue |
        Sort-Object LastWriteTimeUtc -Descending |
        Select-Object -First 1
    if ($null -ne $latestTiming) {
        Write-Host "Build timings: $($latestTiming.FullName)"
    }
}

if (-not $SkipBenchmarks) {
    $hyperfinePath = Get-RequiredCommandPath -Name 'hyperfine'
    $benchmarkCommands = New-BenchmarkCommands -WorkspaceRoot $workspaceRoot -TargetDir $targetDir
    $hyperfineArguments = @(
        '--warmup',
        $Warmup.ToString(),
        '--runs',
        $Runs.ToString(),
        '--shell',
        'none',
        '--export-json',
        $benchmarkJsonPath
    ) + $benchmarkCommands
    Write-Host ''
    Write-Host 'Representative latency benchmarks:'
    Invoke-StrictNative -FilePath $hyperfinePath -ArgumentList $hyperfineArguments
    Write-Host "Benchmark report: $benchmarkJsonPath"
}

if ($IncludeBloat) {
    $largest = $binaryTable | Select-Object -First 3
    New-Item -ItemType Directory -Force -Path $bloatRoot | Out-Null

    foreach ($binary in $largest) {
        $packageName = [System.IO.Path]::GetFileNameWithoutExtension([string]$binary.Name)
        $reportPath = Join-Path $bloatRoot "$packageName.txt"
        Write-Host ''
        Write-Host "cargo-bloat top crates for ${packageName}:"
        & $cargoPath @(
            'bloat',
            '--profile',
            $targetProfile,
            '-p',
            $packageName,
            '--crates',
            '-n',
            '20'
        ) *> $reportPath
        if ($LASTEXITCODE -ne 0) {
            throw "Command failed with exit code ${LASTEXITCODE}: cargo bloat --profile $targetProfile -p $packageName --crates -n 20"
        }
        if ([string]::IsNullOrWhiteSpace($batPath)) {
            Get-Content -LiteralPath $reportPath -TotalCount 24
        }
        else {
            & $batPath '--style=plain' '--paging=never' $reportPath | Select-Object -First 24
        }
        Write-Host "Bloat report: $reportPath"
    }
}

if ($IncludeLlvmLines) {
    $largest = $binaryTable | Select-Object -First 3
    New-Item -ItemType Directory -Force -Path $llvmLinesRoot | Out-Null

    foreach ($binary in $largest) {
        $packageName = [System.IO.Path]::GetFileNameWithoutExtension([string]$binary.Name)
        $reportPath = Join-Path $llvmLinesRoot "$packageName.txt"
        Write-Host ''
        Write-Host "cargo-llvm-lines top items for ${packageName}:"
        & $cargoPath @(
            'llvm-lines',
            '--profile',
            $targetProfile,
            '-p',
            $packageName,
            '--bin',
            $packageName
        ) *> $reportPath
        if ($LASTEXITCODE -ne 0) {
            throw "Command failed with exit code ${LASTEXITCODE}: cargo llvm-lines --profile $targetProfile -p $packageName --bin $packageName"
        }
        if ([string]::IsNullOrWhiteSpace($batPath)) {
            Get-Content -LiteralPath $reportPath -TotalCount 80
        }
        else {
            & $batPath '--style=plain' '--paging=never' $reportPath | Select-Object -First 80
        }
        Write-Host "LLVM lines report: $reportPath"
    }
}

if ($IncludeFlamegraph) {
    [void](Get-RequiredCommandPath -Name 'cargo-flamegraph')
    New-Item -ItemType Directory -Force -Path $flamegraphRoot | Out-Null
    $flamegraphCapability = Get-FlamegraphCapability
    $capabilityPath = Join-Path $flamegraphRoot 'capability.json'
    Write-JsonFile -Path $capabilityPath -Value $flamegraphCapability
    Write-Host "Flamegraph capability: $capabilityPath"
    if (-not [bool]$flamegraphCapability.can_attempt) {
        Write-Warning "Skipping flamegraph collection on this machine: $($flamegraphCapability.reason)"
    }
    elseif (-not [string]::IsNullOrWhiteSpace([string]$flamegraphCapability.note)) {
        Write-Warning $flamegraphCapability.note
    }

    foreach ($package in $FlamegraphPackages) {
        Assert-FlamegraphPackageName -Package $package
        $target = Get-FlamegraphTarget -Package $package -WorkspaceRoot $workspaceRoot
        $outputPath = Join-Path $flamegraphRoot "$($target.Package).svg"
        $logPath = Join-Path $flamegraphRoot "$($target.Package).log"
        $errorPath = Join-Path $flamegraphRoot "$($target.Package)-error.txt"
        $arguments = @(
            'flamegraph',
            '--profile',
            $flamegraphTargetProfile,
            '-p',
            $target.Package,
            '--bin',
            $target.Binary,
            '-o',
            $outputPath,
            '--'
        ) + $target.Arguments

        Write-Host ''
        Write-Host "cargo-flamegraph for $($target.Package):"
        if (-not [bool]$flamegraphCapability.can_attempt) {
            Write-JsonFile -Path $errorPath -Value ([ordered]@{
                    package = $target.Package
                    skipped = $true
                    reason = $flamegraphCapability.reason
                    profiler = $flamegraphCapability
                })
            Write-Host "Flamegraph skipped: $errorPath"
            continue
        }

        try {
            & $cargoPath @arguments *> $logPath
            if ($LASTEXITCODE -ne 0) {
                throw "Command failed with exit code ${LASTEXITCODE}: cargo $($arguments -join ' ')"
            }
            Write-Host "Flamegraph: $outputPath"
        }
        catch {
            Write-Warning "Skipping flamegraph for $($target.Package): $($_.Exception.Message)"
            Write-JsonFile -Path $errorPath -Value ([ordered]@{
                package = $target.Package
                command = @('cargo') + $arguments
                error = $_.Exception.Message
                log_path = $logPath
                profiler = $flamegraphCapability
            })
            Write-Host "Flamegraph error log: $errorPath"
        }
    }
}

$finalBinaryTable = Get-BinaryTable -WorkspaceRoot $workspaceRoot -TargetDir $targetDir
if (($finalBinaryTable | ConvertTo-Json -Depth 4) -cne ($binaryTable | ConvertTo-Json -Depth 4)) {
    Write-Host ''
    Write-Host 'Binary sizes changed during probe; refreshing final size report:'
    $finalBinaryTable | Format-Table -AutoSize
    Write-JsonFile -Path $sizeJsonPath -Value $finalBinaryTable
    Write-Host "Refreshed size report: $sizeJsonPath"
}