param(
    [Parameter(Mandatory=$true)]
    [string]$LiveMountRoot,

    [Parameter(Mandatory=$false)]
    [string]$DatabaseName,

    [Parameter(Mandatory=$true)]
    [ValidateSet("catalog", "restore", "verify")]
    [string]$Action,

    [Parameter(Mandatory=$true)]
    [string]$SqlInstance = "sqlfcsql\TESTINST"
)

# Function to catalog backups
function Catalog-Backups {
    param([string]$Path)

    $backupFiles = Get-ChildItem -Path $Path -Recurse -File | Where-Object { $_.Extension -eq '.bak' -or $_.Extension -eq '.trn' }

    $catalog = @{}

    foreach ($file in $backupFiles) {
        $baseName = $file.BaseName
        $parts = $baseName -split '_'
        if ($parts.Length -lt 4) { continue }

        # Find the type position (FULL, DIFF, LOG)
        $typeIndex = -1
        $validTypes = @('FULL', 'DIFF', 'LOG')
        for ($i = 0; $i -lt $parts.Length; $i++) {
            if ($validTypes -contains $parts[$i]) {
                $typeIndex = $i
                break
            }
        }
        if ($typeIndex -eq -1 -or ($parts.Length - $typeIndex) -lt 3) { continue }

        # Assume parts[0] is prefix if typeIndex > 1, else dbName is parts[0]
        if ($typeIndex -eq 1) {
            $dbName = $parts[0]
        } else {
            $dbName = $parts[1..($typeIndex-1)] -join '_'
        }
        $type = $parts[$typeIndex]
        $dateStr = $parts[$typeIndex + 1]
        $timeStr = $parts[$typeIndex + 2]
        $stripe = 0
        if ($parts.Length -gt $typeIndex + 3) {
            $stripe = [int]$parts[$typeIndex + 3]
        }
        $key = "$dateStr$timeStr"

        if (-not $catalog.ContainsKey($dbName)) { $catalog[$dbName] = @{} }
        if (-not $catalog[$dbName].ContainsKey($type)) { $catalog[$dbName][$type] = @{} }
        if (-not $catalog[$dbName][$type].ContainsKey($key)) { $catalog[$dbName][$type][$key] = @() }
        $catalog[$dbName][$type][$key] += @{File = $file; Stripe = $stripe}
    }

    return $catalog
}

# Function to report catalog
function Report-Catalog {
    param([hashtable]$Catalog)

    Write-Host "Database Backups Catalog:"
    Write-Host "========================="

    foreach ($db in $catalog.Keys | Sort-Object) {
        Write-Host "Database: $db"
        foreach ($type in $catalog[$db].Keys | Sort-Object) {
            Write-Host "  Type: $type"
            foreach ($key in $catalog[$db][$type].Keys | Sort-Object -Descending) {
                Write-Host "    Backup: $key"
                $files = $catalog[$db][$type][$key] | Sort-Object { $_.Stripe }
                foreach ($item in $files) {
                    Write-Host "      $($item.File.FullName)"
                }
            }
        }
        Write-Host ""
    }
}

