Stepper.psm1

function Get-ScriptHash {
    <#
    .SYNOPSIS
        Calculates SHA256 hash of a script file.
    
    .DESCRIPTION
        Reads the content of a script file and returns its SHA256 hash.
        Used to detect if the script has been modified since the last run.
    
    .PARAMETER ScriptPath
        The path to the script file.
    
    .OUTPUTS
        System.String - SHA256 hash of the script content
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$ScriptPath
    )
    
    try {
        $content = Get-Content -Path $ScriptPath -Raw -ErrorAction Stop
        $bytes = [System.Text.Encoding]::UTF8.GetBytes($content)
        $hash = [System.Security.Cryptography.SHA256]::Create().ComputeHash($bytes)
        [System.BitConverter]::ToString($hash).Replace('-', '')
    }
    catch {
        throw "Failed to calculate hash for script '$ScriptPath': $_"
    }
}

function Get-StepIdentifier {
    <#
    .SYNOPSIS
        Gets a unique identifier for the current step based on caller location.

    .DESCRIPTION
        Analyzes the call stack to find the script and line number where New-Step was called.
        Returns an identifier in the format "filepath:line".

    .OUTPUTS
        System.String - Step identifier (e.g., "C:\script.ps1:42")
    #>

    [CmdletBinding()]
    param()

    $callStack = Get-PSCallStack

    # Walk up the call stack to find the first non-module caller
    # Stack typically looks like: [0]=Get-StepIdentifier, [1]=New-Step, [2]=UserScript
    for ($i = 0; $i -lt $callStack.Count; $i++) {
        $frame = $callStack[$i]
        $scriptName = $frame.ScriptName

        # Skip frames without a script name
        if (-not $scriptName) {
            continue
        }

        # Skip frames from the Stepper module directory
        $stepperDir = Split-Path -Path $PSScriptRoot -Parent
        if ($scriptName -like "$stepperDir/Private/*" -or
            $scriptName -like "$stepperDir\Private\*" -or
            $scriptName -like "$stepperDir/Public/*" -or
            $scriptName -like "$stepperDir\Public\*" -or
            $scriptName -like "$stepperDir/Stepper.psm1" -or
            $scriptName -like "$stepperDir\Stepper.psm1") {
            continue
        }

        # This is the user's script - return its location
        $line = $frame.ScriptLineNumber
        return "${scriptName}:${line}"
    }

    throw "Unable to determine step identifier from call stack"
}

function Get-StepperStatePath {
    <#
    .SYNOPSIS
        Gets the path to the Stepper state file for the calling script.
    
    .DESCRIPTION
        Generates a state file path based on the calling script's location.
        State files are stored in the same directory as the script with a .stepper extension.
    
    .PARAMETER ScriptPath
        The path to the script file.
    
    .OUTPUTS
        System.String - Path to the state file
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$ScriptPath
    )
    
    $scriptDir = Split-Path -Path $ScriptPath -Parent
    $scriptName = Split-Path -Path $ScriptPath -Leaf
    $stateFileName = "$scriptName.stepper"
    
    Join-Path -Path $scriptDir -ChildPath $stateFileName
}

function Read-StepperState {
    <#
    .SYNOPSIS
        Reads the Stepper state file.

    .DESCRIPTION
        Reads and deserializes the state file if it exists.
        Returns null if the file doesn't exist or can't be read.

    .PARAMETER StatePath
        The path to the state file.

    .OUTPUTS
        PSCustomObject or $null - The state object containing ScriptHash, LastCompletedStep, Timestamp, and StepperData
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$StatePath
    )

    if (-not (Test-Path -Path $StatePath)) {
        return $null
    }

    try {
        Import-Clixml -Path $StatePath -ErrorAction Stop
    }
    catch {
        Write-Warning "Failed to read state file '$StatePath': $_"
        return $null
    }
}

function Remove-StepperState {
    <#
    .SYNOPSIS
        Removes the Stepper state file.
    
    .DESCRIPTION
        Deletes the state file if it exists.
        Used when starting fresh or when the script completes successfully.
    
    .PARAMETER StatePath
        The path to the state file.
    
    .OUTPUTS
        None
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$StatePath
    )
    
    if (Test-Path -Path $StatePath) {
        try {
            Remove-Item -Path $StatePath -Force -ErrorAction Stop
        }
        catch {
            Write-Warning "Failed to remove state file '$StatePath': $_"
        }
    }
}

