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

function Get-WorkshopRoot {
    return $PSScriptRoot
}

function Get-WorkshopWorkDirectory {
    $path = Join-Path (Get-WorkshopRoot) 'work'
    if (-not (Test-Path -LiteralPath $path)) {
        New-Item -ItemType Directory -Path $path -Force | Out-Null
    }
    return $path
}

function Write-WorkshopJson {
    param(
        [Parameter(Mandatory)] [string] $Path,
        [Parameter(Mandatory)] $Value
    )

    $parent = Split-Path -Parent $Path
    if ($parent -and -not (Test-Path -LiteralPath $parent)) {
        New-Item -ItemType Directory -Path $parent -Force | Out-Null
    }
    $Value | ConvertTo-Json -Depth 50 | Set-Content -LiteralPath $Path -Encoding utf8NoBOM
}

function Invoke-AzureCli {
    param(
        [Parameter(Mandatory)] [string[]] $Arguments,
        [switch] $AllowFailure
    )

    $output = & az @Arguments 2>&1
    $exitCode = $LASTEXITCODE
    if ($exitCode -ne 0 -and -not $AllowFailure) {
        throw "Azure CLI failed (exit $exitCode): az $($Arguments -join ' ')`n$($output -join [Environment]::NewLine)"
    }
    return [pscustomobject]@{
        ExitCode = $exitCode
        Output   = ($output -join [Environment]::NewLine)
    }
}

function Invoke-AzureCliJson {
    param([Parameter(Mandatory)] [string[]] $Arguments)

    $result = Invoke-AzureCli -Arguments ($Arguments + @('--output', 'json'))
    if ([string]::IsNullOrWhiteSpace($result.Output)) { return $null }
    try {
        return $result.Output | ConvertFrom-Json -Depth 100
    }
    catch {
        $jsonMatch = [regex]::Match($result.Output, '(?ms)^[\[{].*$')
        if (-not $jsonMatch.Success) { throw }
        return $jsonMatch.Value | ConvertFrom-Json -Depth 100
    }
}

function Get-ActiveAzureContext {
    $account = Invoke-AzureCliJson -Arguments @('account', 'show')
    if (-not $account.id -or -not $account.tenantId) {
        throw 'Azure CLI is not signed in to an enabled subscription.'
    }
    return $account
}

function Get-SignedInPrincipalObjectId {
    $tokenResult = Invoke-AzureCli -Arguments @(
        'account', 'get-access-token',
        '--resource', 'https://management.azure.com/',
        '--query', 'accessToken',
        '--output', 'tsv'
    )
    $parts = $tokenResult.Output.Trim().Split('.')
    if ($parts.Count -lt 2) { throw 'Unable to decode the Azure access token claims.' }
    $payload = $parts[1].Replace('-', '+').Replace('_', '/')
    while ($payload.Length % 4) { $payload += '=' }
    $claims = [Text.Encoding]::UTF8.GetString([Convert]::FromBase64String($payload)) | ConvertFrom-Json
    if (-not $claims.oid) { throw 'The Azure token does not contain an oid claim.' }
    return [string]$claims.oid
}

function Assert-WorkshopSubscription {
    param([Parameter(Mandatory)] [string] $SubscriptionId)
    $account = Get-ActiveAzureContext
    if ($account.id -ne $SubscriptionId) {
        throw 'The active Azure CLI subscription does not match -SubscriptionId. Run az account set first.'
    }
    return $account
}

function Get-TerraformEnvironment {
    param(
        [Parameter(Mandatory)] [string] $SubscriptionId,
        [string[]] $AllowedSubscriptionIds = @()
    )

    $work = Get-WorkshopWorkDirectory
    $allowedJson = if ($AllowedSubscriptionIds.Count -eq 0) {
        '[]'
    }
    else {
        ConvertTo-Json -InputObject @($AllowedSubscriptionIds) -Compress
    }
    return @{
        TF_DATA_DIR                    = (Join-Path $work '.terraform')
        TF_IN_AUTOMATION               = 'true'
        TF_INPUT                       = 'false'
        TF_VAR_subscription_id         = $SubscriptionId
        TF_VAR_allowed_subscription_ids = $allowedJson
    }
}