# Function to restore database
function Restore-Database {
    param([string]$DbName, [hashtable]$Catalog, [string]$Instance)

    if (-not $catalog.ContainsKey($DbName)) {
        Write-Error "Database $DbName not found in catalog"
        return
    }

    $dbCatalog = $catalog[$DbName]

    # Find the latest FULL backup
    if (-not $dbCatalog.ContainsKey('FULL')) {
        Write-Error "No FULL backup found for database $DbName"
        return
    }

    $latestFullKey = $dbCatalog['FULL'].Keys | Sort-Object -Descending | Select-Object -First 1
    $fullFiles = $dbCatalog['FULL'][$latestFullKey] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }

    # Restore FULL with NORECOVERY
    $fileList = $fullFiles | ForEach-Object { "DISK = '$($_.FullName)'" }
    $restoreQuery = "RESTORE DATABASE [$DbName] FROM $($fileList -join ', ') WITH NORECOVERY"
    Write-Host "Restoring FULL backup for $DbName..."
    Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop

    # Apply DIFF backups after the FULL
    if ($dbCatalog.ContainsKey('DIFF')) {
        $diffKeys = $dbCatalog['DIFF'].Keys | Where-Object { $_ -gt $latestFullKey } | Sort-Object
        foreach ($key in $diffKeys) {
            $diffFiles = $dbCatalog['DIFF'][$key] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
            $fileList = $diffFiles | ForEach-Object { "DISK = '$($_.FullName)'" }
            $restoreQuery = "RESTORE DATABASE [$DbName] FROM $($fileList -join ', ') WITH NORECOVERY"
            Write-Host "Applying DIFF backup $key for $DbName..."
            Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop
        }
    }

    # Apply LOG backups after the FULL
    if ($dbCatalog.ContainsKey('LOG')) {
        $logKeys = $dbCatalog['LOG'].Keys | Where-Object { $_ -gt $latestFullKey } | Sort-Object
        foreach ($key in $logKeys) {
            $logFiles = $dbCatalog['LOG'][$key] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
            $fileList = $logFiles | ForEach-Object { "DISK = '$($_.FullName)'" }
            $restoreQuery = "RESTORE LOG [$DbName] FROM $($fileList -join ', ') WITH NORECOVERY"
            Write-Host "Applying LOG backup $key for $DbName..."
            Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop
        }
    }

    # Final recovery
    $restoreQuery = "RESTORE DATABASE [$DbName] WITH RECOVERY"
    Write-Host "Finalizing restore for $DbName..."
    Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop

    Write-Host "Restore completed for $DbName"
}

# Function to print backup summary
function Print-BackupSummary {
    param([string]$DbName, [hashtable]$Headers)

    Write-Host "Backup Summary for database: $DbName"
    Write-Host "==================================="

    # FULL backups
    if ($Headers.ContainsKey('FULL')) {
        Write-Host "FULL Backups:"
        foreach ($item in $Headers['FULL'] | Sort-Object { $_.Key }) {
            $header = $item.Header
            Write-Host "  Date: $($header.BackupFinishDate) | LSN Range: $($header.FirstLSN) - $($header.LastLSN)"
        }
    }

    # DIFF backups
    if ($Headers.ContainsKey('DIFF')) {
        Write-Host "DIFFERENTIAL Backups:"
        foreach ($item in $Headers['DIFF'] | Sort-Object { $_.Key }) {
            $header = $item.Header
            Write-Host "  Date: $($header.BackupFinishDate) | Base LSN: $($header.DifferentialBaseLSN) | LSN Range: $($header.FirstLSN) - $($header.LastLSN)"
        }
    }

    # LOG backups
    if ($Headers.ContainsKey('LOG')) {
        $logItems = $Headers['LOG'] | Sort-Object { $_.Key }
        if ($logItems.Count -gt 0) {
            $firstLog = $logItems[0].Header
            $lastLog = $logItems[-1].Header
            Write-Host "LOG Backups:"
            Write-Host "  Point-in-Time Range: $($firstLog.BackupStartDate) to $($lastLog.BackupFinishDate)"
            Write-Host "  LSN Range: $($firstLog.FirstLSN) - $($lastLog.LastLSN)"

            # Check for gaps
            $gaps = @()
            for ($i = 1; $i -lt $logItems.Count; $i++) {
                $prevLast = $logItems[$i-1].Header.LastLSN
                $currFirst = $logItems[$i].Header.FirstLSN
                if ($prevLast -ne $currFirst) {
                    $gaps += "Gap between $($logItems[$i-1].Key) (LSN $($prevLast)) and $($logItems[$i].Key) (LSN $($currFirst))"
                }
            }
            if ($gaps.Count -gt 0) {
                Write-Host "  *** MISSING RANGES ***"
                foreach ($gap in $gaps) {
                    Write-Host "    $gap"
                }
            } else {
                Write-Host "  No gaps detected in LOG sequence"
            }
        }
    }

    if (-not ($Headers.ContainsKey('FULL') -or $Headers.ContainsKey('DIFF') -or $Headers.ContainsKey('LOG'))) {
        Write-Host "No backup headers retrieved"
    }
}

