param(
    [string]$RootPath = (Get-Location).Path,
    [string]$BaseUrl = "https://example.com",
    [string[]]$Extensions = @(".html", ".htm", ".php", ".aspx"),
    [switch]$IncludeLastMod
)

# Ensure BaseUrl ends without trailing slash
if ($BaseUrl.EndsWith("/")) { $BaseUrl = $BaseUrl.TrimEnd('/') }

# Collect files
$files = Get-ChildItem -Path $RootPath -Recurse -File |
    Where-Object { $Extensions -contains $_.Extension.ToLower() }

# Prepare XML writer settings
$settings = New-Object System.Xml.XmlWriterSettings
$settings.Indent = $true
$settings.Encoding = [System.Text.Encoding]::UTF8

$sitemapPath = Join-Path $RootPath "sitemap.xml"
$writer = [System.Xml.XmlWriter]::Create($sitemapPath, $settings)

# Write XML header and root element
$writer.WriteStartDocument()
$writer.WriteStartElement("urlset", "http://www.sitemaps.org/schemas/sitemap/0.9")

foreach ($file in $files) {
    # Compute relative path
    $relativePath = $file.FullName.Substring($RootPath.Length).TrimStart('\','/')

    # Convert backslashes to forward slashes for URLs
    $relativeUrl = $relativePath -replace '\\','/'

    # Handle default documents (e.g., index.html) as root
    if ($relativeUrl -match '^(.*?/)?index\.(html|htm|php|aspx)$') {
        $loc = $BaseUrl + '/'
    } else {
        $loc = "$BaseUrl/$relativeUrl".TrimEnd('/')
    }

    $writer.WriteStartElement("url")
    $writer.WriteElementString("loc", $loc)

    if ($IncludeLastMod) {
        $lastMod = $file.LastWriteTimeUtc.ToString("yyyy-MM-ddTHH:mm:ssZ")
        $writer.WriteElementString("lastmod", $lastMod)
    }

    $writer.WriteEndElement() # </url>
}

# Close root element and document
$writer.WriteEndElement() # </urlset>
$writer.WriteEndDocument()
$writer.Flush()
$writer.Close()

Write-Host "Sitemap generated at $sitemapPath"