param(
    [Parameter(Mandatory = $true)]
    [string]$TraceId,
    [string]$LokiUrl = "http://localhost:3100",
    [int]$LookbackMinutes = 30,
    [int]$Limit = 200,
    [string]$LogPath,
    [switch]$PreferLogFile,
    [switch]$Raw,
    [switch]$AsJson
)

$ErrorActionPreference = "Stop"

function Convert-TraceLineToEntry {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Line,
        [string]$Timestamp,
        [string]$DefaultLevel,
        [string]$FallbackTraceId
    )

    $parsed = $null
    try {
        $parsed = $Line | ConvertFrom-Json
    } catch {
        $parsed = $null
    }

    if ($null -eq $parsed) {
        return [pscustomobject]@{
            timestamp = $Timestamp
            level = $DefaultLevel
            method = $null
            route = $null
            status = $null
            duration_ms = $null
            outcome = $null
            tenant = $null
            request_id = $null
            trace_id = $FallbackTraceId
            message = $Line
            raw = $Line
        }
    }

    return [pscustomobject]@{
        timestamp = if ($parsed.timestamp) { $parsed.timestamp } else { $Timestamp }
        level = if ($parsed.level) { $parsed.level } else { $DefaultLevel }
        method = $parsed.method
        route = $parsed.route
        status = $parsed.status
        duration_ms = $parsed.duration_ms
        outcome = $parsed.outcome
        tenant = $parsed.tenant
        request_id = $parsed.request_id
        trace_id = if ($parsed.trace_id) { $parsed.trace_id } else { $FallbackTraceId }
        message = $parsed.message
        raw = $Line
    }
}

function Find-TraceInLogFile {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path,
        [Parameter(Mandatory = $true)]
        [string]$TraceId
    )

    if (-not (Test-Path -LiteralPath $Path)) {
        throw "Trace log file '$Path' does not exist."
    }

    $matches = foreach ($line in Get-Content -LiteralPath $Path) {
        if ($line -like "*$TraceId*") {
            Convert-TraceLineToEntry -Line $line -FallbackTraceId $TraceId
        }
    }

    return @($matches)
}

function Find-TraceInLoki {
    param(
        [Parameter(Mandatory = $true)]
        [string]$TraceId,
        [Parameter(Mandatory = $true)]
        [string]$LokiUrl,
        [Parameter(Mandatory = $true)]
        [int]$LookbackMinutes,
        [Parameter(Mandatory = $true)]
        [int]$Limit
    )

    $startNs = [DateTimeOffset]::UtcNow.AddMinutes(-$LookbackMinutes).ToUnixTimeMilliseconds() * 1000000
    $endNs = [DateTimeOffset]::UtcNow.ToUnixTimeMilliseconds() * 1000000
    $query = "{service=`"athena-auth`"} |= `"$TraceId`""
    $encodedQuery = [Uri]::EscapeDataString($query)
    $uri = "$LokiUrl/loki/api/v1/query_range?query=$encodedQuery&limit=$Limit&direction=BACKWARD&start=$startNs&end=$endNs"

    try {
        $response = Invoke-RestMethod -Method Get -Uri $uri
    } catch {
        throw "Failed to query Loki at $LokiUrl. Ensure the local observability stack is running. $($_.Exception.Message)"
    }

    $entries = foreach ($stream in $response.data.result) {
        foreach ($value in $stream.values) {
            $timestampNs = [Int64]$value[0]
            $timestamp = [DateTimeOffset]::FromUnixTimeMilliseconds([int64]($timestampNs / 1000000))
            $line = [string]$value[1]
            Convert-TraceLineToEntry `
                -Line $line `
                -Timestamp $timestamp.UtcDateTime.ToString("o") `
                -DefaultLevel $stream.stream.level `
                -FallbackTraceId $TraceId
        }
    }

    return @($entries)
}

$entries = @()
$source = $null
$lokiError = $null

if (-not $PreferLogFile) {
    try {
        $entries = Find-TraceInLoki -TraceId $TraceId -LokiUrl $LokiUrl -LookbackMinutes $LookbackMinutes -Limit $Limit
        if ($entries) {
            $source = "loki"
        }
    } catch {
        $lokiError = $_.Exception.Message
    }
}

if (-not $entries -and $LogPath) {
    $entries = Find-TraceInLogFile -Path $LogPath -TraceId $TraceId
    if ($entries) {
        if ($lokiError) {
            Write-Warning "$lokiError Falling back to log file '$LogPath'."
        }
        $source = "log-file"
    }
}

if (-not $entries -and $PreferLogFile -and -not $LogPath) {
    throw "PreferLogFile was requested but no -LogPath was provided."
}

if (-not $entries -and $lokiError -and -not $LogPath) {
    throw $lokiError
}

if (-not $entries) {
    if ($LogPath) {
        Write-Warning "No matching trace lines were found for '$TraceId' in Loki or in '$LogPath'."
    } else {
        Write-Warning "No Loki log lines matched trace ID '$TraceId' in the last $LookbackMinutes minutes."
    }
    exit 0
}

$entries = $entries | Sort-Object timestamp
Write-Host "Source: $source"

if ($AsJson) {
    $entries | ConvertTo-Json -Depth 20
    exit 0
}

if ($Raw) {
    $entries | Select-Object timestamp, raw | Format-Table -Wrap -AutoSize
    exit 0
}

$displayEntries = $entries |
    Select-Object timestamp, level, message, method, route, status, duration_ms, outcome, tenant, request_id, trace_id

if (@($displayEntries).Count -le 5) {
    $displayEntries | Format-List
    exit 0
}

$displayEntries | Format-Table -Wrap -AutoSize