function Write-StepperState {
    <#
    .SYNOPSIS
        Writes the Stepper state file.

    .DESCRIPTION
        Serializes and writes the state object to disk.

    .PARAMETER StatePath
        The path to the state file.

    .PARAMETER ScriptHash
        SHA256 hash of the script content.

    .PARAMETER LastCompletedStep
        Identifier of the last successfully completed step (format: "filepath:line").

    .PARAMETER StepperData
        The $Stepper hashtable to persist.

    .OUTPUTS
        None
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$StatePath,

        [Parameter(Mandatory)]
        [string]$ScriptHash,

        [Parameter(Mandatory)]
        [string]$LastCompletedStep,

        [Parameter()]
        [hashtable]$StepperData
    )

    $state = [PSCustomObject]@{
        ScriptHash        = $ScriptHash
        LastCompletedStep = $LastCompletedStep
        Timestamp         = (Get-Date).ToString('o')
        StepperData       = $StepperData
    }

    try {
        Export-Clixml -Path $StatePath -InputObject $state -ErrorAction Stop
    }
    catch {
        Write-Warning "Failed to write state file '$StatePath': $_"
    }
}

function New-Step {
    <#
    .SYNOPSIS
        Executes a step in a resumable script.

    .DESCRIPTION
        New-Step allows scripts to be resumed from the last successfully completed step.
        On first execution, it checks for an existing state file and offers to resume.
        Each step is automatically tracked by its location (file:line) in the script.

        The script content is hashed to detect modifications. If the script changes
        between runs, the state is invalidated and execution starts fresh.

    .PARAMETER ScriptBlock
        The code to execute for this step.

    .EXAMPLE
        New-Step {
            Write-Host "Downloading files..."
            Start-Sleep -Seconds 2
        }

        New-Step {
            Write-Host "Processing data..."
            Start-Sleep -Seconds 2
        }

        If the script fails during processing, the next run will skip the download step.

    .NOTES
        State files are stored alongside the script with a .stepper extension.
        Call Stop-Stepper at the end of your script to remove the state file upon successful completion.
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory, Position = 0)]
        [scriptblock]$ScriptBlock
    )

    # Inherit verbose preference by walking the call stack
    $callStack = Get-PSCallStack
    foreach ($frame in $callStack) {
        if ($frame.InvocationInfo.BoundParameters.ContainsKey('Verbose') -and
            $frame.InvocationInfo.BoundParameters['Verbose']) {
            $VerbosePreference = 'Continue'
            break
        }
    }

    # Get step identifier and script info
    $stepId = Get-StepIdentifier
    # Extract script path from identifier (format: "path:line")
    $lastColonIndex = $stepId.LastIndexOf(':')
    $scriptPath = $stepId.Substring(0, $lastColonIndex)
    $currentHash = Get-ScriptHash -ScriptPath $scriptPath
    $statePath = Get-StepperStatePath -ScriptPath $scriptPath

    # Initialize $Stepper hashtable in calling script scope if it doesn't exist
    $callingScope = $PSCmdlet.SessionState
    try {
        $existingStepper = $callingScope.PSVariable.Get('Stepper')
        if (-not $existingStepper) {
            $callingScope.PSVariable.Set('Stepper', @{})
            Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Initialized `$Stepper hashtable"
        }
    }
    catch {
        $callingScope.PSVariable.Set('Stepper', @{})
        Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Initialized `$Stepper hashtable"
    }

    # Initialize module state on first call in this session
    if (-not $script:StepperSessionState) {
        # Verify the script contains Stop-Stepper
        $scriptContent = Get-Content -Path $scriptPath -Raw
        if ($scriptContent -notmatch 'Stop-Stepper') {
            Write-Host "Script '$scriptPath' does not call Stop-Stepper." -ForegroundColor Yellow
            $response = Read-Host "Add 'Stop-Stepper' to the end of the script? (Y/N)"

            if ($response -eq 'Y' -or $response -eq 'y') {
                Write-Host "Adding 'Stop-Stepper' to the end of the script..." -ForegroundColor Yellow

                # Add Stop-Stepper to the end of the script
                $updatedContent = $scriptContent.TrimEnd()
                if (-not $updatedContent.EndsWith("`n")) {
                    $updatedContent += "`n"
                }
                $updatedContent += "`nStop-Stepper`n"

                Set-Content -Path $scriptPath -Value $updatedContent -NoNewline

                Write-Host "Stop-Stepper added. Please run the script again." -ForegroundColor Green

                # Exit this execution - the script will need to be run again
                throw "Script modified to include Stop-Stepper. Please run the script again."
            }
            else {
                Write-Warning "Continuing without Stop-Stepper. State file will not be cleaned up automatically."
            }
        }

        $script:StepperSessionState = @{
            Initialized       = $false
            RestoreMode       = $false
            TargetStep        = $null
            CurrentScriptPath = $scriptPath
            CurrentScriptHash = $currentHash
            StatePath         = $statePath
        }
    }

    # First step: Check for existing state and prompt user
    if (-not $script:StepperSessionState.Initialized) {
        $script:StepperSessionState.Initialized = $true

        $existingState = Read-StepperState -StatePath $statePath

        # Try to load persisted $Stepper data from state
        if ($existingState -and $existingState.StepperData) {
            try {
                $callingScope.PSVariable.Set('Stepper', $existingState.StepperData)
                $variableNames = ($existingState.StepperData.Keys | Sort-Object) -join ', '
                Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Loaded `$Stepper data from disk ($variableNames)"
            }
            catch {
                Write-Warning "Failed to load persisted `$Stepper data: $_"
            }
        }

        if ($existingState) {
            # Check if script has been modified
            if ($existingState.ScriptHash -ne $currentHash) {
                Write-Host "Script has been modified since last run. Starting fresh." -ForegroundColor Yellow
                Remove-StepperState -StatePath $statePath
            }
            else {
                # Count total steps in the script by finding all New-Step calls
                $scriptContent = Get-Content -Path $scriptPath -Raw
                $stepMatches = [regex]::Matches($scriptContent, '^\s*New-Step\s+\{', [System.Text.RegularExpressions.RegexOptions]::Multiline)
                $totalSteps = $stepMatches.Count

                # Find all step line numbers to determine which step number we're on
                $stepLines = @()
                $lineNumber = 1
                foreach ($line in (Get-Content -Path $scriptPath)) {
                    if ($line -match '^\s*New-Step\s+\{') {
                        $stepLines += "${scriptPath}:${lineNumber}"
                    }
                    $lineNumber++
                }

                # Find the index of the last completed step
                $lastStep = $existingState.LastCompletedStep
                $lastStepIndex = $stepLines.IndexOf($lastStep)
                $nextStepNumber = $lastStepIndex + 2  # +1 for next step, +1 because index is 0-based

                $timestamp = [DateTime]::Parse($existingState.Timestamp).ToString('yyyy-MM-dd HH:mm:ss')

                # Get available variable names from StepperData
                $availableVars = if ($existingState.StepperData -and $existingState.StepperData.Count -gt 0) {
                    ($existingState.StepperData.Keys | Sort-Object) -join ', '
                }
                else {
                    'None'
                }

                Write-Host ""
                Write-Host "Incomplete script run detected!" -ForegroundColor Cyan
                Write-Host ""
                Write-Host "Total Steps: $totalSteps" -ForegroundColor Cyan
                Write-Host "Steps Completed: $($lastStepIndex + 1)" -ForegroundColor Cyan
                Write-Host "Variables: $availableVars" -ForegroundColor Cyan
                Write-Host "Last Activity: $timestamp" -ForegroundColor Cyan
                Write-Host ""

                if ($nextStepNumber -le $totalSteps) {
                    # Get the script name and next step line number
                    $scriptName = Split-Path $scriptPath -Leaf
                    $nextStepId = $stepLines[$lastStepIndex + 1]
                    $nextStepLine = ($nextStepId -split ':')[-1]

                    $response = Read-Host "Resume $scriptName from Line ${nextStepLine}? (Y/n)"

                    if ($response -eq '' -or $response -eq 'Y' -or $response -eq 'y') {
                        Write-Host "Resuming from step $nextStepNumber..." -ForegroundColor Green
                        $script:StepperSessionState.RestoreMode = $true
                        $script:StepperSessionState.TargetStep = $lastStep
                    }
                    else {
                        Write-Host "Starting fresh..." -ForegroundColor Yellow
                        Remove-StepperState -StatePath $statePath
                    }
                }
                else {
                    Write-Host "All steps were completed. Starting fresh..." -ForegroundColor Yellow
                    Remove-StepperState -StatePath $statePath
                }
                Write-Host ""
            }
        }
    }

    # Determine if we should execute this step
    $shouldExecute = $true

    if ($script:StepperSessionState.RestoreMode) {
        # Format step identifier for display messages
        $stepIdParts = $stepId -split ':'
        $scriptName = Split-Path $stepIdParts[0] -Leaf
        $displayStepId = "${scriptName}:$($stepIdParts[1])"

        # In restore mode: skip steps until we reach the target
        if ($stepId -eq $script:StepperSessionState.TargetStep) {
            # This is the last completed step, skip it and disable restore mode
            $script:StepperSessionState.RestoreMode = $false
            $shouldExecute = $false
        }
        elseif ($script:StepperSessionState.RestoreMode) {
            # Still skipping
            Write-Host "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Skipping step: $displayStepId" -ForegroundColor DarkGray
            $shouldExecute = $false
        }
    }

    # Execute the step if needed
    if ($shouldExecute) {
        # Format step identifier for display (scriptname:line instead of full path)
        $stepIdParts = $stepId -split ':'
        $scriptName = Split-Path $stepIdParts[0] -Leaf
        $displayStepId = "${scriptName}:$($stepIdParts[1])"

        # Calculate step number (X/Y)
        $scriptContent = Get-Content -Path $scriptPath -Raw
        $stepMatches = [regex]::Matches($scriptContent, '^\s*New-Step\s+\{', [System.Text.RegularExpressions.RegexOptions]::Multiline)
        $totalSteps = $stepMatches.Count

        # Find all step line numbers
        $stepLines = @()
        $lineNumber = 1
        foreach ($line in (Get-Content -Path $scriptPath)) {
            if ($line -match '^\s*New-Step\s+\{') {
                $stepLines += "${scriptPath}:${lineNumber}"
            }
            $lineNumber++
        }
        $currentStepNumber = $stepLines.IndexOf($stepId) + 1

        Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Executing step $currentStepNumber/$totalSteps ($displayStepId)"

        # Show current $Stepper data
        try {
            $stepperData = $callingScope.PSVariable.Get('Stepper').Value
            if ($stepperData -and $stepperData.Count -gt 0) {
                $variableNames = ($stepperData.Keys | Sort-Object) -join ', '
                Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Available `$Stepper data ($variableNames)"
            }
        }
        catch {
            # Ignore if unable to read $Stepper
        }

        try {
            & $ScriptBlock

            # Update state file after successful execution (including $Stepper data)
            $stepperData = $callingScope.PSVariable.Get('Stepper').Value
            Write-StepperState -StatePath $statePath -ScriptHash $currentHash -LastCompletedStep $stepId -StepperData $stepperData
            Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Step $currentStepNumber/$totalSteps completed ($displayStepId)"

            if ($stepperData -and $stepperData.Count -gt 0) {
                Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Saved `$Stepper data ($($stepperData.Count) items)"
            }
        }
        catch {
            Write-Error "Step failed at $stepId : $_"
            throw
        }
    }
}

