param(
    [string]$Root = "D:\AI\ComfyUI_windows_portable",
    [int]$Port = 8188,
    [switch]$Json
)

$ErrorActionPreference = "SilentlyContinue"

function Get-GitState {
    param([string]$Path)

    if (-not (Test-Path -LiteralPath $Path)) {
        return $null
    }
    if (-not (Test-Path -LiteralPath (Join-Path $Path ".git"))) {
        return [ordered]@{
            branch = $null
            commit = $null
            dirty = $null
        }
    }

    $oldLocation = Get-Location
    try {
        Set-Location -LiteralPath $Path
        $branch = ((git rev-parse --abbrev-ref HEAD 2>$null) -join "").Trim()
        $commit = ((git rev-parse --short HEAD 2>$null) -join "").Trim()
        $dirtyOutput = git status --porcelain 2>$null
        return [ordered]@{
            branch = $branch
            commit = $commit
            dirty = [bool]$dirtyOutput
        }
    }
    finally {
        Set-Location $oldLocation
    }
}

function Get-DirectoryStats {
    param([string]$Path)

    if (-not (Test-Path -LiteralPath $Path)) {
        return [ordered]@{
            exists = $false
            files = 0
            size_gb = 0
        }
    }

    $files = @(Get-ChildItem -LiteralPath $Path -File -Recurse -ErrorAction SilentlyContinue |
        Where-Object { $_.FullName -notmatch '\\\.cache\\' })
    $sum = ($files | Measure-Object -Property Length -Sum).Sum
    if (-not $sum) {
        $sum = 0
    }

    return [ordered]@{
        exists = $true
        files = $files.Count
        size_gb = [math]::Round($sum / 1GB, 3)
    }
}

function Read-TorchInfo {
    param(
        [string]$Root
    )

    $version = $null
    $cudaRuntime = $null
    $torchVersionPy = Join-Path $Root "python_embeded\Lib\site-packages\torch\version.py"
    if (Test-Path -LiteralPath $torchVersionPy) {
        $versionLine = Select-String -LiteralPath $torchVersionPy -Pattern "^__version__" | Select-Object -First 1
        if ($versionLine -and $versionLine.Line -match "['""]([^'""]+)['""]") {
            $version = $Matches[1]
        }
        $cudaLine = Select-String -LiteralPath $torchVersionPy -Pattern "^cuda:" | Select-Object -First 1
        if ($cudaLine -and $cudaLine.Line -match "['""]([^'""]+)['""]") {
            $cudaRuntime = $Matches[1]
        }
    }

    return [ordered]@{
        torch = $version
        cuda_runtime = $cudaRuntime
        cuda_check = "not_checked_to_avoid_driver_runtime_crashes"
    }
}

$comfyRoot = Join-Path $Root "ComfyUI"
$python = Join-Path $Root "python_embeded\python.exe"
$modelsRoot = Join-Path $comfyRoot "models"
$customNodesRoot = Join-Path $comfyRoot "custom_nodes"
$versionPath = Join-Path $comfyRoot "comfyui_version.py"

$version = $null
if (Test-Path -LiteralPath $versionPath) {
    $versionLine = Select-String -LiteralPath $versionPath -Pattern "__version__" | Select-Object -First 1
    if ($versionLine -and $versionLine.Line -match '["'']([^"'']+)["'']') {
        $version = $Matches[1]
    }
}

$nvidia = $null
if (Get-Command nvidia-smi -ErrorAction SilentlyContinue) {
    $nvidia = ((nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader 2>$null) -join "`n").Trim()
}

$portState = $null
if (Get-Command Get-NetTCPConnection -ErrorAction SilentlyContinue) {
    $portState = @(Get-NetTCPConnection -LocalPort $Port -State Listen -ErrorAction SilentlyContinue).Count -gt 0
}

$httpStatus = $null
try {
    $response = Invoke-WebRequest -Uri "http://127.0.0.1:$Port/" -UseBasicParsing -TimeoutSec 2
    $httpStatus = [int]$response.StatusCode
}
catch {
    $httpStatus = $null
}