function Invoke-Terraform {
    param(
        [Parameter(Mandatory)] [string[]] $Arguments,
        [Parameter(Mandatory)] [string] $SubscriptionId,
        [string[]] $AllowedSubscriptionIds = @(),
        [switch] $AllowFailure
    )

    $environment = Get-TerraformEnvironment -SubscriptionId $SubscriptionId -AllowedSubscriptionIds $AllowedSubscriptionIds
    $previous = @{}
    try {
        foreach ($key in $environment.Keys) {
            $previous[$key] = [Environment]::GetEnvironmentVariable($key, 'Process')
            [Environment]::SetEnvironmentVariable($key, [string]$environment[$key], 'Process')
        }
        $output = & terraform "-chdir=$(Get-WorkshopRoot)" @Arguments 2>&1
        $exitCode = $LASTEXITCODE
    }
    finally {
        foreach ($key in $environment.Keys) {
            [Environment]::SetEnvironmentVariable($key, $previous[$key], 'Process')
        }
    }

    if ($exitCode -ne 0 -and -not $AllowFailure) {
        throw "Terraform failed (exit $exitCode): terraform $($Arguments -join ' ')`n$($output -join [Environment]::NewLine)"
    }
    return [pscustomobject]@{
        ExitCode = $exitCode
        Output   = ($output -join [Environment]::NewLine)
    }
}

function Initialize-WorkshopTerraform {
    param([Parameter(Mandatory)] [string] $SubscriptionId)
    $result = Invoke-Terraform -SubscriptionId $SubscriptionId -Arguments @('init', '-upgrade=false')
    return $result
}

function Get-WorkshopTerraformOutputs {
    param([Parameter(Mandatory)] [string] $SubscriptionId)
    $result = Invoke-Terraform -SubscriptionId $SubscriptionId -Arguments @('output', '-json')
    $json = $result.Output | ConvertFrom-Json -Depth 100
    $values = [ordered]@{}
    foreach ($property in $json.PSObject.Properties) {
        $values[$property.Name] = $property.Value.value
    }
    return [pscustomobject]$values
}

function Get-WorkshopRuntime {
    $path = Join-Path (Get-WorkshopWorkDirectory) 'runtime.json'
    if (-not (Test-Path -LiteralPath $path)) {
        throw 'The private runtime file does not exist. Run the Preflight stage first.'
    }
    return Get-Content -Raw -LiteralPath $path | ConvertFrom-Json -Depth 100
}

function Save-WorkshopRuntime {
    param([Parameter(Mandatory)] $Runtime)
    Write-WorkshopJson -Path (Join-Path (Get-WorkshopWorkDirectory) 'runtime.json') -Value $Runtime
}

function Get-WorkshopBaselineSnapshot {
    param([Parameter(Mandatory)] [string] $SubscriptionId)

    $scope = "/subscriptions/$SubscriptionId"
    $resources = Invoke-AzureCliJson -Arguments @('resource', 'list', '--subscription', $SubscriptionId, '--query', '[].id')
    $roles = Invoke-AzureCliJson -Arguments @('role', 'assignment', 'list', '--scope', $scope, '--query', '[].id')
    $diagnosticsResult = Invoke-AzureCli -Arguments @(
        'rest', '--method', 'get',
        '--url', "https://management.azure.com$scope/providers/Microsoft.Insights/diagnosticSettings?api-version=2021-05-01-preview"
    ) -AllowFailure
    $diagnosticIds = @()
    if ($diagnosticsResult.ExitCode -eq 0 -and $diagnosticsResult.Output) {
        $diagnostics = $diagnosticsResult.Output | ConvertFrom-Json -Depth 100
        $diagnosticIds = @($diagnostics.value.id)
    }

    return [ordered]@{
        capturedAtUtc = [DateTime]::UtcNow.ToString('o')
        resources     = @($resources | Sort-Object)
        roleAssignments = @($roles | Sort-Object)
        subscriptionDiagnosticSettings = @($diagnosticIds | Sort-Object)
    }
}