function Stop-Stepper {
    <#
    .SYNOPSIS
        Stops Stepper and clears the state file for the calling script.

    .DESCRIPTION
        Removes the state file, typically called at the end of a script
        when it completes successfully. This ensures the next run starts fresh.

    .EXAMPLE
        # At the end of your script:
        New-Step { Write-Host "Final step" }
        Stop-Stepper

    .NOTES
        This function automatically determines which script called it and
        removes the corresponding state file.
    #>

    [CmdletBinding()]
    param()

    # Inherit verbose preference by walking the call stack
    $callStack = Get-PSCallStack
    foreach ($frame in $callStack) {
        if ($frame.InvocationInfo.BoundParameters.ContainsKey('Verbose') -and
            $frame.InvocationInfo.BoundParameters['Verbose']) {
            $VerbosePreference = 'Continue'
            break
        }
    }

    Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] All steps complete. Cleaning up..."

    try {
        $callStack = Get-PSCallStack

        # Find the calling script (skip this function)
        for ($i = 1; $i -lt $callStack.Count; $i++) {
            $frame = $callStack[$i]
            $scriptPath = $frame.ScriptName

            # Skip frames without a script name
            if (-not $scriptPath) {
                continue
            }

            # Skip frames from within the Stepper module
            if ($scriptPath -like '*Stepper.psm1' -or
                $scriptPath -like '*\Private\*.ps1' -or
                $scriptPath -like '*\Public\*.ps1') {
                continue
            }

            # Found the user's script
            $statePath = Get-StepperStatePath -ScriptPath $scriptPath
            Remove-StepperState -StatePath $statePath
            $scriptName = Split-Path $scriptPath -Leaf
            Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Cleared Stepper state for $scriptName"

            # Clear the session state
            if ($script:StepperSessionState) {
                $script:StepperSessionState = $null
            }

            Write-Verbose "[$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')][Stepper] Cleanup complete!"

            return
        }

        Write-Warning "Unable to determine calling script from call stack"
    }
    catch {
        Write-Error "Failed to clear Stepper state: $_"
    }
}



# Export functions and aliases as required
Export-ModuleMember -Function @('New-Step', 'Stop-Stepper') -Alias @()