<#
.SYNOPSIS
Azure VM の Allocation Failure に対応し、複数の SKU で起動を試みる Azure Automation Runbook

.DESCRIPTION
割り当て解除状態の Azure VM に対して、指定された SKU の優先順位に従い起動を試みます。
Allocation Error が発生した場合、次の SKU で再試行し、全 SKU を試行してもダメな場合は
最初の SKU から再度ループします。

.NOTES
実行環境: Azure Automation Runbook (PowerShell Runtime 7.2)
前提条件: Managed Identity で Azure に認証可能なこと

#>

# ==================== 設定値 ====================
$WaitTimeSeconds = 60
$LoopWaitSeconds = 5

# ==================== 固定値（完全自動実行） ====================
$SubscriptionId = "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
$ResourceGroupName = "RG1,RG2"
$MaxRetryLoops = 1

# ==================== RG 別 SKU マッピング（固定値） ====================
$rgSkuMap = @{
    "RG1" = "Standard_NV6ads_A10_v5,Standard_D2as_v5"
    "RG2" = "Standard_D2as_v5,Standard_NV6ads_A10_v5"
}

# ==================== ログ出力用フォーマット関数 ====================
function Write-LogMessage {
    param(
        [string]$Message,
        [string]$Level = "INFO"
    )
    $timestamp = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    Write-Output "[$timestamp] [$Level] $Message"
}

# ==================== バッチ分割関数 ====================
function Split-IntoBatches {
    param(
        [array]$Items,
        [int]$BatchSize
    )
    $batches = @()
    for ($i = 0; $i -lt $Items.Count; $i += $BatchSize) {
        $batch = $Items[$i..([Math]::Min($i + $BatchSize - 1, $Items.Count - 1))]
        $batches += , $batch
    }
    return $batches
}

# スクリプト開始
Write-Output "========== Runbook 実行開始 =========="

