# Migration Data Execution Module
# This module provides functions to run data migration using EDB Migration Toolkit
# Usage: Import-Module Helper-Run-Migration-Data.psm1; Start-MigrationData
# Import the utilities module for logging and common functions



# Script-level variables for deduplication
$script:LoggedMTKErrors = @()
$script:LoggedFailedTables = @()
$script:IgnoreFilePath = $null

# Function to reset deduplication tracking
function Reset-DeduplicationTracking {
    $script:LoggedMTKErrors = @()
    $script:LoggedFailedTables = @()
    $script:IgnoreFilePath = $null
}

# Function to execute SQL file with parameter replacement and error handling
function Invoke-SqlFileExecution {
    param(
        [string]$SqlFileName,
        [hashtable]$Properties,
        [hashtable]$Parameters = @{},
        [string]$Database = "postgres",
        [string]$User = $Properties['POSTGRES_ADMIN_USER'],
        [string]$Password = $Properties['POSTGRES_ADMIN_PASSWORD'],
        [bool]$ContinueOnError = $false
    )
    
    $sqlFilePath = Get-HelperScriptPath $SqlFileName
    
    if (-not (Test-Path $sqlFilePath)) {
        Show-Error -Message "SQL file not found: $sqlFilePath" -LogPrefix "MigrationData" 
        return $false
    }
    
    Show-Info -Message "Using SQL file: $sqlFilePath" -LogPrefix "MigrationData" 
    Show-Info -Message "SQL file contents:" -LogPrefix "MigrationData" 
    Get-Content $sqlFilePath | ForEach-Object { Show-Info -Message "  $_" -LogPrefix "MigrationData"  }
    
    # Read SQL content and replace parameters
    $sqlContent = Get-Content $sqlFilePath -Raw
    
    # Replace parameters in SQL content
    foreach ($key in $Parameters.Keys) {
        $sqlContent = $sqlContent -replace "\{$key\}", $Parameters[$key]
    }
    
    # Set environment variable for password
    $env:PGPASSWORD = $Password
    
    # Build psql command arguments for direct execution
    $psqlArgs = @(
        "-U", $User,
        "-h", $Properties['POSTGRES_SERVER'],
        "-p", $Properties['POSTGRES_SERVER_PORT'],
        "-d", $Database,
        "-q"
    )
    
    Show-Info -Message "Executing SQL file: $SqlFileName" -LogPrefix "MigrationData" 
    
    try {
        # Execute SQL content directly using echo and pipe to psql
        $result = $sqlContent | & psql @psqlArgs 2>&1
        $exitCode = $LASTEXITCODE
        
        # Check if there are any ERROR messages in the output
        $hasErrors = $result | Where-Object { $_ -match "ERROR\:" }
        
        if ($exitCode -eq 0 -and -not $hasErrors) {
            Show-Success -Message "Successfully executed SQL file: $SqlFileName" -LogPrefix "MigrationData" 
            return $true
        } else {
            if ($hasErrors) {
                Show-Warning -Message "SQL errors detected during execution of $SqlFileName" -LogPrefix "MigrationData" 
                $hasErrors | ForEach-Object { Show-Warning -Message "  $_" -LogPrefix "MigrationData"  }
                
                if ($ContinueOnError) {
                    Show-Info -Message "Continuing despite SQL errors (ContinueOnError = true)" -LogPrefix "MigrationData" 
                    return $true
                } else {
                    Show-Error -Message "Cannot proceed due to SQL errors. Exiting..." -LogPrefix "MigrationData" 
                    return $false
                }
            } else {
                Show-Error -Message "Command failed with exit code: $exitCode" -LogPrefix "MigrationData" 
                if (-not $ContinueOnError) {
                    Show-Error -Message "Cannot proceed due to command failure. Exiting..." -LogPrefix "MigrationData" 
                    return $false
                }
                return $true
            }
        }
    } catch {
        Show-Error -Message "Error executing SQL file $SqlFileName - $($_.Exception.Message)" -LogPrefix "MigrationData" 
        if (-not $ContinueOnError) {
            Show-Error -Message "Cannot proceed due to execution error. Exiting..." -LogPrefix "MigrationData" 
            return $false
        }
        return $true
    } finally {
        # No cleanup needed - no temp files used
    }
}

function Get-MtkInstallationPath {
    param(
        [string]$LogPrefix = "MigrationData"
    )
    $mtkPath = "C:\Program Files\edb\mtk\bin\runMTK.bat"
    if (-not (Test-Path $mtkPath)) {
        Show-Error -Message "EDB Migration Toolkit not found at: $mtkPath" -LogPrefix $LogPrefix 
        Show-Error -Message "Please ensure EDB Migration Toolkit is installed." -LogPrefix $LogPrefix 
        Invoke-CleanUpAndExit -ExitCode 1
    }
    Show-Success -Message "EDB Migration Toolkit found at: $mtkPath" -LogPrefix $LogPrefix 
    return $mtkPath
}

