param(
    [ValidateSet('Fixed', 'Leaf', 'FullPath')]
    [string]$Mode = 'Fixed',
    [string]$Name = 'memo_folder',
    [switch]$DryRun,
    [switch]$Check,
    [switch]$Remove,
    [string]$ProfilePath,
    [switch]$AllHosts,
    [switch]$SessionOnly
)

$ErrorActionPreference = 'Stop'
$ScriptName = Split-Path -Leaf $PSCommandPath
$BlockName = 'penwebtools-powershell-prompt'
$Begin = "# >>> $BlockName"
$End = "# <<< $BlockName"

function Write-Log {
    param([string]$Message)
    Write-Host "[$ScriptName] $Message"
}

if ($Name -match '[\r\n\p{C}]') {
    throw '-Name must not contain control characters or newlines.'
}

function Get-PromptBody {
    param([string]$PromptMode, [string]$PromptName)

    switch ($PromptMode) {
        'Fixed' {
            $safeName = $PromptName.Replace("'", "''")
            return @"
function global:prompt {
    '$safeName$ '
}
"@
        }
        'Leaf' {
            return @'
function global:prompt {
    "$(Split-Path -Leaf (Get-Location))$ "
}
'@
        }
        'FullPath' {
            return @'
function global:prompt {
    "PS $($executionContext.SessionState.Path.CurrentLocation)$('>' * ($nestedPromptLevel + 1)) "
}
'@
        }
    }
}

$PromptBody = Get-PromptBody -PromptMode $Mode -PromptName $Name

if ($SessionOnly) {
    if ($Remove) {
        throw '-SessionOnly cannot be combined with -Remove.'
    }
    if ($Check) {
        Write-Log '-SessionOnly has no persistent profile state to check.'
        exit 0
    }
    if ($DryRun) {
        Write-Log "Would set current-session prompt mode: $Mode"
        exit 0
    }
    Invoke-Expression $PromptBody
    Write-Log "Current-session prompt mode set: $Mode"
    exit 0
}

if ([string]::IsNullOrWhiteSpace($ProfilePath)) {
    if ($AllHosts) {
        $ProfilePath = $PROFILE.CurrentUserAllHosts
    } else {
        $ProfilePath = $PROFILE.CurrentUserCurrentHost
    }
}

function Get-ProfileWithoutBlock {
    param([string[]]$Lines)

    $output = New-Object System.Collections.Generic.List[string]
    $inside = $false
    $beginCount = 0
    $endCount = 0

    foreach ($line in $Lines) {
        if ($line -eq $Begin) {
            if ($inside) { throw "Managed block markers are nested in $ProfilePath." }
            $inside = $true
            $beginCount++
            continue
        }
        if ($line -eq $End) {
            if (-not $inside) { throw "Managed block end marker is missing a begin marker in $ProfilePath." }
            $inside = $false
            $endCount++
            continue
        }
        if (-not $inside) {
            $output.Add($line)
        }
    }

    if ($inside -or $beginCount -ne $endCount -or $beginCount -gt 1) {
        throw "Managed block markers are incomplete or duplicated in $ProfilePath."
    }

    return ,$output.ToArray()
}

if ((-not (Test-Path -LiteralPath $ProfilePath)) -and $Check) {
    Write-Log "PowerShell prompt is not installed: $ProfilePath"
    exit 1
}

$existing = @()
if (Test-Path -LiteralPath $ProfilePath) {
    $existing = Get-Content -LiteralPath $ProfilePath -ErrorAction Stop
}

$stripped = Get-ProfileWithoutBlock -Lines $existing
while ($stripped.Count -gt 0 -and [string]::IsNullOrWhiteSpace($stripped[$stripped.Count - 1])) {
    if ($stripped.Count -eq 1) {
        $stripped = @()
    } else {
        $stripped = @($stripped[0..($stripped.Count - 2)])
    }
}
if ($Remove) {
    $target = $stripped
} else {
    $block = @($Begin) + ($PromptBody -split "`r?`n") + @($End)
    if ($stripped.Count -gt 0) {
        $target = @($stripped) + @('') + $block
    } else {
        $target = $block
    }
}

$currentText = ($existing -join "`n")
$targetText = ($target -join "`n")

if ($currentText -eq $targetText) {
    if ($Remove) {
        Write-Log 'PowerShell prompt block is already absent.'
    } else {
        Write-Log 'PowerShell prompt block is already current.'
    }
    exit 0
}

if ($Check) {
    Write-Log 'PowerShell prompt block needs update.'
    exit 1
}

if ($DryRun) {
    if ($Remove) {
        Write-Log "Would remove managed block from $ProfilePath."
    } else {
        Write-Log "Would update managed block in $ProfilePath with mode: $Mode."
    }
    exit 0
}

$profileDir = Split-Path -Parent $ProfilePath
if (-not [string]::IsNullOrWhiteSpace($profileDir)) {
    New-Item -ItemType Directory -Force -Path $profileDir | Out-Null
}
if (Test-Path -LiteralPath $ProfilePath) {
    Copy-Item -LiteralPath $ProfilePath -Destination "$ProfilePath.bak.$(Get-Date -Format yyyyMMddHHmmss)" -Force
}

Set-Content -LiteralPath $ProfilePath -Value $target -Encoding UTF8

if ($Remove) {
    Write-Log "PowerShell prompt block removed: $ProfilePath"
} else {
    Write-Log "PowerShell prompt block updated: $ProfilePath"
}