function Compare-WorkshopBaseline {
    param(
        [Parameter(Mandatory)] $Before,
        [Parameter(Mandatory)] $After
    )

    $resourceDiff = Compare-Object @($Before.resources) @($After.resources)
    $roleDiff = Compare-Object @($Before.roleAssignments) @($After.roleAssignments)
    $diagnosticDiff = Compare-Object @($Before.subscriptionDiagnosticSettings) @($After.subscriptionDiagnosticSettings)
    return [ordered]@{
        resourcesUnchanged                  = ($null -eq $resourceDiff)
        roleAssignmentsUnchanged            = ($null -eq $roleDiff)
        subscriptionDiagnosticsUnchanged    = ($null -eq $diagnosticDiff)
        changedResourceCount                = @($resourceDiff).Count
        changedRoleAssignmentCount          = @($roleDiff).Count
        changedSubscriptionDiagnosticCount  = @($diagnosticDiff).Count
    }
}

function Protect-WorkshopText {
    param(
        [Parameter(Mandatory)] [AllowEmptyString()] [string] $Text,
        $Runtime,
        $Outputs
    )

    $safe = $Text
    $safe = [regex]::Replace($safe, '(?i)Bearer\s+[A-Za-z0-9._~+/-]+=*', 'Bearer [redacted-token]')
    $safe = [regex]::Replace($safe, '(?i)\b[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\b', '[redacted-guid]')
    $safe = [regex]::Replace($safe, '(?i)\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b', '[redacted-email]')
    $safe = [regex]::Replace($safe, '(?<![A-Za-z0-9])(?:\d{1,3}\.){3}\d{1,3}(?:/\d{1,2})?(?![A-Za-z0-9])', '[redacted-ip]')
    $safe = [regex]::Replace($safe, '(?i)[A-Z]:\\[^\r\n\"'']+', '[redacted-local-path]')
    $safe = [regex]::Replace($safe, '(?is)-----BEGIN [^-]+PRIVATE KEY-----.*?-----END [^-]+PRIVATE KEY-----', '[redacted-private-key]')

    if ($Runtime) {
        foreach ($value in @($Runtime.subscriptionId, $Runtime.tenantId, $Runtime.principalObjectId)) {
            if ($value) { $safe = $safe.Replace([string]$value, '[redacted-id]') }
        }
    }
    if ($Outputs) {
        $replacements = @{
            storage_account_name = '[storage-account]'
            key_vault_name       = '[key-vault]'
            workspace_name       = '[log-workspace]'
            test_vm_name         = '[test-vm]'
        }
        foreach ($key in $replacements.Keys) {
            if ($Outputs.PSObject.Properties.Name -contains $key -and $Outputs.$key) {
                $safe = $safe.Replace([string]$Outputs.$key, $replacements[$key])
            }
        }
    }
    return $safe
}

function Test-PublicArtifactText {
    param([Parameter(Mandatory)] [string] $Path)

    $content = Get-Content -Raw -LiteralPath $Path
    $violations = [System.Collections.Generic.List[string]]::new()
    if ($content -match '(?i)\b[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\b') { $violations.Add('GUID') }
    if ($content -match '(?i)\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b') { $violations.Add('Email') }
    if ($content -match '(?i)Bearer\s+[A-Za-z0-9._~+/-]{20,}') { $violations.Add('BearerToken') }
    if ($content -match '(?i)-----BEGIN [^-]+PRIVATE KEY-----') { $violations.Add('PrivateKey') }
    if ($content -match '(?i)[A-Z]:\\Users\\') { $violations.Add('LocalPath') }
    return @($violations)
}
