#Requires -RunAsAdministrator

# --- Paramètres du script ---
param(
    # Définit le chemin d'installation.
    # Si le script est lancé sans cet argument, il utilisera son propre dossier par défaut.
    [string]$InstallPath = $PSScriptRoot
)

<#
.SYNOPSIS
    A professional installer for ComfyUI using Git and a Python venv.
.DESCRIPTION
    This script provides a clean, modern installation of ComfyUI by:
    1. Checking for and installing Python 3.12 if needed.
    2. Cloning the official ComfyUI repository.
    3. Creating a dedicated Python virtual environment (venv).
    4. Installing all dependencies into the venv.
    5. Creating a launcher to run the application easily.
.NOTES
    Author: Code Partner
    Version: 7.0 (Git + Venv Architecture)
#>

#===========================================================================
# SECTION 1: SCRIPT CONFIGURATION & HELPER FUNCTIONS
#===========================================================================

# --- Nettoyage et configuration des chemins ---
$InstallPath = $InstallPath.Trim('"')
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12

# Le chemin de base est maintenant le dossier d'installation principal.
$comfyPath = Join-Path $InstallPath "ComfyUI"
$logPath = Join-Path $InstallPath "logs"
$logFile = Join-Path $logPath "install_log.txt"
$sevenZipPath = "C:\Program Files\7-Zip\7z.exe"

# Variable qui contiendra le chemin vers le python du venv
$venvPython = Join-Path $comfyPath "venv\Scripts\python.exe"

# --- Create Log Directory ---
if (-not (Test-Path $logPath)) {
    New-Item -ItemType Directory -Force -Path $logPath | Out-Null
}

# --- Helper functions ---
function Write-Log {
    param([string]$Message, [string]$Color = "White")
    $timestamp = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    $formattedMessage = "[$timestamp] $Message"
    Write-Host $Message -ForegroundColor $Color
    Add-Content -Path $logFile -Value $formattedMessage
}