# Function to execute psql command
function Invoke-PostgresCommand {
    param(
        [string]$Command,
        [hashtable]$Properties
    )
    
    # Set the PGPASSWORD environment variable to avoid password prompt
    $env:PGPASSWORD = $Properties['POSTGRES_ADMIN_PASSWORD']
   
    $psqlCmd = "psql -U $($Properties['POSTGRES_ADMIN_USER']) -h $($Properties['POSTGRES_SERVER']) -p $($Properties['POSTGRES_SERVER_PORT']) -d postgres -q -c '$Command'"
    
    Show-Info -Message "Executing psql command: $psqlCmd" -LogPrefix "MigrationData" 
    
    try {
        Invoke-Expression $psqlCmd
        if ($LASTEXITCODE -eq 0) {
            Show-Success -Message "psql command executed successfully" -LogPrefix "MigrationData" 
            return $true
        } else {
            Show-Error -Message "psql command failed with exit code: $LASTEXITCODE" -LogPrefix "MigrationData" 
            return $false
        }
    } catch {
        Show-Error -Message "Failed to execute psql command: $($_.Exception.Message)" -LogPrefix "MigrationData" 
        return $false
    }
}

# Function to execute MTK command
function Start-MTKCommand {
    param(
        [string]$Arguments,
        [string]$WorkingDirectory = "C:\Program Files\edb\mtk\bin"
    )
    
    $mtkCmd = ".\runMTK.bat $Arguments"
    
    Show-Info -Message "Executing MTK command: $mtkCmd" -LogPrefix "MigrationData" 
    Show-Info -Message "Working directory: $WorkingDirectory" -LogPrefix "MigrationData" 
    Show-Info -Message "Starting MTK process..." -LogPrefix "MigrationData" 
    
    try {
        Push-Location $WorkingDirectory
        
        # Use Start-Process with Wait to ensure the process completes
        $processInfo = Start-Process -FilePath "cmd.exe" -ArgumentList "/c", $mtkCmd -Wait -PassThru -NoNewWindow
        $exitCode = $processInfo.ExitCode
        
        Pop-Location
        
        if ($exitCode -eq 0) {
            Show-Success -Message "MTK command completed successfully with exit code: $exitCode" -LogPrefix "MigrationData" 
            return $true
        } else {
            Show-Error -Message "MTK command failed with exit code: $exitCode" -LogPrefix "MigrationData" 
            return $false
        }
    } catch {
        Pop-Location
        Show-Error -Message "Failed to execute MTK command: $($_.Exception.Message)" -LogPrefix "MigrationData" 
        return $false
    }
}