$processes = @(Get-CimInstance Win32_Process -Filter "name = 'python.exe' OR name = 'python3.exe'" |
    Where-Object { $_.CommandLine -like "*ComfyUI*" -and $_.CommandLine -like "*main.py*" } |
    Select-Object ProcessId, CommandLine)

$modelDirs = @(
    "checkpoints",
    "clip_vision",
    "diffusion_models",
    "loras",
    "text_encoders",
    "vae",
    "vae_approx"
)

$modelStats = [ordered]@{}
foreach ($dir in $modelDirs) {
    $modelStats[$dir] = Get-DirectoryStats -Path (Join-Path $modelsRoot $dir)
}

$wanModels = @()
$diffusionPath = Join-Path $modelsRoot "diffusion_models"
if (Test-Path -LiteralPath $diffusionPath) {
    $wanModels = @(Get-ChildItem -LiteralPath $diffusionPath -File -ErrorAction SilentlyContinue |
        Where-Object { $_.Name -match "wan|Wan" } |
        Sort-Object Name |
        ForEach-Object {
            [ordered]@{
                name = $_.Name
                size_gb = [math]::Round($_.Length / 1GB, 3)
            }
        })
}

$knownNodes = @(
    "ComfyUI-Manager",
    "ComfyUI-GGUF",
    "ComfyUI-WanVideoWrapper",
    "ComfyUI-KJNodes",
    "ComfyUI-VideoHelperSuite"
)

$customNodes = @()
foreach ($node in $knownNodes) {
    $nodePath = Join-Path $customNodesRoot $node
    $customNodes += [ordered]@{
        name = $node
        exists = Test-Path -LiteralPath $nodePath
        path = $nodePath
        git = Get-GitState -Path $nodePath
    }
}

$result = [ordered]@{
    root = $Root
    comfy_root = $comfyRoot
    exists = Test-Path -LiteralPath $comfyRoot
    version = $version
    git = Get-GitState -Path $comfyRoot
    python = $python
    torch = Read-TorchInfo -Root $Root
    nvidia = $nvidia
    port = [ordered]@{
        number = $Port
        listening = $portState
        http_status = $httpStatus
    }
    processes = $processes
    custom_nodes = $customNodes
    model_stats = $modelStats
    wan_diffusion_models = $wanModels
}

if ($Json) {
    $result | ConvertTo-Json -Depth 8
    exit
}

Write-Host "ComfyUI audit"
Write-Host "Root: $($result.root)"
Write-Host "Exists: $($result.exists)"
Write-Host "Version: $($result.version)"
if ($result.git) {
    Write-Host "Git: $($result.git.branch) $($result.git.commit) dirty=$($result.git.dirty)"
}
Write-Host "Python: $($result.python)"
if ($result.torch) {
    Write-Host "Torch: $($result.torch.torch) cuda_runtime=$($result.torch.cuda_runtime) cuda_check=$($result.torch.cuda_check)"
}
Write-Host "NVIDIA: $($result.nvidia)"
Write-Host "Port ${Port}: listening=$($result.port.listening) http_status=$($result.port.http_status)"
Write-Host "ComfyUI processes: $($result.processes.Count)"

Write-Host ""
Write-Host "Custom nodes"
foreach ($node in $result.custom_nodes) {
    $commit = $null
    $dirty = $null
    if ($node.git) {
        $commit = $node.git.commit
        $dirty = $node.git.dirty
    }
    Write-Host ("- {0}: exists={1} commit={2} dirty={3}" -f $node.name, $node.exists, $commit, $dirty)
}

Write-Host ""
Write-Host "Model directories"
foreach ($dir in $modelDirs) {
    $stats = $result.model_stats[$dir]
    Write-Host ("- {0}: exists={1} files={2} size_gb={3}" -f $dir, $stats.exists, $stats.files, $stats.size_gb)
}

Write-Host ""
Write-Host "Wan diffusion models: $($result.wan_diffusion_models.Count)"
foreach ($model in $result.wan_diffusion_models) {
    Write-Host ("- {0} ({1} GB)" -f $model.name, $model.size_gb)
}