# Define the FolderIconSetter class
Add-Type -TypeDefinition @"
using System;
using System.Runtime.InteropServices;

public class FolderIconSetter {
    [DllImport("shell32.dll", CharSet = CharSet.Unicode)]
    public static extern int SHGetSetFolderCustomSettings(ref SHFOLDERCUSTOMSETTINGS pfcs, string pszPath, int dwReadWrite);

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct SHFOLDERCUSTOMSETTINGS {
        public int dwSize;
        public uint dwMask;
        public IntPtr pvid;
        public string pszWebViewTemplate;
        public uint cchWebViewTemplate;
        public string pszWebViewTemplateVersion;
        public string pszInfoTip;
        public uint cchInfoTip;
        public IntPtr pclsid;
        public uint dwFlags;
        public string pszIconFile;
        public uint cchIconFile;
        public int iIconIndex;
        public string pszLogo;
        public uint cchLogo;
    }

    public const uint FCSM_ICONFILE = 0x00000010;
    public const uint FCS_FORCEWRITE = 0x00000002;
    public const uint FCS_READ = 0x00000001;

    public static void SetFolderIcon(string folderPath, string iconFileName) {
        SHFOLDERCUSTOMSETTINGS settings = new SHFOLDERCUSTOMSETTINGS();
        settings.dwSize = Marshal.SizeOf(settings);
        settings.dwMask = FCSM_ICONFILE;
        settings.pszIconFile = iconFileName;
        settings.cchIconFile = (uint)iconFileName.Length;

        int result = SHGetSetFolderCustomSettings(ref settings, folderPath, (int)(FCS_READ | FCS_FORCEWRITE));
    }
}
"@

