param(
    [string]$BundleRoot = "C:\Users\aryan\Google Drive\My Drive\nes2net_remote_bundle",
    [string]$ProjectRoot = "C:\E\Project\Project B.tech\Multimodal Deepfake Detection",
    [switch]$IncludeDFDataset,
    [switch]$IncludeLA,
    [switch]$IncludeDemoWav,
    [switch]$IncludeInTheWildCheckpoint,
    [switch]$ZipCodeOnly,
    [switch]$DryRun
)

$ErrorActionPreference = "Stop"

function Write-Step {
    param([string]$Message)
    Write-Host "[Bundle] $Message"
}

function Ensure-Dir {
    param([string]$Path)
    if ($DryRun) {
        Write-Step "Would create directory: $Path"
        return
    }
    New-Item -ItemType Directory -Force -Path $Path | Out-Null
}

function Copy-FileSafe {
    param(
        [string]$Source,
        [string]$DestinationDir
    )
    if (-not (Test-Path $Source)) {
        throw "Missing source file: $Source"
    }

    if ($DryRun) {
        Write-Step "Would copy file: $Source -> $DestinationDir"
        return
    }

    Copy-Item -Path $Source -Destination $DestinationDir -Force
}

function Copy-DirSafe {
    param(
        [string]$Source,
        [string]$DestinationParent
    )
    if (-not (Test-Path $Source)) {
        throw "Missing source directory: $Source"
    }

    $leaf = Split-Path -Path $Source -Leaf
    $destination = Join-Path $DestinationParent $leaf

    if ($DryRun) {
        Write-Step "Would copy directory: $Source -> $destination"
        return
    }

    Ensure-Dir $destination

    $robocopyArgs = @(
        $Source,
        $destination,
        "/E",
        "/Z",
        "/R:5",
        "/W:5",
        "/NFL",
        "/NDL",
        "/NJH",
        "/NJS",
        "/NP"
    )

    & robocopy @robocopyArgs | Out-Host
    $robocopyExit = $LASTEXITCODE

    if ($robocopyExit -gt 7) {
        throw "robocopy failed for $Source -> $destination with exit code $robocopyExit"
    }
}

function Zip-CodeFolder {
    param(
        [string]$Source,
        [string]$DestinationZip
    )
    if ($DryRun) {
        Write-Step "Would zip code folder: $Source -> $DestinationZip"
        return
    }

    if (Test-Path $DestinationZip) {
        Remove-Item -Path $DestinationZip -Force
    }

    Compress-Archive -Path $Source -DestinationPath $DestinationZip -Force
}

$codeDir = Join-Path $ProjectRoot "audio_detection\Nes2Net_ASVspoof_ITW"
$checkpointsDir = Join-Path $ProjectRoot "audio_detection\checkpoints"
$audioDataDir = Join-Path $ProjectRoot "data\audio"
$notebookPath = Join-Path $ProjectRoot "notebooks\audio_remote_validation_colab.ipynb"

$bundleCodeDir = Join-Path $BundleRoot "code"
$bundleCheckpointsDir = Join-Path $BundleRoot "checkpoints"
$bundleNotebooksDir = Join-Path $BundleRoot "notebooks"
$bundleAudioDir = Join-Path $BundleRoot "data\audio"

Write-Step "Preparing bundle root: $BundleRoot"
Ensure-Dir $bundleCodeDir
Ensure-Dir $bundleCheckpointsDir
Ensure-Dir $bundleNotebooksDir

if ($IncludeDFDataset -or $IncludeLA -or $IncludeDemoWav) {
    Ensure-Dir $bundleAudioDir
}

Write-Step "Copying code bundle"
Copy-DirSafe -Source $codeDir -DestinationParent $bundleCodeDir

Write-Step "Copying Colab notebook"
Copy-FileSafe -Source $notebookPath -DestinationDir $bundleNotebooksDir

Write-Step "Copying required checkpoints"
Copy-FileSafe -Source (Join-Path $checkpointsDir "xlsr2_300m.pt") -DestinationDir $bundleCheckpointsDir
Copy-FileSafe -Source (Join-Path $checkpointsDir "ASVspoof_2021_wav2vec2_Nes2Net_X_e100_bz12_lr2.5e_07_algo4_avg_ckpt_ep56_60_62_76_95.pth") -DestinationDir $bundleCheckpointsDir

if ($IncludeInTheWildCheckpoint) {
    Write-Step "Copying optional In-the-Wild checkpoint"
    Copy-FileSafe -Source (Join-Path $checkpointsDir "In_the_Wild_wav2vec2_Nes2Net_X_e100_bz12_lr2.5e_07_algo4_seed12345_avg_ckpt_ep59_61_69.pth") -DestinationDir $bundleCheckpointsDir
}

if ($IncludeDFDataset) {
    Write-Step "Copying optional DF dataset"
    Copy-DirSafe -Source (Join-Path $audioDataDir "ASVspoof2021_DF_eval") -DestinationParent $bundleAudioDir
}

if ($IncludeLA) {
    Write-Step "Copying optional LA dataset"
    Copy-DirSafe -Source (Join-Path $audioDataDir "ASVspoof2021_LA_eval") -DestinationParent $bundleAudioDir
}

if ($IncludeDemoWav) {
    Write-Step "Copying optional demo WAV folder"
    Copy-DirSafe -Source (Join-Path $audioDataDir "demo_wav") -DestinationParent $bundleAudioDir
}

if ($ZipCodeOnly) {
    $zipPath = Join-Path $BundleRoot "Nes2Net_ASVspoof_ITW.zip"
    Write-Step "Creating optional code zip"
    Zip-CodeFolder -Source $codeDir -DestinationZip $zipPath
}

Write-Step "Bundle preparation complete."
if ($DryRun) {
    Write-Step "Dry run only. No files were copied."
}