function Invoke-AndLog {
    param(
        [string]$File,
        [string]$Arguments
    )
    
    # Chemin vers un fichier de log temporaire unique
    $tempLogFile = Join-Path $env:TEMP ([System.Guid]::NewGuid().ToString() + ".tmp")

    try {
        # Exécute la commande et redirige TOUTE sa sortie vers le fichier temporaire
        $commandToRun = "`"$File`" $Arguments"
        $cmdArguments = "/C `"$commandToRun > `"`"$tempLogFile`"`" 2>&1`""
        Start-Process -FilePath "cmd.exe" -ArgumentList $cmdArguments -Wait -WindowStyle Hidden
        
        # Une fois la commande terminée, on lit le fichier temporaire
        if (Test-Path $tempLogFile) {
            $output = Get-Content $tempLogFile
            # Et on l'ajoute au log principal en toute sécurité
            Add-Content -Path $logFile -Value $output
        }
    } catch {
        Write-Log "FATAL ERROR trying to execute command: $commandToRun" -Color Red
    } finally {
        # On s'assure que le fichier temporaire est toujours supprimé
        if (Test-Path $tempLogFile) {
            Remove-Item $tempLogFile
        }
    }
}

function Download-File {
    param([string]$Uri, [string]$OutFile)
    if (Test-Path $OutFile) {
        Write-Log "Skipping: $((Split-Path $OutFile -Leaf)) (already exists)." -Color Gray
    } else {
        $fileName = Split-Path -Path $Uri -Leaf
        if (Get-Command 'aria2c' -ErrorAction SilentlyContinue) {
            Write-Log "Downloading: $fileName"
            $aria_args = "-c -x 16 -s 16 -k 1M --dir=`"$((Split-Path $OutFile -Parent))`" --out=`"$((Split-Path $OutFile -Leaf))`" `"$Uri`""
            Invoke-AndLog "aria2c" $aria_args
        } else {
            Write-Log "Aria2 not found. Falling back to standard download: $fileName" -Color Yellow
            Invoke-WebRequest -Uri $Uri -OutFile $OutFile
        }
    }
}

function Install-Aria2-Binary {
    Write-Log "--- Starting Aria2 binary installation ---" -Color Magenta
    $destFolder = "C:\Tools\aria2"
    if (-not (Test-Path $destFolder)) { New-Item -ItemType Directory -Force -Path $destFolder | Out-Null }
    $aria2Url = "https://github.com/aria2/aria2/releases/download/release-1.36.0/aria2-1.36.0-win-64bit-build1.zip"
    $zipPath  = Join-Path $env:TEMP "aria2_temp.zip"
    Download-File -Uri $aria2Url -OutFile $zipPath
    Write-Log "Extracting zip file to $destFolder..."
    Expand-Archive -Path $zipPath -DestinationPath $destFolder -Force
    $extractedSubfolder = Join-Path $destFolder "aria2-1.36.0-win-64bit-build1"
    if (Test-Path $extractedSubfolder) {
        Move-Item -Path (Join-Path $extractedSubfolder "*") -Destination $destFolder -Force
        Remove-Item -Path $extractedSubfolder -Recurse -Force
    }
    $configFile = Join-Path $destFolder "aria2.conf"
    $configContent = "continue=true`nmax-connection-per-server=16`nsplit=16`nmin-split-size=1M`nfile-allocation=none"
    $configContent | Out-File $configFile -Encoding UTF8
    $envScope = "User"
    $oldPath = [System.Environment]::GetEnvironmentVariable("Path", $envScope)
    if ($oldPath -notlike "*$destFolder*") {
        Write-Log "Adding '$destFolder' to user PATH..."
        $newPath = $oldPath + ";$destFolder"
        [System.Environment]::SetEnvironmentVariable("Path", $newPath, $envScope)
        $env:Path = $newPath
        Write-Log "PATH updated. Aria2 will be available immediately."
    }
    Write-Log "--- Aria2 binary installation complete ---" -Color Magenta
}


#===========================================================================
# SECTION 2: MAIN SCRIPT EXECUTION
#===========================================================================

Clear-Host
# --- Bannière ---
Write-Log "-------------------------------------------------------------------------------"
$asciiBanner = @'
                      __  __               ___    _ ____  ______
                     / / / /___ ___  ___  /   |  (_) __ \/_  __/
                    / / / / __ `__ \/ _ \/ /| | / / /_/ / / / 
                   / /_/ / / / / / /  __/ ___ |/ / _, _/ / /
                   \____/_/ /_/ /_/\___/_/  |_/_/_/ |_| /_/
'@
Write-Host $asciiBanner -ForegroundColor Cyan
Write-Log "-------------------------------------------------------------------------------"
Write-Log "                 ComfyUI - Git & Venv Based Installer                          " -Color Yellow
Write-Log "                                  Version 7.0                                  " -Color White
Write-Log "-------------------------------------------------------------------------------"

# --- Étape 1: Vérification et Installation de Python ---
Write-Log "Step 1: Checking for Python 3.12..." -Color Yellow
$pythonExe = Get-Command python -ErrorAction SilentlyContinue
$pythonVersionOK = $false
if ($pythonExe) {
    # Capturer la sortie (va souvent sur le flux d'erreur)
    $versionString = (python --version 2>&1)
    Write-Log "  - Found Python: $versionString"
    if ($versionString -like "Python 3.12*") {
        Write-Log "  - Correct version already installed." -Color Green
        $pythonVersionOK = $true
    }
}

if (-not $pythonVersionOK) {
    Write-Log "  - Python 3.12 not found. Downloading and installing..." -Color Yellow
    $pythonInstallerPath = Join-Path $env:TEMP "python-3.12-installer.exe"
    Download-File -Uri "https://www.python.org/ftp/python/3.12.4/python-3.12.4-amd64.exe" -OutFile $pythonInstallerPath
    Write-Log "  - Running Python installer silently... (This may take a moment)"
    # Installation silencieuse, pour tous les utilisateurs, et ajout au PATH
    Start-Process -FilePath $pythonInstallerPath -ArgumentList "/quiet InstallAllUsers=1 PrependPath=1" -Wait
    Remove-Item $pythonInstallerPath
    Write-Log "  - Python 3.12 installation complete." -Color Green
}

# --- Étape 2: Installation des dépendances (Aria2, 7-Zip, Git) ---
Write-Log "`nStep 2: Checking for required tools..." -Color Yellow
if (-not (Get-Command 'aria2c' -ErrorAction SilentlyContinue)) { Install-Aria2-Binary } else { Write-Log "  - Aria2 is already installed." -Color Green }
if (-not (Test-Path $sevenZipPath)) {
    $sevenZipInstaller = Join-Path $env:TEMP "7z-installer.exe"; Download-File -Uri "https://www.7-zip.org/a/7z2201-x64.exe" -OutFile $sevenZipInstaller; Start-Process -FilePath $sevenZipInstaller -ArgumentList "/S" -Wait; Remove-Item $sevenZipInstaller
} else { Write-Log "  - 7-Zip is already installed." -Color Green }
if (-not (Get-Command git.exe -ErrorAction SilentlyContinue)) {
    $gitInstaller = Join-Path $env:TEMP "Git-Installer.exe"; Download-File -Uri "https://github.com/git-for-windows/git/releases/download/v2.41.0.windows.3/Git-2.41.0.3-64-bit.exe" -OutFile $gitInstaller; Start-Process -FilePath $gitInstaller -ArgumentList "/VERYSILENT" -Wait; Remove-Item $gitInstaller
}
Invoke-AndLog "git" "config --system core.longpaths true"
Write-Log "  - Git is ready." -Color Green


# --- Étape 3: Cloner ComfyUI et créer le Venv ---
Write-Log "`nStep 3: Cloning ComfyUI and creating Virtual Environment..." -Color Yellow
if (-not (Test-Path $comfyPath)) {
    Write-Log "  - Cloning ComfyUI repository..."
    Invoke-AndLog "git" "clone https://github.com/comfyanonymous/ComfyUI.git `"$comfyPath`""
} else {
    Write-Log "  - ComfyUI directory already exists. Skipping clone." -Color Green
}

if (-not (Test-Path (Join-Path $comfyPath "venv"))) {
    Write-Log "  - Creating Python virtual environment (venv)..."
    Push-Location $comfyPath
    Invoke-AndLog "python" "-m venv venv"
    Pop-Location
    Write-Log "  - Venv created successfully." -Color Green
} else {
    Write-Log "  - Venv already exists. Skipping creation." -Color Green
}

# --- Étape 4: Installation des dépendances dans le Venv ---
Write-Log "`nStep 4: Installing all Python dependencies into the venv..." -Color Yellow
Invoke-AndLog "$venvPython" "-m pip install --upgrade pip wheel"
Invoke-AndLog "$venvPython" "-m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121"
Invoke-AndLog "$venvPython" "-m pip install -r `"$comfyPath\requirements.txt`""

# --- Étape 5: Installation des Custom Nodes ---
Write-Log "`nStep 5: Installing custom nodes..." -Color Yellow
$customNodesPath = Join-Path $InstallPath "custom_nodes" # Définit le chemin pour les noeuds
# ... (La liste des noeuds est inchangée)
$customNodes = @(
    @{Name="ComfyUI-Manager"; Repo="https://github.com/ltdrdata/ComfyUI-Manager.git"; HasRequirements=$false},
    @{Name="ComfyUI-Impact-Pack"; Repo="https://github.com/ltdrdata/ComfyUI-Impact-Pack"; HasRequirements=$true},
    @{Name="ComfyUI-Impact-Subpack"; Repo="https://github.com/ltdrdata/ComfyUI-Impact-Subpack"; Subfolder="ComfyUI-Impact-Pack/impact_subpack"; HasRequirements=$true},
    @{Name="ComfyUI-GGUF"; Repo="https://github.com/city96/ComfyUI-GGUF"; HasRequirements=$true},
    @{Name="ComfyUI-mxToolkit"; Repo="https://github.com/Smirnov75/ComfyUI-mxToolkit"; HasRequirements=$false},
    @{Name="ComfyUI-Custom-Scripts"; Repo="https://github.com/pythongosssss/ComfyUI-Custom-Scripts"; HasRequirements=$false},
    @{Name="ComfyUI-KJNodes"; Repo="https://github.com/kijai/ComfyUI-KJNodes"; HasRequirements=$true},
    @{Name="ComfyUI-WanVideoWrapper"; Repo="https://github.com/kijai/ComfyUI-WanVideoWrapper"; HasRequirements=$true},
    @{Name="ComfyUI-VideoHelperSuite"; Repo="https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"; HasRequirements=$true},
    @{Name="ComfyUI-Frame-Interpolation"; Repo="https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"; RequirementsFile="requirements-with-cupy.txt"; HasRequirements=$true},
    @{Name="rgthree-comfy"; Repo="https://github.com/rgthree/rgthree-comfy"; HasRequirements=$true},
    @{Name="ComfyUI-Easy-Use"; Repo="https://github.com/yolain/ComfyUI-Easy-Use"; HasRequirements=$true},
    @{Name="ComfyUI_PuLID_Flux_ll"; Repo="https://github.com/lldacing/ComfyUI_PuLID_Flux_ll"; HasRequirements=$true},
    @{Name="ComfyUI-HunyuanVideoMultiLora"; Repo="https://github.com/facok/ComfyUI-HunyuanVideoMultiLora"; HasRequirements=$false},
    @{Name="was-node-suite-comfyui"; Repo="https://github.com/WASasquatch/was-node-suite-comfyui"; HasRequirements=$true},
    @{Name="ComfyUI-Florence2"; Repo="https://github.com/kijai/ComfyUI-Florence2"; HasRequirements=$true},
    @{Name="ComfyUI-MultiGPU"; Repo="https://github.com/pollockjj/ComfyUI-MultiGPU"; HasRequirements=$false},
    @{Name="ComfyUI-WanStartEndFramesNative"; Repo="https://github.com/Flow-two/ComfyUI-WanStartEndFramesNative"; HasRequirements=$false},
    @{Name="ComfyUI-Image-Saver"; Repo="https://github.com/alexopus/ComfyUI-Image-Saver"; HasRequirements=$true},
    @{Name="ComfyUI_UltimateSDUpscale"; Repo="https://github.com/ssitu/ComfyUI_UltimateSDUpscale"; HasRequirements=$false},
    @{Name="comfyui_controlnet_aux"; Repo="https://github.com/Fannovel16/comfyui_controlnet_aux"; HasRequirements=$true},
    @{Name="x-flux-comfyui"; Repo="https://github.com/XLabs-AI/x-flux-comfyui"; HasRequirements=$true},
    @{Name="ComfyUI-RMBG"; Repo="https://github.com/1038lab/ComfyUI-RMBG"; HasRequirements=$true},
    @{Name="ComfyUI-Detail-Daemon"; Repo="https://github.com/Jonseed/ComfyUI-Detail-Daemon"; HasRequirements=$true},
    @{Name="ComfyUI-TeaCache"; Repo="https://github.com/welltop-cn/ComfyUI-TeaCache"; HasRequirements=$true},
    @{Name="ComfyUI-Crystools"; Repo="https://github.com/crystian/ComfyUI-Crystools"; HasRequirements=$true}
)
foreach ($node in $customNodes) {
    $nodePath = if ($node.Subfolder) { Join-Path $customNodesPath $node.Subfolder } else { Join-Path $customNodesPath $node.Name }
    if (-not (Test-Path $nodePath)) {
        Write-Log "  - Installing $($node.Name)..."
        $cloneTargetPath = if ($node.Subfolder) { (Split-Path $nodePath -Parent) } else { $nodePath }
        if ($node.Name -eq 'ComfyUI-Impact-Subpack') { $clonePath = Join-Path $cloneTargetPath "impact_subpack" } else { $clonePath = $cloneTargetPath }
        Invoke-AndLog "git" "clone $($node.Repo) `"$clonePath`""
        if ($node.HasRequirements) {
            $requirementsFile = if ($node.RequirementsFile) { $node.RequirementsFile } else { "requirements.txt" }
            $reqPath = Join-Path $nodePath $requirementsFile
            if (Test-Path $reqPath) { Invoke-AndLog "$venvPython" "-m pip install -r `"$reqPath`"" }
        }
    } else { Write-Log "  - $($node.Name) (already exists, skipping)" -Color Green }
}

# --- Étape 6: Installation des modules Python supplémentaires ---
Write-Log "`nStep 6: Installing supplementary Python modules..." -Color Yellow
Invoke-AndLog "$venvPython" "-m pip install -U xformers --index-url https://download.pytorch.org/whl/cu121"
# ... (Ajoutez ici l'installation de Triton, SageAttention etc. si nécessaire, en utilisant $venvPython) ...

# --- Étape 7: Téléchargement des Workflows et Settings ---
Write-Log "`nStep 7: Downloading Workflows & Settings..." -Color Yellow
$userDefaultPath = Join-Path $comfyPath "user\default"; New-Item -Path $userDefaultPath -ItemType Directory -Force | Out-Null
$workflowPath = Join-Path $userDefaultPath "workflows"; New-Item -Path $workflowPath -ItemType Directory -Force | Out-Null
$workflowCloneDest = Join-Path $workflowPath "UmeAiRT-Workflow"
$settingsFilePath = Join-Path $userDefaultPath "comfy.settings.json"
Download-File -Uri "https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/others/comfy.settings.json" -OutFile $settingsFilePath
if (-not (Test-Path $workflowCloneDest)) { Invoke-AndLog "git" "clone https://github.com/UmeAiRT/ComfyUI-Workflows `"$workflowCloneDest`"" }

# --- Étape 9: Téléchargement optionnel des packs de modèles ---
Write-Log "`nStep 9: Optional Model Pack Downloads..." -Color Yellow
$modelPacks = @( @{Name="FLUX"; ScriptName="Download-FLUX-Models.ps1"} )
$scriptsSubFolder = Join-Path $InstallPath "scripts"
foreach ($pack in $modelPacks) {
    $scriptPath = Join-Path $scriptsSubFolder $pack.ScriptName
    if (-not (Test-Path $scriptPath)) { Write-Log "Model downloader script not found: '$($pack.ScriptName)'. Skipping." -Color Red; continue }
    $validInput = $false
    while (-not $validInput) {
        Write-Host "`nWould you like to download $($pack.Name) models? (Y/N)"
        $choice = Read-Host
        if ($choice -eq 'Y' -or $choice -eq 'y') {
            Write-Log "  - Launching downloader for $($pack.Name) models..." -Color Green
            & $scriptPath -InstallPath $InstallPath
            $validInput = $true
        } elseif ($choice -eq 'N' -or $choice -eq 'n') {
            Write-Log "  - Skipping download for $($pack.Name) models." -Color Gray
            $validInput = $true
        } else { Write-Host "  Invalid choice. Please enter Y or N." }
    }
}

#===========================================================================
# FINALIZATION
#===========================================================================
Write-Log "-------------------------------------------------------------------------------" -Color Green
Write-Log "Installation of ComfyUI and all nodes is complete!" -Color Green
Read-Host "Press Enter to close this window."