try {
    Write-LogMessage "スクリプト初期化中..."

    Write-LogMessage "自動実行モード（固定値）"

    # パラメータログ出力
    Write-LogMessage "SubscriptionId: $SubscriptionId"
    Write-LogMessage "ResourceGroupName: $ResourceGroupName"
    Write-LogMessage "MaxRetryLoops: $MaxRetryLoops"
    Write-LogMessage "RG-SKU マッピング数: $($rgSkuMap.Count)"

    # パラメータ検証
    if ([string]::IsNullOrWhiteSpace($SubscriptionId)) {
        Write-LogMessage "エラー: SubscriptionId が空です" "ERROR"
        Write-Output "========== 実行完了（パラメータエラー） =========="
        return
    }

    if ([string]::IsNullOrWhiteSpace($ResourceGroupName)) {
        Write-LogMessage "エラー: ResourceGroupName が空です" "ERROR"
        Write-Output "========== 実行完了（パラメータエラー） =========="
        return
    }

    if ($rgSkuMap.Count -eq 0) {
        Write-LogMessage "エラー: RG 別の SKU マッピングが空です" "ERROR"
        Write-Output "========== 実行完了（パラメータエラー） =========="
        return
    }

    if ($MaxRetryLoops -le 0) {
        Write-LogMessage "エラー: MaxRetryLoops は 1 以上である必要があります" "ERROR"
        Write-Output "========== 実行完了（パラメータエラー） =========="
        return
    }

    # リソースグループ リストを配列に変換（カンマ区切りまたは単一値対応）
    $rgArray = @()
    if ($ResourceGroupName -like "*,*") {
        $rgArray = $ResourceGroupName -split ','
    } else {
        $rgArray = @($ResourceGroupName)
    }
    $rgArray = $rgArray | ForEach-Object { $_.Trim() }

    Write-LogMessage "リソースグループ配列生成完了: $($rgArray -join ', ')"
    Write-LogMessage "リソースグループ数: $($rgArray.Count)"

    # Azure 接続
    Write-LogMessage "Azure Automation 接続中..."
    try {
        $context = Get-AzContext -ErrorAction SilentlyContinue
        if ($null -eq $context) {
            Write-LogMessage "Managed Identity で認証中..."
            $connResult = Connect-AzAccount -Identity
            Write-LogMessage "Azure 接続成功"
        } else {
            Write-LogMessage "既存の Azure セッションが有効です"
        }
    }
    catch {
        Write-LogMessage "Azure 接続エラー: $_" "ERROR"
        Write-Output "========== 実行完了（接続エラー） =========="
        return
    }

    # サブスクリプション切り替え
    Write-LogMessage "サブスクリプション切り替え中: $SubscriptionId"
    try {
        Set-AzContext -SubscriptionId $SubscriptionId -ErrorAction Stop
        Write-LogMessage "サブスクリプション切り替え完了"
    }
    catch {
        Write-LogMessage "サブスクリプション切り替えエラー: $_" "ERROR"
        Write-Output "========== 実行完了（サブスクリプションエラー） =========="
        return
    }

    # ==================== 全 RG の統計 ====================
    $totalSuccessCount = 0
    $allFailureDetails = @()
    $rgResults = @()

    # ==================== リソースグループごとの処理 ====================
    foreach ($currentRg in $rgArray) {
        Write-LogMessage ""
        Write-LogMessage "========== リソースグループ: $currentRg の処理開始 =========="

        # RG 別 SKU リストの取得
        if (-not $rgSkuMap.ContainsKey($currentRg)) {
            Write-LogMessage "エラー: $currentRg に対応する SKU リストが未定義です" "ERROR"
            Write-LogMessage "リソースグループ $currentRg をスキップして次に進みます"
            continue
        }

        $skuListForRg = $rgSkuMap[$currentRg]
        $skuArray = @()
        if ($skuListForRg -like "*,*") {
            $skuArray = $skuListForRg -split ','
        } else {
            $skuArray = @($skuListForRg)
        }
        $skuArray = $skuArray | ForEach-Object { $_.Trim() }
        Write-LogMessage "SKU 配列生成完了: $($skuArray -join ', ')"
        Write-LogMessage "SKU 数: $($skuArray.Count)"

        # リソースグループ確認
        Write-LogMessage "リソースグループ確認中: $currentRg"
        try {
            $rg = Get-AzResourceGroup -Name $currentRg -ErrorAction Stop
            Write-LogMessage "リソースグループ確認完了: Location=$($rg.Location)"
        }
        catch {
            Write-LogMessage "リソースグループが見つかりません: $_" "ERROR"
            Write-LogMessage "リソースグループ $currentRg をスキップして次に進みます"
            continue
        }

        # VM 検出
        Write-LogMessage "割り当て解除状態の VM を検出中..."
        try {
            $vmList = @(Get-AzVM -ResourceGroupName $currentRg -Status -ErrorAction Stop |
                        Where-Object { $_.PowerState -eq "VM deallocated" })

            Write-LogMessage "検出された deallocated VM: $($vmList.Count) 台"
            
            if ($vmList.Count -eq 0) {
                Write-LogMessage "割り当て解除状態の VM が見つかりません" "WARN"
                Write-LogMessage "リソースグループ $currentRg の処理を完了します"
                continue
            }

            $vmList | ForEach-Object {
                Write-LogMessage "  - $($_.Name) (PowerState: $($_.PowerState))"
            }
        }
        catch {
            Write-LogMessage "VM 検出エラー: $_" "ERROR"
            Write-LogMessage "リソースグループ $currentRg をスキップして次に進みます"
            continue
        }

        # ==================== VM 起動リトライ処理 ====================
        $successCount = 0
        $failureDetails = @()

        for ($loopIndex = 1; $loopIndex -le $MaxRetryLoops; $loopIndex++) {
            Write-LogMessage ""
            Write-LogMessage "========== ループ $loopIndex / $MaxRetryLoops =========="

            # リトライ対象 VM を再取得
            $vmListToRetry = @(Get-AzVM -ResourceGroupName $currentRg -Status -ErrorAction SilentlyContinue |
                               Where-Object { $_.PowerState -eq "VM deallocated" })

            if ($vmListToRetry.Count -eq 0) {
                Write-LogMessage "割り当て解除状態の VM がすべて起動完了しました"
                break
            }

            Write-LogMessage "リトライ対象 VM: $($vmListToRetry.Count) 台"

        # SKU ごとに試行
        foreach ($sku in $skuArray) {
            Write-LogMessage "SKU: $sku で起動試行中..."

            $skuSuccessCount = 0
            
            # VM リストを 100 台ずつのバッチに分割
            $batches = Split-IntoBatches -Items $vmListToRetry -BatchSize 100
            $batchCount = $batches.Count
            $batchNumber = 0

            foreach ($batch in $batches) {
                $batchNumber++
                $batchStartNum = ($batchNumber - 1) * 100 + 1
                $batchEndNum = [Math]::Min($batchNumber * 100, $vmListToRetry.Count)
                
                Write-LogMessage "バッチ処理開始: VM $batchStartNum～$batchEndNum / 全 $($vmListToRetry.Count) 台 (SKU: $sku)"

                $batchSuccessCount = 0

                foreach ($vm in $batch) {
                    $vmName = $vm.Name

                    try {
                        Write-LogMessage "  VM: $vmName の起動を試みています..."

                        # VM のハードウェアプロファイルを SKU で更新
                        $vmResource = Get-AzVM -ResourceGroupName $currentRg -Name $vmName -ErrorAction Stop
                        $currentSku = $vmResource.HardwareProfile.VmSize
                        
                        if ($currentSku -ne $sku) {
                            Write-LogMessage "  SKU 変更: $currentSku -> $sku"
                            $vmResource.HardwareProfile.VmSize = $sku
                            Update-AzVM -VM $vmResource -ResourceGroupName $currentRg -ErrorAction Stop | Out-Null
                            Write-LogMessage "  ハードウェアプロファイル更新完了"
                        }

                        # VM 起動
                        Write-LogMessage "  VM を起動中..."
                        Start-AzVM -ResourceGroupName $currentRg -Name $vmName -ErrorAction Stop -WarningAction SilentlyContinue | Out-Null

                        Write-LogMessage "✓ VM: $vmName 起動コマンド成功"
                        $batchSuccessCount++
                        $successCount++
                        $skuSuccessCount++
                    }
                    catch {
                        $errorMessage = $_.Exception.Message
                        
                        # 無効な SKU エラーの場合はスキップして次の SKU を試す
                        if ($errorMessage -like "*not valid*" -or $errorMessage -like "*not valid. The valid sizes*") {
                            Write-LogMessage "スキップ - VM: $vmName (SKU: $sku - このリージョンで利用できません)" "WARN"
                            continue
                        }
                        
                        if ($errorMessage -like "*AllocationFailed*" -or $errorMessage -like "*Allocation*" -or $errorMessage -like "*capacity*") {
                            Write-LogMessage "✗ VM: $vmName 起動失敗（Allocation Error）: $errorMessage" "WARN"
                            $failureDetails += @{
                                VMName = $vmName
                                SKU    = $sku
                                Error  = $errorMessage
                                Loop   = $loopIndex
                            }
                        }
                        else {
                            Write-LogMessage "✗ VM: $vmName 起動失敗（エラー）: $errorMessage" "ERROR"
                            $failureDetails += @{
                                VMName = $vmName
                                SKU    = $sku
                                Error  = $errorMessage
                                Loop   = $loopIndex
                            }
                        }
                    }

                    if ($batch.Count -gt 1) {
                        Start-Sleep -Seconds $LoopWaitSeconds
                    }
                }

                Write-LogMessage "バッチ $batchNumber が完了しました（成功: $batchSuccessCount 台）。$WaitTimeSeconds 秒待機中..."
                Start-Sleep -Seconds $WaitTimeSeconds

                # バッチ待機後、状態確認
                Write-LogMessage "バッチ $batchNumber の VM 状態を確認中..."
                foreach ($vm in $batch) {
                    try {
                        $vmStatus = Get-AzVM -ResourceGroupName $currentRg -Name $vm.Name -Status -ErrorAction Stop
                        if ($vmStatus.PowerState -eq "VM running") {
                            Write-LogMessage "✓ VM: $($vm.Name) 状態確認完了: $($vmStatus.PowerState)"
                        }
                        else {
                            Write-LogMessage "⚠ VM: $($vm.Name) 状態: $($vmStatus.PowerState)" "WARN"
                        }
                    }
                    catch {
                        Write-LogMessage "✗ VM: $($vm.Name) の状態確認エラー: $_" "ERROR"
                    }
                }
            }

            Write-LogMessage "SKU: $sku での全 $batchCount バッチの試行が完了しました（成功: $skuSuccessCount 台）"

            # SKU ごとの試行後、再度 deallocated VM を確認
            $vmListToRetry = @(Get-AzVM -ResourceGroupName $currentRg -Status -ErrorAction SilentlyContinue |
                               Where-Object { $_.PowerState -eq "VM deallocated" })

            if ($vmListToRetry.Count -eq 0) {
                Write-LogMessage "割り当て解除状態の VM がすべて起動完了しました"
                break
            }
        }

        # ループ間の待機
        if ($loopIndex -lt $MaxRetryLoops) {
            $vmListToRetry = @(Get-AzVM -ResourceGroupName $currentRg -Status -ErrorAction SilentlyContinue |
                               Where-Object { $_.PowerState -eq "VM deallocated" })

            if ($vmListToRetry.Count -gt 0) {
                Write-LogMessage "ループ $loopIndex 完了。残り $($vmListToRetry.Count) 台。$LoopWaitSeconds 秒待機後、ループ $($loopIndex + 1) を開始します..."
                Start-Sleep -Seconds $LoopWaitSeconds
            }
        }
    }

        # ==================== リソースグループごとのレポート ====================
        Write-Output ""
        Write-Output "========== リソースグループ $currentRg の処理完了 =========="
        Write-LogMessage "成功: $successCount 台"

        if ($failureDetails.Count -gt 0) {
            Write-Output "========== リソースグループ $currentRg の失敗詳細 =========="
            $failureDetails | ForEach-Object {
                Write-LogMessage "失敗 - VM: $($_.VMName) | SKU: $($_.SKU) | ループ: $($_.Loop) | エラー: $($_.Error)" "ERROR"
            }
        }

        # 最終 VM 状態確認
        Write-Output ""
        Write-Output "========== リソースグループ $currentRg の最終 VM 状態 =========="
        $finalVMStatus = Get-AzVM -ResourceGroupName $currentRg -Status -ErrorAction SilentlyContinue
        if ($finalVMStatus) {
            $finalVMStatus | ForEach-Object {
                Write-LogMessage "VM: $($_.Name) | PowerState: $($_.PowerState)"
            }
        }

        # 全体統計に追加
        $totalSuccessCount += $successCount
        $allFailureDetails += $failureDetails
        $rgResults += @{
            ResourceGroupName = $currentRg
            SuccessCount     = $successCount
            FailureCount     = $failureDetails.Count
        }

        Write-LogMessage "リソースグループ $currentRg の処理を終了します"
    }

    # ==================== 全リソースグループの統合レポート ====================
    Write-Output ""
    Write-Output "========== 全体の実行完了 =========="
    Write-LogMessage "全リソースグループの合計成功: $totalSuccessCount 台"

    if ($rgResults.Count -gt 0) {
        Write-Output ""
        Write-Output "========== リソースグループ別の結果 =========="
        $rgResults | ForEach-Object {
            Write-LogMessage "RG: $($_.ResourceGroupName) | 成功: $($_.SuccessCount) 台 | 失敗: $($_.FailureCount) 件"
        }
    }

    if ($allFailureDetails.Count -gt 0) {
        Write-Output ""
        Write-Output "========== 全失敗詳細 =========="
        $allFailureDetails | ForEach-Object {
            Write-LogMessage "失敗 - VM: $($_.VMName) | SKU: $($_.SKU) | ループ: $($_.Loop) | RG: (該当RG)" "ERROR"
        }
    }

    Write-Output "========== Runbook 実行終了 =========="
}
catch {
    Write-Output ""
    Write-Output "========== 予期しないエラーが発生しました =========="
    Write-LogMessage "エラー: $_" "ERROR"
    if ($_.Exception.InnerException) {
        Write-LogMessage "内部エラー: $($_.Exception.InnerException.Message)" "ERROR"
    }
    if ($_.ScriptStackTrace) {
        Write-LogMessage "スタック: $($_.ScriptStackTrace)" "ERROR"
    }
    Write-Output "========== Runbook エラー終了 =========="
}