# Function to find the latest MTK log file
function Get-LatestMTKLogFile {    
    $currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name
    $userName = $currentUser.Split('\')[-1]  # Extract username from domain\username format
    $logPath = "C:\Users\$userName\.enterprisedb\migration-toolkit\logs"
    
    if (-not (Test-Path $logPath)) {
        Show-Error -Message "MTK logs directory not found: $logPath" -LogPrefix "MigrationData" 
        Show-Error -Message "Please ensure the MTK logs directory exists and is accessible." -LogPrefix "MigrationData" 
        exit 1
    }
    
    $logFiles = Get-ChildItem -Path $logPath -Filter "mtk_*.log" | Sort-Object LastWriteTime -Descending
    
    if ($logFiles.Count -eq 0) {
        Show-Error -Message "No MTK log files found in: $logPath" -LogPrefix "MigrationData" 
        Show-Error -Message "Please ensure MTK has been run and log files are available." -LogPrefix "MigrationData" 
        exit 1
    }
    
    $latestLog = $logFiles[0]
    Show-Info -Message "===============================================" -LogPrefix "MTKLogFileName" 
    Show-Info -Message "MTK log file: $($latestLog.Name)" -LogPrefix "MTKLogFile" 
    Show-Info -Message "===============================================" -LogPrefix "MTKLogFileName" 

    return $latestLog.FullName
}

# Function to parse failed tables from log file
function Get-FailedTablesFromLog {
    param(
        [string]$LogFilePath
    )
    
    if (-not (Test-Path $LogFilePath)) {
        Show-Error -Message "Log file not found: $LogFilePath" -LogPrefix "MigrationData" 
        Show-Error -Message "Please ensure the log file exists and is accessible." -LogPrefix "MigrationData" 
        exit 1
    }
    
    Show-Info -Message "Parsing failed tables from log file: $LogFilePath" -LogPrefix "MigrationData" 
    
    $content = Get-Content $LogFilePath
    $failedTables = @()
    $inFailedSection = $false
    
    foreach ($line in $content) {
                
        # Original logic for parsing "List of failed objects" section
        if ($line -contains "List of failed objects") {
            $inFailedSection = $true
            continue
        }
        
        if ($line -contains "List of skipped objects") {
            $inFailedSection = $false
            break
        }
        
        if ($inFailedSection -and $line -contains "Tables") {
            continue
        }
        
        if ($inFailedSection -and $line -match "^-+$") {
            continue
        }

        # Check for MTK-17001 errors: "MTK-17001: Error Loading Data into Table: schema.table"
        if ($line -match "MTK-17001:\s*Error Loading Data into Table:\s*([A-Za-z0-9_]+)\.([^\.]+)") {
            $schemaName = $matches[1]
            $tableName = $matches[2]
            $errorKey = "$schemaName.$tableName"
            
            # Only log if we haven't seen this error before in this session
            if ($script:LoggedMTKErrors -notcontains $errorKey) {
                Show-Warning -Message "Found MTK-17001 error for table: $schemaName.$tableName" -LogPrefix "MigrationData" 
                $script:LoggedMTKErrors += $errorKey
            }
            
            # Add the table name to failed tables list
            # Apply case sensitivity logic: if no quotes, uppercase; if quotes, case sensitive
            if ($failedTables -notcontains $tableName) {
                if ($tableName -contains '"') {
                    # Quoted table name - keep case sensitive
                    $failedTables += $tableName
                } else {
                    # Unquoted table name - convert to uppercase
                    $failedTables += $tableName.ToUpper()
                }
            }
            continue
        }
        else {
            if ($inFailedSection -and $line.Trim() -ne "" -and $line -notmatch "^=") {
                $tableName = $line.Trim()
                if ($tableName -ne "Tables" -and $tableName -ne "--------------------" -and $tableName -notmatch "^\*+$") {
                    # Strip out numbered prefixes and schema names
                    # Remove numbered prefix (e.g., "1. SCHEMA.TABLE" -> "SCHEMA.TABLE")
                    if ($tableName -match "^\d+\.\s*(.+)$") {
                        $tableName = $matches[1]
                    }
                    
                    # Remove schema prefix (e.g., "AUTODEVNUNIT.TABLE" -> "TABLE" or "ax7xdev.TABLE" -> "TABLE")
                    if ($tableName -match "^[A-Za-z0-9_]+\.(.+)$") {
                        $tableName = $matches[1]
                    }
                }
            }  
        }
    }
    
    Show-Info -Message "Found $($failedTables.Count) failed tables" -LogPrefix "MigrationData" 
    if ($failedTables.Count -gt 0) {
        Show-Info -Message "Failed tables: $($failedTables -join ', ')" -LogPrefix "MigrationData" 
    }
    
    return $failedTables
}

function Remove-Tables {
    param(
        [string[]]$Tables,
        [string[]]$ExcludeTables
    )
    #display the list of both tables and exclude tables. Display if any tables are removed.
    $originalCount = $Tables.Count
    $filteredTables = $Tables | Where-Object { $_ -notin $ExcludeTables }

    # Find which tables were actually removed (tables that were in the original list but not in the filtered list)
    $removedTables = $Tables | Where-Object { $_ -in $ExcludeTables }

    $filteredCount = $originalCount - $filteredTables.Count

    if ($filteredCount -gt 0) {
        Show-Info -Message "Filtering out $filteredCount tables: $($removedTables -join ', ')" -LogPrefix "MigrationData" 
    }

    return $filteredTables

}

# Function to get table list from Oracle objects CSV
function Get-TableListFromOracleObjects {
    param(
        [hashtable]$Properties,
        [bool]$LobTables = $false
    )
    
    try {
        # Find the Oracle objects CSV file
        $oracleFile = Get-OracleCsvFilePath -OracleUser $Properties['ORACLE_USER']

        if (-not (Test-Path $oracleFile)) {
            Show-Warning -Message "Expected Oracle objects CSV file not found: $oracleFile" -LogPrefix "MigrationData" 
            return @()
        }
        
        Show-Info -Message "Reading Oracle objects from: $oracleFile" -LogPrefix "MigrationData" 
        
        # Import CSV and filter for tables
        $oracleObjects = Import-Csv -Path $oracleFile
        $tableType = if ($LobTables) { "TABLE_LOB" } else { "TABLE" }
        $tables = $oracleObjects | Where-Object { $_.objecttype -eq "$tableType" } | ForEach-Object { $_.objectname }
        
        Show-Info -Message "Found $($tables.Count) tables in Oracle schema" -LogPrefix "MigrationData" 
        return $tables
        
    } catch {
        Show-Error -Message "Failed to get table list from Oracle objects: $($_.Exception.Message)" -LogPrefix "MigrationData" 
        return @()
    }
}

# Function to get tables from Oracle using SQL file
function Get-TablesFromOracleSQL {
    param(
        [hashtable]$Properties,
        [string]$SqlFileName,
        [string]$Description = "tables"
    )
    
    try {
        # Validate input parameters
        if (-not $Properties) {
            Show-Error -Message "Properties parameter is null or empty" -LogPrefix "MigrationData" 
            return @()
        }
        
        if ([string]::IsNullOrWhiteSpace($SqlFileName)) {
            Show-Error -Message "SqlFileName parameter is null or empty" -LogPrefix "MigrationData" 
            return @()
        }
        
        $sqlFilePath = Get-HelperScriptPath $SqlFileName
        
        if (-not (Test-Path $sqlFilePath)) {
            Show-Error -Message "SQL file not found: $sqlFilePath" -LogPrefix "MigrationData" 
            return @()
        }
        
        Show-Info -Message "Reading SQL file: $SqlFileName" -LogPrefix "MigrationData" 
        
        # Read SQL content and replace parameters
        $sqlContent = Get-Content $sqlFilePath -Raw
        $sqlContent = $sqlContent -replace "&1", $Properties['ORACLE_USER']
        

        # Build SQL*Plus command for direct execution
        $sqlplusCmd = Get-OracleSqlPlusCommand -Properties $Properties

        Show-Info -Message "Executing SQL*Plus command to get $Description..." -LogPrefix "MigrationData"

        # Test Oracle connection first with a simple query
        $testCmd = "echo SELECT 1 FROM DUAL; | $(Get-OracleSqlPlusCommand -Properties $Properties)"
        Show-Info -Message "Testing Oracle connection..." -LogPrefix "MigrationData" 
        
        $testJob = Start-Job -ScriptBlock {
            param($cmd)
            & cmd /c $cmd 2>&1
        } -ArgumentList $testCmd
        
        $testCompleted = Wait-Job -Job $testJob -Timeout 10
        if (-not $testCompleted) {
            Show-Warning -Message "Oracle connection test timed out. This may indicate connection issues." -LogPrefix "MigrationData" 
            Stop-Job -Job $testJob -ErrorAction SilentlyContinue
            Remove-Job -Job $testJob -ErrorAction SilentlyContinue
        } else {
            $testOutput = Receive-Job -Job $testJob
            $testExitCode = $testJob.ExitCode
            Remove-Job -Job $testJob -ErrorAction SilentlyContinue
            
            # If exit code is null or empty, check if we got valid output
            if ([string]::IsNullOrEmpty($testExitCode) -or $testExitCode -eq "") {
                # Check if we got "1" in the output (from SELECT 1 FROM DUAL)
                if ($testOutput -match "1") {
                    $testExitCode = 0  # Assume success if we got "1"
                } else {
                    $testExitCode = 1  # Assume failure if no valid output
                }
            }
            
            if ($testExitCode -ne 0) {
                Show-Warning -Message "Oracle connection test failed. Exit code: $testExitCode" -LogPrefix "MigrationData" 
                Show-Info -Message "Connection test output: $($testOutput -join '; ')" -LogPrefix "MigrationData" 
            } else {
                Show-Info -Message "Oracle connection test successful" -LogPrefix "MigrationData" 
            }
        }
        
        # No cleanup needed - no temp files used
        
        # Execute SQL*Plus with timeout to prevent hanging
        $job = Start-Job -ScriptBlock {
            param($cmd, $sqlContent)
            $sqlContent | & cmd /c $cmd 2>&1
        } -ArgumentList $sqlplusCmd, $sqlContent
        
        # Wait for job to complete with timeout (30 seconds)
        $timeout = 30
        $startTime = Get-Date
        $completed = $false
        
        while (-not $completed -and ((Get-Date) - $startTime).TotalSeconds -lt $timeout) {
            if ($job.State -eq "Completed") {
                $completed = $true
            } else {
                Start-Sleep -Milliseconds 100
            }
        }
        
        if (-not $completed) {
            Show-Warning -Message "SQL*Plus command timed out after $timeout seconds. Stopping job..." -LogPrefix "MigrationData" 
            
            # Try to get any output that might have been generated before timeout
            try {
                $partialOutput = Receive-Job -Job $job -ErrorAction SilentlyContinue
                if ($partialOutput) {
                    Show-Info -Message "Partial output received before timeout: $($partialOutput -join '; ')" -LogPrefix "MigrationData" 
                }
            } catch {
                Show-Info -Message "No partial output available" -LogPrefix "MigrationData" 
            }
            
            Stop-Job -Job $job -ErrorAction SilentlyContinue
            Remove-Job -Job $job -ErrorAction SilentlyContinue
            return @()
        }
        
        $output = Receive-Job -Job $job
        $exitCode = $job.ExitCode
        Remove-Job -Job $job -ErrorAction SilentlyContinue
        
        # If exit code is null or empty, check if we got valid output
        if ([string]::IsNullOrEmpty($exitCode) -or $exitCode -eq "") {
            # Check if we got any valid table names in the output
            $validTables = $output | Where-Object { $_ -match '^\s*[A-Za-z_$][A-Za-z0-9_$]*\s*$' } | ForEach-Object { $_.Trim() }
            if ($validTables.Count -gt 0) {
                $exitCode = 0  # Assume success if we got valid table names
            } else {
                $exitCode = 1  # Assume failure if no valid output
            }
        }
        
         if ($exitCode -eq 0) {
             # Parse the output to extract table names
             # Handle both line-by-line and semicolon-separated formats
             $allOutput = $output -join " "
             $tableNames = $allOutput -split '[;\s]+' | Where-Object { $_ -match '^[A-Za-z_$][A-Za-z0-9_$]*$' } | ForEach-Object { $_.Trim() } | Where-Object { $_ -ne "" }
             
             # Remove duplicates and sort (case-sensitive)
             $tables = $tableNames | Sort-Object -Unique -CaseSensitive
             
             Show-Success -Message "Successfully found $($tables.Count) $Description from SQL" -LogPrefix "MigrationData" 
             
             return $tables
         } else {
             Show-Warning -Message "Failed to execute SQL to get $Description. Exit code: $exitCode" -LogPrefix "MigrationData" 
             Show-Info -Message "SQL*Plus output: $($output -join '; ')" -LogPrefix "MigrationData" 
            
            return @()
        }
        
    } catch {
        Show-Error -Message "Failed to get $Description from SQL: $($_.Exception.Message)" -LogPrefix "MigrationData" 
                
        return @()
    }
}
# Helper function to split comma-separated table names
function Split-TableNames {
    param([string[]]$TableNames)
    
    return $TableNames | 
        Where-Object { $null -ne $_ } | 
        ForEach-Object { ($_.ToString().Trim() -split ',') | ForEach-Object { $_.Trim() } } | 
        Where-Object { $_ -ne "" }
}

# Helper function to initialize ignore file path
function Initialize-IgnoreFilePath {
    param([hashtable]$Properties)
    
    if ($null -eq $script:IgnoreFilePath) {
        $timestamp = Get-Date -Format "yyyyMMdd_HHmmss"
        $schemaName = $Properties['POSTGRES_SCHEMA']
        $script:IgnoreFilePath = Join-Path $global:MigrationDirectory "ignoreoracleobjects_${schemaName}_${timestamp}.txt"
        Show-Info "Initializing ignore file: $script:IgnoreFilePath" -LogPrefix "MigrationData"
    }
    
    return $script:IgnoreFilePath
}

# Helper function to read existing tables from ignore file
function Read-IgnoreFileTables {
    param([string]$FilePath)
    
    $existingTables = @()
    
    if (Test-Path $FilePath) {
        try {
            $fileContent = Get-Content -Path $FilePath -Encoding ASCII -Raw
            if ($fileContent) {
                $existingTables = ($fileContent -split "`r?`n" | ForEach-Object { $_.Trim() } | Where-Object { $_ -ne "" })
            }
            Show-Info "Found $($existingTables.Count) existing tables in ignore file" -LogPrefix "MigrationData"
        } catch {
            Show-Warning "Could not read existing ignore file: $($_.Exception.Message)" -LogPrefix "MigrationData"
        }
    }
    
    return $existingTables
}

# Function to add successful tables to ignore file
function Add-SuccessfulTablesToIgnoreFile {
    param(
        [string[]]$ProcessedTables,
        [string[]]$FailedTables,
        [hashtable]$Properties
    )
    
    # Split comma-separated table names into individual tables
    $splitTables = Split-TableNames -TableNames $ProcessedTables
    
    # Calculate successful tables (exclude failed tables)
    $successfulTables = $splitTables | Where-Object { 
        $_ -notin $FailedTables 
    }
    
    if ($successfulTables.Count -eq 0) {
        Show-Info "No successful tables found to add to ignore file" -LogPrefix "MigrationData"
        return
    }
    
    Show-Info "Found $($successfulTables.Count) successful tables to add to ignore file" -LogPrefix "MigrationData"
    Show-Info "Successful tables: $($successfulTables -join ', ')" -LogPrefix "MigrationData"
    
    # Initialize ignore file path on first call
    $ignoreFilePath = Initialize-IgnoreFilePath -Properties $Properties
    
    # Read existing tables from file
    $existingTables = Read-IgnoreFileTables -FilePath $ignoreFilePath
    
    # Merge existing tables with new successful tables
    $allIgnoreTables = ($existingTables + $successfulTables) | Sort-Object -Unique
    
    # Write updated ignore file
    try {
        [System.IO.File]::WriteAllLines($ignoreFilePath, $allIgnoreTables, [System.Text.Encoding]::ASCII)
        Show-Success "Successfully updated ignore file with $($successfulTables.Count) new successful tables (total: $($allIgnoreTables.Count)): $ignoreFilePath" -LogPrefix "MigrationData"
    } catch {
        Show-Error "Failed to write to ignore file: $($_.Exception.Message)" -LogPrefix "MigrationData"
    }
}

# Function to process tables in batches
function Invoke-TablesInBatches {
    param(
        [string[]]$Tables,
        [int]$MaxTablesPerBatch,
        [hashtable]$Properties,
        [string]$BaseArgs
    )
    
    $allFailedTables = @()
    $totalTables = $Tables.Count
    $processedTables = 0
    
    Show-Info -Message "Processing $totalTables tables in batches of $MaxTablesPerBatch" -LogPrefix "MigrationData" 
    
    for ($i = 0; $i -lt $totalTables; $i += $MaxTablesPerBatch) {
        $batchTables = $Tables[$i..([Math]::Min($i + $MaxTablesPerBatch - 1, $totalTables - 1))]
        $batchNumber = [Math]::Floor($i / $MaxTablesPerBatch) + 1
        $totalBatches = [Math]::Ceiling($totalTables / $MaxTablesPerBatch)
        
        Show-Section -Message "Processing Batch $batchNumber of $totalBatches" -LogPrefix "MigrationData" 
        Show-Info -Message "Processing tables: $($batchTables -join ', ')" -LogPrefix "MigrationData" 
        
        # Build MTK arguments for this batch
        $tablesList = $batchTables -join ","
        $batchArgs = "$BaseArgs -tables $tablesList -targetSchema $($Properties['POSTGRES_SCHEMA']) $($Properties['ORACLE_USER'])"
        
        # Run MTK command for this batch and wait for completion
        Show-Info -Message "Starting batch $batchNumber..." -LogPrefix "MigrationData" 
        $batchStartTime = Get-Date
        
        $batchFailedTables = @()
        if (-not (Start-MTKCommand -Arguments $batchArgs)) {
            Show-Error -Message "MTK command failed for batch $batchNumber" -LogPrefix "MigrationData" 
            # Add all tables in this batch to failed tables
            $batchFailedTables = $batchTables
            $allFailedTables += $batchTables
        } else {
            # Wait a moment to ensure the log file is written
            Start-Sleep -Seconds 2
            
            # Check log for failed tables in this batch
            $logFile = Get-LatestMTKLogFile
            if ($logFile) {
                $failedTablesInBatch = Get-FailedTablesFromLog -LogFilePath $logFile
                if ($failedTablesInBatch.Count -gt 0) {
                    Show-Warning -Message "Found $($failedTablesInBatch.Count) failed tables in batch $batchNumber" -LogPrefix "MigrationData" 
                    Show-Info -Message "Failed tables in batch: $($failedTablesInBatch -join ', ')" -LogPrefix "MigrationData" 
                    $batchFailedTables = $failedTablesInBatch
                    $allFailedTables += $failedTablesInBatch
                }
            }
        }
        
        # Add successful tables after each batch
        Add-SuccessfulTablesToIgnoreFile -ProcessedTables $batchTables -FailedTables $batchFailedTables -Properties $Properties
        
        $batchEndTime = Get-Date
        $batchDuration = $batchEndTime - $batchStartTime
        $processedTables += $batchTables.Count
        
        Show-Info -Message "Batch $batchNumber completed in $($batchDuration.TotalMinutes.ToString('F2')) minutes" -LogPrefix "MigrationData" 
        Show-Info -Message "Completed $processedTables of $totalTables tables" -LogPrefix "MigrationData" 
        
        # Wait before starting the next batch to ensure proper completion
        if ($batchNumber -lt $totalBatches) {
            Show-Info -Message "Waiting 5 seconds before starting next batch..." -LogPrefix "MigrationData" 
            Start-Sleep -Seconds 5
        }
    }
    
    Show-Success -Message "All batches completed. Total failed tables: $($allFailedTables.Count)" -LogPrefix "MigrationData" 
    return $allFailedTables
}

# Function to confirm if user wants to import axiUm audit logs
function Request-ImportLogs {
    Write-Host ""
    Show-Info "Log Import Configuration:"
    Write-Host "Do you want to import axiUm audit logs from Oracle?" -ForegroundColor Cyan
    Write-Host "  This will include log tables (ELMAH`$ERROR, AXIUM_API_ERR_LOG, etc.) in the migration." -ForegroundColor Gray
    Write-Host "  Note: Log tables can be large and may significantly increase migration time." -ForegroundColor Gray
    Write-Host ""
    $importLogs = Read-Host "Import audit logs? (y/N, default: N)"
    $importLogs = $importLogs.ToLower()
    $includeLogTables = ($importLogs -eq "y" -or $importLogs -eq "yes")

    if ($includeLogTables) {
        Show-Info "Log tables will be included in migration"
    } else {
        Show-Info "Log tables will be excluded from migration"
    }
    return $includeLogTables
}

# Main function to run migration data
function Start-MigrationData {
    param(
        [hashtable]$Properties
    )
    $null = Get-MtkInstallationPath -LogPrefix "MigrationData"

    # Reset deduplication tracking for this session
    Reset-DeduplicationTracking
    
    Show-Header "Migration Data Execution"
    Show-Info -Message "This step will run data migration using EDB Migration Toolkit." -LogPrefix "MigrationData" 
    
    # Validate that properties are provided
    if ($null -eq $Properties) {
        Show-Error -Message "Properties parameter is required. Please provide the migration properties." -LogPrefix "MigrationData" 
        return $false
    }
    
    # Check if required properties are available
    $requiredProperties = @(
        "POSTGRES_ADMIN_USER", "POSTGRES_SERVER", "POSTGRES_SERVER_PORT",
        "POSTGRES_USER", "POSTGRES_SCHEMA", "ORACLE_USER", "POSTGRES_DATABASE"
    )
    
    $missingProperties = @()
    foreach ($prop in $requiredProperties) {
        if (-not $Properties.ContainsKey($prop) -or $Properties[$prop] -eq "") {
            $missingProperties += $prop
        }
    }
    
    if ($missingProperties.Count -gt 0) {
        Show-Error "Missing required properties: $($missingProperties -join ', ')" 
        Show-Error "Please run Step 1 first to configure all required properties." 
        return $false
    }
    
    # Display migration configuration
    Show-Section -Message "Data Import Configuration" -LogPrefix "MigrationData" 
    Show-Info -Message "" -LogPrefix "MigrationData" 
    Show-Info -Message "  Migration Configuration:" -LogPrefix "MigrationData" 
    Show-Info -Message "  - PostgreSQL User: $($Properties['POSTGRES_USER'])" -LogPrefix "MigrationData" 
    Show-Info -Message "  - Oracle User and Schema: $($Properties['ORACLE_USER'])" -LogPrefix "MigrationData" 
    Show-Info -Message "  - PostgreSQL Database: $($Properties['POSTGRES_DATABASE'])" -LogPrefix "MigrationData" 
    Show-Info -Message "" -LogPrefix "MigrationData" 
    
    if (Request-ImportLogs) {
        $includeLogTables = $true
    } else {
        $includeLogTables = $false
    }

    # Display log import preference
    if ($IncludeLogTables) {
        Show-Info "Log tables will be included in migration" 
    } else {
        Show-Info "Log tables will be excluded from migration" 
    }
    
    Show-Info "Proceeding with data import using $($Properties['POSTGRES_ADMIN_USER'])" 
    
    # Perform column renaming operations BEFORE migration
    Show-Section "Performing Column Renaming Operations (Pre-Migration)"
    
    # Track which setup steps succeeded for cleanup
    $grantSuperSuccess = $false
    $uppercaseColSuccess = $false
    $alterRoleSuccess = $false
    $migrationSuccess = $false
    
    # Step 1: Grant superuser privileges to POSTGRES_USER using POSTGRES_ADMIN_USER
    Show-Info "Step 1: Granting superuser privileges to $($Properties['POSTGRES_USER'])" 
    
    $grantParams = @{ "USERNAME" = $Properties['POSTGRES_USER'] }
    if (Invoke-SqlFileExecution -SqlFileName "SQL-Grant-Superuser-Privileges.sql" -Properties $Properties -Parameters $grantParams -Database $Properties['POSTGRES_DATABASE']) {
        $grantSuperSuccess = $true
    } else {
        return $false
    }
    
    # Step 2: Execute column renaming operations using POSTGRES_USER
    Show-Info "Step 2: Renaming UPPERCASE columns to lowercase using $($Properties['POSTGRES_USER'])" 
    
    if (Invoke-SqlFileExecution -SqlFileName "SQL-Rename-EDI-Columns.sql" -Properties $Properties -Database $Properties['POSTGRES_DATABASE'] -User $Properties['POSTGRES_USER'] -Password $Properties['POSTGRES_USER_PASSWORD'] -ContinueOnError $true) {
        $uppercaseColSuccess = $true
    } else {
        # Step 1 succeeded, so we need to run cleanup step 8
        if ($grantSuperSuccess) {
            Show-Info "Step 8 (Cleanup): Revoking superuser privileges from $($Properties['POSTGRES_USER'])" 
            $revokeParams = @{ "USERNAME" = $Properties['POSTGRES_USER'] }
            Invoke-SqlFileExecution -SqlFileName "SQL-Revoke-Superuser-Privileges.sql" -Properties $Properties -Parameters $revokeParams -Database $Properties['POSTGRES_DATABASE'] | Out-Null
        }
        return $false
    }
    
    # Set session_replication_role to replica
    Show-Section "Step 3: Setting session_replication_role to replica"
    $psqlCommand = "ALTER ROLE \`"$($Properties['POSTGRES_USER'])\`" SET session_replication_role = replica"
    if (Invoke-PostgresCommand -Command $psqlCommand -Properties $Properties) {
        $alterRoleSuccess = $true
    } else {
        Show-Error "Failed to set session_replication_role to replica" 
        # Step 2 succeeded, so we need to run cleanup step 7
        if ($uppercaseColSuccess) {
            Show-Info "Step 7 (Cleanup): Reversing column renaming, using $($Properties['POSTGRES_USER'])" 
            Invoke-SqlFileExecution -SqlFileName "SQL-Reverse-Rename-EDI-Columns.sql" -Properties $Properties -Database $Properties['POSTGRES_DATABASE'] -User $Properties['POSTGRES_USER'] -Password $Properties['POSTGRES_USER_PASSWORD'] -ContinueOnError $true | Out-Null
        }
        # Step 1 succeeded, so we need to run cleanup step 8
        if ($grantSuperSuccess) {
            Show-Info "Step 8 (Cleanup): Revoking superuser privileges from $($Properties['POSTGRES_USER'])" 
            $revokeParams = @{ "USERNAME" = $Properties['POSTGRES_USER'] }
            Invoke-SqlFileExecution -SqlFileName "SQL-Revoke-Superuser-Privileges.sql" -Properties $Properties -Parameters $revokeParams -Database $Properties['POSTGRES_DATABASE'] | Out-Null
        }
        return $false
    }
    
    # Run MTK command based on user choice
    Show-Section "Step 4: Running EDB Migration Toolkit"

    try {
    Show-Info -Message "Attempting to get all non-lob tables from Oracle (reading from csv)..." -LogPrefix "MigrationData" 
    # Initialize ignoreTables variable
    $ignoreTables = @()
    if (-not $IncludeLogTables) {
        # Get ignore tables from SQL-Get-All-Ignore-Tables.sql
        Show-Info -Message "Attempting to get ignore tables from Oracle..." -LogPrefix "MigrationData" 
        $ignoreTables = Get-TablesFromOracleSQL -Properties $Properties -SqlFileName "SQL-Get-All-Ignore-Tables.sql" -Description "Oracle tables to exclude from data migration"
            
        # Output ignore tables
        Show-Info -Message "Tables to exclude from data migration ($($ignoreTables.Count) tables): $($ignoreTables -join ', ')" -LogPrefix "MigrationData" 

    } else {
        Show-Info -Message "Log tables will be included, so no tables will be excluded from data migration" -LogPrefix "MigrationData" 
    }
    # Check MAX_TABLE_IMPORT property
    $maxTableImport = 0
    if ($Properties.ContainsKey('MAX_TABLE_IMPORT')) {
        $maxTableImport = [int]$Properties['MAX_TABLE_IMPORT']
    }
        

    $nonLobTables = Get-TableListFromOracleObjects -Properties $Properties -LobTables $false
    $nonLobTables = Remove-Tables -Tables $nonLobTables -ExcludeTables $ignoreTables

    $lobTables = Get-TableListFromOracleObjects -Properties $Properties -LobTables $true
    $lobTables = Remove-Tables -Tables $lobTables -ExcludeTables $ignoreTables

    $allProcessedTables = $nonLobTables + $lobTables
    if ($nonLobTables.Count -eq 0) {
        Show-Error "No tables found to migrate" 
        $migrationSuccess = $false
    } else {
        # Initialize failed tables variable
        $failedTables = @()
        
        if ($maxTableImport -eq 0) {
            # Use original logic with -allTables and -excludeTables
            Show-Info "MAX_TABLE_IMPORT is not set, selecting non-lob tables to migrate" 
            $mtkArgs = "-copyViaDBLinkOra -dataOnly -tables $nonLobTables -truncLoad -fastCopy -fetchSize 1000 -lobBatchSize 200 -cpBatchSize 16 -targetSchema $($Properties['POSTGRES_SCHEMA']) $($Properties['ORACLE_USER'])"
            if (-not (Start-MTKCommand -Arguments $mtkArgs)) {
                Show-Error "MTK command failed" 
                $migrationSuccess = $false
            } else {
                # Check log file for additional failed tables
                $logFile = Get-LatestMTKLogFile
                if ($logFile) {
                    $failedTables = Get-FailedTablesFromLog -LogFilePath $logFile
                }
                $migrationSuccess = $true
                
      # Add successful tables after all tables if maxTableImport -eq 0
     Add-SuccessfulTablesToIgnoreFile -ProcessedTables $nonLobTables -FailedTables $failedTables -Properties $Properties
            }
        } else {
            # Use table list approach
            Show-Info "MAX_TABLE_IMPORT is $maxTableImport, processing tables in batches" 
            # Build base MTK arguments (without table specification)
            $baseArgs = "-copyViaDBLinkOra -dataOnly -truncLoad -fastCopy -fetchSize 1000 -lobBatchSize 200 -cpBatchSize 16"
            
            # Process tables in batches
            $failedTables = Invoke-TablesInBatches -Tables $nonLobTables -MaxTablesPerBatch $maxTableImport -Properties $Properties -BaseArgs $baseArgs
            $migrationSuccess = $true
        }
            
        if ($failedTables.Count -gt 0) {
            Show-Warning "Found $($failedTables.Count) failed tables during batch processing" 
            Show-Info "Failed tables: $($failedTables -join ', ')" 
            Show-Info "Will retry failed tables with lob tables..." 
            $lobTables = $lobTables + $failedTables
        }
        
        # Process LOB tables and failed tables
        if ($maxTableImport -eq 0) {
            # Use original logic when MAX_TABLE_IMPORT is not set
            $lobTablesString = $lobTables -join ','
            $nonLobAndFailedTableArgs = "-dataOnly -tables $lobTablesString -truncLoad -vacuumAnalyze -lobBatchSize 200 -cpBatchSize 16 -targetSchema $($Properties['POSTGRES_SCHEMA']) $($Properties['ORACLE_USER'])"
            
            if (-not (Start-MTKCommand -Arguments $nonLobAndFailedTableArgs)) {
                Show-Error "Failed to retry failed tables" 
                $migrationSuccess = $false
            } else {
                $migrationSuccess = $true
            }
        } else {
            # Use batch processing approach
            Show-Info "MAX_TABLE_IMPORT is $maxTableImport, processing LOB and failed tables in batches" 
            # Build base MTK arguments for LOB tables (without table specification)
            $lobBaseArgs = "-dataOnly -truncLoad -vacuumAnalyze -lobBatchSize 200 -cpBatchSize 16"
            
            # Process LOB and failed tables in batches
            $lobFailedTables = Invoke-TablesInBatches -Tables $lobTables -MaxTablesPerBatch $maxTableImport -Properties $Properties -BaseArgs $lobBaseArgs
            
            # Note: Additional failed tables from LOB processing are tracked but not retried again
            if ($lobFailedTables.Count -gt 0) {
                Show-Warning "Found $($lobFailedTables.Count) failed tables during LOB batch processing" 
                Show-Info "Failed LOB tables: $($lobFailedTables -join ', ')" 
                Show-Warning -Message "IMPORTANT: Failed tables are usually due to the table not existing in PostgreSQL." -LogPrefix "MigrationData" 
                Show-Warning -Message "This typically happens when:" -LogPrefix "MigrationData" 
                Show-Warning -Message "    1. The table structure migration (Step 3) was not completed successfully" -LogPrefix "MigrationData" 
                Show-Warning -Message "    2. The table was excluded from the structure migration" -LogPrefix "MigrationData" 
                Show-Warning -Message "    3. There was an error during table creation in PostgreSQL" -LogPrefix "MigrationData" 
                Show-Info -Message "" -LogPrefix "MigrationData" 
            }
            $migrationSuccess = $true
        }
    }
    } finally {
        # Cleanup steps - run in reverse order if their corresponding setup steps succeeded
        
        # Step 5: Reset session_replication_role to origin (if step 3 succeeded)
        if ($alterRoleSuccess) {
            Show-Section "Step 5 (Cleanup): Resetting session_replication_role to origin"
            $psqlCommand = "ALTER ROLE \`"$($Properties['POSTGRES_USER'])\`" SET session_replication_role = origin"
            Invoke-PostgresCommand -Command $psqlCommand -Properties $Properties | Out-Null
        }

        # Step 6: Reverse the changes - rename edi back to EDI (uppercase) using POSTGRES_USER (if step 2 succeeded)
        if ($uppercaseColSuccess) {
            Show-Info "Step 6 (Cleanup): Reversing column renaming, using $($Properties['POSTGRES_USER'])" 
            Invoke-SqlFileExecution -SqlFileName "SQL-Reverse-Rename-EDI-Columns.sql" -Properties $Properties -Database $Properties['POSTGRES_DATABASE'] -User $Properties['POSTGRES_USER'] -Password $Properties['POSTGRES_USER_PASSWORD'] -ContinueOnError $true | Out-Null
        }
        
        # Step 7: Revoke superuser privileges from POSTGRES_USER using POSTGRES_ADMIN_USER (if step 1 succeeded)
        if ($grantSuperSuccess) {
            Show-Info "Step 7 (Cleanup): Revoking superuser privileges from $($Properties['POSTGRES_USER'])" 
            $revokeParams = @{ "USERNAME" = $Properties['POSTGRES_USER'] }
            if (-not (Invoke-SqlFileExecution -SqlFileName "SQL-Revoke-Superuser-Privileges.sql" -Properties $Properties -Parameters $revokeParams -Database $Properties['POSTGRES_DATABASE'])) {
                Show-Error "Security warning: Superuser privileges may still be active!" 
            }
        }
    }
    
    if ($migrationSuccess) {
        Show-Success "Column renaming operations completed!" 
        Show-Success "Migration data execution completed successfully!" 
        return $true
    } else {
        Show-Error "Migration data execution completed with errors. Cleanup steps have been executed." 
        return $false
    }
}

# Export the functions
Export-ModuleMember -Function Start-MigrationData, Get-MtkInstallationPath, Invoke-PostgresCommand, Start-MTKCommand, Get-LatestMTKLogFile, Get-FailedTablesFromLog, Invoke-SqlFileExecution, Get-TableListFromOracleObjects, Invoke-TablesInBatches, Reset-DeduplicationTracking, Get-TablesFromOracleSQL