# Function to verify backups
function Verify-Backups {
    param([hashtable]$Catalog, [string]$Instance)

    foreach ($db in $catalog.Keys | Sort-Object) {
        Write-Host "Verifying backups for database: $db"
        $headers = @{}
        foreach ($type in $catalog[$db].Keys | Sort-Object) {
            foreach ($key in $catalog[$db][$type].Keys | Sort-Object) {
                $files = $catalog[$db][$type][$key] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
                $fileList = $files | ForEach-Object { "DISK = '$($_.FullName)'" }
                $verifyQuery = "RESTORE VERIFYONLY FROM $($fileList -join ', ')"
                Write-Host "Verifying $type backup $key for $db..."
                $verified = $false
                try {
                    Invoke-Sqlcmd -ServerInstance $Instance -Query $verifyQuery -QueryTimeout 0 -ErrorAction Stop
                    Write-Host "Verification successful for $type $key"
                    $verified = $true
                } catch {
                    Write-Host "Verification failed for $type $key : $($_.Exception.Message)"
                }

                # Get backup header information only if verified
                if ($verified) {
                    $header = Get-BackupInfo -Files $files -Instance $Instance
                    if ($header) {
                        Write-Host "  Backup Details:"
                        Write-Host "    Start Date: $($header.BackupStartDate)"
                        Write-Host "    Finish Date: $($header.BackupFinishDate)"
                        Write-Host "    First LSN: $($header.FirstLSN)"
                        Write-Host "    Last LSN: $($header.LastLSN)"
                        if ($header.DatabaseBackupLSN) {
                            Write-Host "    Database Backup LSN: $($header.DatabaseBackupLSN)"
                        }
                        if ($header.DifferentialBaseLSN) {
                            Write-Host "    Differential Base LSN: $($header.DifferentialBaseLSN)"
                        }

                        # Collect headers for summary
                        if (-not $headers.ContainsKey($type)) { $headers[$type] = @() }
                        $headers[$type] += @{ Key = $key; Header = $header }
                    }
                }
            }
        }

        # Print summary
        Print-BackupSummary -DbName $db -Headers $headers
        Write-Host ""
    }
}

# Function to get backup header information
function Get-BackupInfo {
    param([System.IO.FileInfo[]]$Files, [string]$Instance)

    $fileList = $Files | ForEach-Object { "DISK = '$($_.FullName)'" }
    $headerQuery = "RESTORE HEADERONLY FROM $($fileList -join ', ')"

    try {
        $header = Invoke-Sqlcmd -ServerInstance $Instance -Query $headerQuery -QueryTimeout 0 -ErrorAction Stop
        return $header
    } catch {
        Write-Warning "Failed to get header for $($Files[0].Name): $($_.Exception.Message)"
        return $null
    }
}

# Main script
if ($Action -eq "catalog") {
    $catalog = Catalog-Backups -Path $LiveMountRoot
    if ($DatabaseName) {
        $filteredCatalog = @{}
        if ($catalog.ContainsKey($DatabaseName)) {
            $filteredCatalog[$DatabaseName] = $catalog[$DatabaseName]
        }
        Report-Catalog -Catalog $filteredCatalog
    } else {
        Report-Catalog -Catalog $catalog
    }
} elseif ($Action -eq "restore") {
    if (-not $DatabaseName) {
        Write-Error "DatabaseName is required for restore action"
        exit 1
    }

    $catalog = Catalog-Backups -Path $LiveMountRoot
    Restore-Database -DbName $DatabaseName -Catalog $catalog -Instance $SqlInstance
} elseif ($Action -eq "verify") {
    $catalog = Catalog-Backups -Path $LiveMountRoot
    if ($DatabaseName) {
        $filteredCatalog = @{}
        if ($catalog.ContainsKey($DatabaseName)) {
            $filteredCatalog[$DatabaseName] = $catalog[$DatabaseName]
        }
    } else {
        $filteredCatalog = $catalog
    }
    Verify-Backups -Catalog $filteredCatalog -Instance $SqlInstance
}