# Function to set folder icon and create or update Desktop.ini
function Set-CustomFolderIcons {
    # Get all subfolders
    $subfolders = Get-ChildItem -Path $scriptDirectory -Recurse -Directory

    $customizedFoldersCount = 0
    $nonConformingFolders = @()  # Array to store non-conforming folder names
    $nonConformingFoldersFixedCount = 0  # Count of non-conforming folders fixed
    $alreadyCustomizedFoldersCount = 0  # Count of folders already customized

    # Counter for progress indication
    $foldersProcessed = 0
    $totalFoldersWithIcons = 0

    # Initialize start time
    $start = Get-Date

    foreach ($folder in $subfolders) {
        try {
            $iconFiles = Get-ChildItem -Path $folder.FullName -Filter "*.ico" -Force
            if ($iconFiles.Count -gt 0) {
                $totalFoldersWithIcons++
                $folderIndex = $totalFoldersWithIcons

                $folderName = $folder.Name
                $truncatedFolderName = $folderName.Substring(0, [Math]::Min($folderName.Length, 100))
                if ($folderName.Length -gt 90) {
                    $truncatedFolderName += "..."
                }

                Write-Host "Processing folder ${folderIndex}/${totalFoldersWithIcons}" -ForegroundColor Cyan
                Write-Host "  Folder Name: " -ForegroundColor Green -NoNewline
                Write-Host "${truncatedFolderName}" -ForegroundColor White

                # Select a random icon from the folder
                $selectedIcon = Get-Random -InputObject $iconFiles
                $iconFileName = $selectedIcon.Name

                if ($iconFileName -eq "icon.ico") {
                    Write-Host "  Selected Icon: " -ForegroundColor Green -NoNewline
                    Write-Host "${iconFileName} " -ForegroundColor White -NoNewline
                    Write-Host "(default icon 'icon.ico' found)" -ForegroundColor Yellow
                } else {
                    Write-Host "  Selected Icon: " -ForegroundColor Green -NoNewline
                    Write-Host "${iconFileName}" -ForegroundColor White
                }

                Write-Host "" # Add a blank line after Selected Icon

                $folderPath = $folder.FullName

                # Check if Desktop.ini file exists
                $desktopIniPath = Join-Path $folderPath "Desktop.ini"
                if (!(Test-Path $desktopIniPath -PathType Leaf)) {
                    # Set the folder icon using the custom function
                    [FolderIconSetter]::SetFolderIcon($folderPath, $iconFileName)

                    # Create Desktop.ini file
                    @"
[.ShellClassInfo]
IconResource=$iconFileName,0
"@ | Out-File -Encoding ascii -FilePath $desktopIniPath -Force

                    # Set attributes for Desktop.ini file
                    Set-ItemProperty -Path $desktopIniPath -Name Attributes -Value "ReadOnly, Hidden"

                    $customizedFoldersCount++
                } else {
                    $alreadyCustomizedFoldersCount++
                    # Read existing Desktop.ini content
                    $desktopIniContent = Get-Content -Path $desktopIniPath -Raw

                    # Check if Desktop.ini content matches expected format
                    if ($desktopIniContent -notmatch "\[.ShellClassInfo\]`r?`nIconResource=.+") {
                        # Add folder name to non-conforming folders list
                        $nonConformingFolders += $folderName

                        # Check if existing Desktop.ini contains localizedresourcename
                        $localizedResourceName = $null
                        if ($desktopIniContent -match "(?m)localizedresourcename=.+") {
                            $localizedResourceName = $matches[0]
                        }

                        # Delete existing Desktop.ini
                        Remove-Item -Path $desktopIniPath -Force

                        # Create new Desktop.ini file
                        @"
[.ShellClassInfo]
IconResource=$iconFileName,0
$localizedResourceName
[ViewState]
FolderType=Videos
"@ | Out-File -Encoding ascii -FilePath $desktopIniPath -Force

                        # Set attributes for Desktop.ini file
                        Set-ItemProperty -Path $desktopIniPath -Name Attributes -Value "ReadOnly, Hidden"

                        # Increment count of non-conforming folders fixed
                        $nonConformingFoldersFixedCount++
                    }
                }
            }
        } catch {
            # Ignore errors related to file access for Desktop.ini
            if ($_.Exception.Message -match "Access to the path.*Desktop.ini.*is denied.") {
                continue
            } else {
                Write-Output "Error setting folder icon for: $folderName"
                Write-Output $_.Exception.Message
            }
        }
    }

    # Calculate new folders customized
    $newFoldersCount = $totalFoldersWithIcons - $alreadyCustomizedFoldersCount

    # Output summary information
    Write-Host "Finished customizing folders!" -ForegroundColor Green
    $end = Get-Date
    $timeTaken = $end - $start
    $formattedTime = ""
    if ($timeTaken.TotalSeconds -ge 60) {
        $minutes = [Math]::Floor($timeTaken.TotalMinutes)
        $seconds = $timeTaken.TotalSeconds % 60
        if ($minutes -lt 2) {
            $formattedTime = "$minutes minute"
        } else {
            $formattedTime = "$minutes minutes"
        }
        if ($seconds -ne 0) {
            $formattedTime += " and " + "{0:N2}" -f $seconds + " seconds"
        }
    } else {
        $formattedTime = "{0:N2}" -f $timeTaken.TotalSeconds + " seconds"
    }
    Write-Host "Time taken: $formattedTime" -ForegroundColor Green
    Write-Host "Folders processed: $totalFoldersWithIcons" -ForegroundColor Green
    Write-Host "New folders customized: $newFoldersCount" -ForegroundColor Green
    Write-Host "Folders already customized: $alreadyCustomizedFoldersCount" -ForegroundColor Green
    Write-Host "Non-conforming folders fixed: $nonConformingFoldersFixedCount" -ForegroundColor Green
    if ($nonConformingFolders.Count -gt 0) {
        Write-Host ""  # Blank line here
        Write-Host "Non-conforming folders:" -ForegroundColor Yellow
        $nonConformingFolders | ForEach-Object {
            $truncatedName = $_.Substring(0, [Math]::Min($_.Length, 100))
            if ($_.Length -gt 100) {
                $truncatedName += "..."
            }
            Write-Host $truncatedName -ForegroundColor Yellow
        }
    }
}

# Get the directory containing the script
$scriptDirectory = Split-Path -Parent $PSCommandPath

# Call the function to set custom folder icons and create/update Desktop.ini files
Set-CustomFolderIcons